diff --git a/.github/workflows/notify_slack_about_release.yml b/.github/workflows/notify_slack_about_release.yml
index 6c0b96954e81..7751827d81f5 100644
--- a/.github/workflows/notify_slack_about_release.yml
+++ b/.github/workflows/notify_slack_about_release.yml
@@ -15,7 +15,7 @@ jobs:
- name: Setup Python
uses: actions/setup-python@v6
with:
- python-version: '3.8'
+ python-version: '3.10'
- name: Notify Slack about the release
env:
diff --git a/.github/workflows/pr_dependency_test.yml b/.github/workflows/pr_dependency_test.yml
index dfc35c41066f..ba5cd1c76cbc 100644
--- a/.github/workflows/pr_dependency_test.yml
+++ b/.github/workflows/pr_dependency_test.yml
@@ -22,7 +22,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v6
with:
- python-version: "3.8"
+ python-version: "3.10"
- name: Install dependencies
run: |
pip install -e .
diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml
index c0dfa89e776d..3e69dfd05eae 100644
--- a/.github/workflows/pr_tests.yml
+++ b/.github/workflows/pr_tests.yml
@@ -35,7 +35,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v6
with:
- python-version: "3.8"
+ python-version: "3.10"
- name: Install dependencies
run: |
pip install --upgrade pip
@@ -55,7 +55,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v6
with:
- python-version: "3.8"
+ python-version: "3.10"
- name: Install dependencies
run: |
pip install --upgrade pip
@@ -92,7 +92,6 @@ jobs:
runner: aws-general-8-plus
image: diffusers/diffusers-pytorch-cpu
report: torch_example_cpu
-
name: ${{ matrix.config.name }}
runs-on:
@@ -115,8 +114,7 @@ jobs:
- name: Install dependencies
run: |
uv pip install -e ".[quality]"
- #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1
+ uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps
- name: Environment
@@ -218,8 +216,6 @@ jobs:
run_lora_tests:
needs: [check_code_quality, check_repository_consistency]
- strategy:
- fail-fast: false
name: LoRA tests with PEFT main
@@ -247,9 +243,8 @@ jobs:
uv pip install -U peft@git+https://github.com/huggingface/peft.git --no-deps
uv pip install -U tokenizers
uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps
- #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1
-
+ uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
+
- name: Environment
run: |
python utils/print_env.py
@@ -275,6 +270,6 @@ jobs:
if: ${{ always() }}
uses: actions/upload-artifact@v6
with:
- name: pr_main_test_reports
+ name: pr_lora_test_reports
path: reports
diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml
index dd20bbe93250..b79d80f71c09 100644
--- a/.github/workflows/pr_tests_gpu.yml
+++ b/.github/workflows/pr_tests_gpu.yml
@@ -36,7 +36,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v6
with:
- python-version: "3.8"
+ python-version: "3.10"
- name: Install dependencies
run: |
pip install --upgrade pip
@@ -56,7 +56,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v6
with:
- python-version: "3.8"
+ python-version: "3.10"
- name: Install dependencies
run: |
pip install --upgrade pip
@@ -131,8 +131,7 @@ jobs:
run: |
uv pip install -e ".[quality]"
uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git
- #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1
+ uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- name: Environment
run: |
@@ -202,8 +201,7 @@ jobs:
uv pip install -e ".[quality]"
uv pip install peft@git+https://github.com/huggingface/peft.git
uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git
- #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1
+ uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- name: Environment
run: |
@@ -264,8 +262,7 @@ jobs:
nvidia-smi
- name: Install dependencies
run: |
- #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1
+ uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
uv pip install -e ".[quality,training]"
- name: Environment
diff --git a/.github/workflows/pr_torch_dependency_test.yml b/.github/workflows/pr_torch_dependency_test.yml
index 4a7e5ab37e47..79569488ae21 100644
--- a/.github/workflows/pr_torch_dependency_test.yml
+++ b/.github/workflows/pr_torch_dependency_test.yml
@@ -22,7 +22,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v6
with:
- python-version: "3.8"
+ python-version: "3.10"
- name: Install dependencies
run: |
pip install -e .
diff --git a/.github/workflows/push_tests.yml b/.github/workflows/push_tests.yml
index 4456f18c95bc..61bb9f0ef679 100644
--- a/.github/workflows/push_tests.yml
+++ b/.github/workflows/push_tests.yml
@@ -76,8 +76,7 @@ jobs:
run: |
uv pip install -e ".[quality]"
uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git
- #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1
+ uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- name: Environment
run: |
python utils/print_env.py
@@ -129,8 +128,7 @@ jobs:
uv pip install -e ".[quality]"
uv pip install peft@git+https://github.com/huggingface/peft.git
uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git
- #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1
+ uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- name: Environment
run: |
@@ -182,8 +180,7 @@ jobs:
- name: Install dependencies
run: |
uv pip install -e ".[quality,training]"
- #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1
+ uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git
- name: Environment
run: |
python utils/print_env.py
diff --git a/.github/workflows/push_tests_mps.yml b/.github/workflows/push_tests_mps.yml
index cc16d5f82cd0..7f8ce9a4b99d 100644
--- a/.github/workflows/push_tests_mps.yml
+++ b/.github/workflows/push_tests_mps.yml
@@ -41,7 +41,7 @@ jobs:
shell: arch -arch arm64 bash {0}
run: |
${CONDA_RUN} python -m pip install --upgrade pip uv
- ${CONDA_RUN} python -m uv pip install -e ".[quality,test]"
+ ${CONDA_RUN} python -m uv pip install -e ".[quality]"
${CONDA_RUN} python -m uv pip install torch torchvision torchaudio
${CONDA_RUN} python -m uv pip install accelerate@git+https://github.com/huggingface/accelerate.git
${CONDA_RUN} python -m uv pip install transformers --upgrade
diff --git a/.github/workflows/pypi_publish.yaml b/.github/workflows/pypi_publish.yaml
index 214b996f5381..1e8ce4599544 100644
--- a/.github/workflows/pypi_publish.yaml
+++ b/.github/workflows/pypi_publish.yaml
@@ -20,7 +20,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v6
with:
- python-version: '3.8'
+ python-version: '3.10'
- name: Fetch latest branch
id: fetch_latest_branch
@@ -47,7 +47,7 @@ jobs:
- name: Setup Python
uses: actions/setup-python@v6
with:
- python-version: "3.8"
+ python-version: "3.10"
- name: Install dependencies
run: |
diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml
index b0a90e278550..76dd48d09931 100644
--- a/.github/workflows/stale.yml
+++ b/.github/workflows/stale.yml
@@ -20,7 +20,7 @@ jobs:
- name: Setup Python
uses: actions/setup-python@v6
with:
- python-version: 3.8
+ python-version: 3.10
- name: Install requirements
run: |
diff --git a/LICENSE b/LICENSE
index 261eeb9e9f8b..038e32f6445e 100644
--- a/LICENSE
+++ b/LICENSE
@@ -144,7 +144,7 @@
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
- implied, including, without limitation, any warranties or conditions
+ implied, including, without limitation, Any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
diff --git a/benchmarks/benchmarking_utils.py b/benchmarks/benchmarking_utils.py
index c8c1a10ef899..141850e64f2e 100644
--- a/benchmarks/benchmarking_utils.py
+++ b/benchmarks/benchmarking_utils.py
@@ -6,7 +6,7 @@
import threading
from contextlib import nullcontext
from dataclasses import dataclass
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable
import pandas as pd
import torch
@@ -91,10 +91,10 @@ def model_init_fn(model_cls, group_offload_kwargs=None, layerwise_upcasting=Fals
class BenchmarkScenario:
name: str
model_cls: ModelMixin
- model_init_kwargs: Dict[str, Any]
+ model_init_kwargs: dict[str, Any]
model_init_fn: Callable
get_model_input_dict: Callable
- compile_kwargs: Optional[Dict[str, Any]] = None
+ compile_kwargs: dict[str, Any] | None = None
@require_torch_gpu
@@ -176,7 +176,7 @@ def run_benchmark(self, scenario: BenchmarkScenario):
result["fullgraph"], result["mode"] = None, None
return result
- def run_bencmarks_and_collate(self, scenarios: Union[BenchmarkScenario, list[BenchmarkScenario]], filename: str):
+ def run_bencmarks_and_collate(self, scenarios: BenchmarkScenario | list[BenchmarkScenario], filename: str):
if not isinstance(scenarios, list):
scenarios = [scenarios]
record_queue = queue.Queue()
@@ -214,10 +214,10 @@ def _run_phase(
*,
model_cls: ModelMixin,
init_fn: Callable,
- init_kwargs: Dict[str, Any],
+ init_kwargs: dict[str, Any],
get_input_fn: Callable,
- compile_kwargs: Optional[Dict[str, Any]],
- ) -> Dict[str, float]:
+ compile_kwargs: dict[str, Any] | None = None,
+ ) -> dict[str, float]:
# setup
self.pre_benchmark()
diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml
index 64a4222845b0..098660ec3f39 100644
--- a/docs/source/en/_toctree.yml
+++ b/docs/source/en/_toctree.yml
@@ -625,8 +625,7 @@
title: Image-to-image
- local: api/pipelines/stable_diffusion/inpaint
title: Inpainting
- - local: api/pipelines/stable_diffusion/k_diffusion
- title: K-Diffusion
+
- local: api/pipelines/stable_diffusion/latent_upscale
title: Latent upscaler
- local: api/pipelines/stable_diffusion/ldm3d_diffusion
diff --git a/docs/source/en/api/pipelines/cosmos.md b/docs/source/en/api/pipelines/cosmos.md
index 60ecce660303..2302ed2c4a6c 100644
--- a/docs/source/en/api/pipelines/cosmos.md
+++ b/docs/source/en/api/pipelines/cosmos.md
@@ -46,6 +46,20 @@ output = pipe(
output.save("output.png")
```
+## Cosmos2_5_TransferPipeline
+
+[[autodoc]] Cosmos2_5_TransferPipeline
+ - all
+ - __call__
+
+
+## Cosmos2_5_PredictBasePipeline
+
+[[autodoc]] Cosmos2_5_PredictBasePipeline
+ - all
+ - __call__
+
+
## CosmosTextToWorldPipeline
[[autodoc]] CosmosTextToWorldPipeline
@@ -70,12 +84,6 @@ output.save("output.png")
- all
- __call__
-## Cosmos2_5_PredictBasePipeline
-
-[[autodoc]] Cosmos2_5_PredictBasePipeline
- - all
- - __call__
-
## CosmosPipelineOutput
[[autodoc]] pipelines.cosmos.pipeline_output.CosmosPipelineOutput
diff --git a/docs/source/en/api/pipelines/qwenimage.md b/docs/source/en/api/pipelines/qwenimage.md
index ee3dd3b28e4d..c0994c8685d0 100644
--- a/docs/source/en/api/pipelines/qwenimage.md
+++ b/docs/source/en/api/pipelines/qwenimage.md
@@ -29,7 +29,7 @@ Qwen-Image comes in the following variants:
| Qwen-Image-Edit Plus | [Qwen/Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) |
> [!TIP]
-> [Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs.
+> See the [Caching](../../optimization/cache) guide to speed up inference by storing and reusing intermediate outputs.
## LoRA for faster inference
@@ -190,6 +190,12 @@ For detailed benchmark scripts and results, see [this gist](https://gist.github.
- all
- __call__
+## QwenImageLayeredPipeline
+
+[[autodoc]] QwenImageLayeredPipeline
+ - all
+ - __call__
+
## QwenImagePipelineOutput
[[autodoc]] pipelines.qwenimage.pipeline_output.QwenImagePipelineOutput
\ No newline at end of file
diff --git a/docs/source/en/api/pipelines/stable_diffusion/k_diffusion.md b/docs/source/en/api/pipelines/stable_diffusion/k_diffusion.md
deleted file mode 100644
index 75f052b08f13..000000000000
--- a/docs/source/en/api/pipelines/stable_diffusion/k_diffusion.md
+++ /dev/null
@@ -1,30 +0,0 @@
-
-
-> [!WARNING]
-> This pipeline is deprecated but it can still be used. However, we won't test the pipeline anymore and won't accept any changes to it. If you run into any issues, reinstall the last Diffusers version that supported this model.
-
-# K-Diffusion
-
-[k-diffusion](https://github.com/crowsonkb/k-diffusion) is a popular library created by [Katherine Crowson](https://github.com/crowsonkb/). We provide `StableDiffusionKDiffusionPipeline` and `StableDiffusionXLKDiffusionPipeline` that allow you to run Stable DIffusion with samplers from k-diffusion.
-
-Note that most the samplers from k-diffusion are implemented in Diffusers and we recommend using existing schedulers. You can find a mapping between k-diffusion samplers and schedulers in Diffusers [here](https://huggingface.co/docs/diffusers/api/schedulers/overview)
-
-
-## StableDiffusionKDiffusionPipeline
-
-[[autodoc]] StableDiffusionKDiffusionPipeline
-
-
-## StableDiffusionXLKDiffusionPipeline
-
-[[autodoc]] StableDiffusionXLKDiffusionPipeline
\ No newline at end of file
diff --git a/docs/source/en/modular_diffusers/guiders.md b/docs/source/en/modular_diffusers/guiders.md
index 6abe4fad2736..ffe039f41556 100644
--- a/docs/source/en/modular_diffusers/guiders.md
+++ b/docs/source/en/modular_diffusers/guiders.md
@@ -89,10 +89,8 @@ t2i_pipeline.guider
## Changing guider parameters
-The guider parameters can be adjusted with either the [`~ComponentSpec.create`] method or with [`~ModularPipeline.update_components`]. The example below changes the `guidance_scale` value.
+The guider parameters can be adjusted with the [`~ComponentSpec.create`] method and [`~ModularPipeline.update_components`]. The example below changes the `guidance_scale` value.
-
-
```py
guider_spec = t2i_pipeline.get_component_spec("guider")
@@ -100,18 +98,6 @@ guider = guider_spec.create(guidance_scale=10)
t2i_pipeline.update_components(guider=guider)
```
-
-
-
-```py
-guider_spec = t2i_pipeline.get_component_spec("guider")
-guider_spec.config["guidance_scale"] = 10
-t2i_pipeline.update_components(guider=guider_spec)
-```
-
-
-
-
## Uploading custom guiders
Call the [`~utils.PushToHubMixin.push_to_hub`] method on a custom guider to share it to the Hub.
diff --git a/docs/source/en/modular_diffusers/modular_diffusers_states.md b/docs/source/en/modular_diffusers/modular_diffusers_states.md
index eb55b524e491..657b088fe485 100644
--- a/docs/source/en/modular_diffusers/modular_diffusers_states.md
+++ b/docs/source/en/modular_diffusers/modular_diffusers_states.md
@@ -25,9 +25,7 @@ This guide explains how states work and how they connect blocks.
The [`~modular_pipelines.PipelineState`] is a global state container for all blocks. It maintains the complete runtime state of the pipeline and provides a structured way for blocks to read from and write to shared data.
-There are two dict's in [`~modular_pipelines.PipelineState`] for structuring data.
-
-- The `values` dict is a **mutable** state containing a copy of user provided input values and intermediate output values generated by blocks. If a block modifies an `input`, it will be reflected in the `values` dict after calling `set_block_state`.
+[`~modular_pipelines.PipelineState`] stores all data in a `values` dict, which is a **mutable** state containing user provided input values and intermediate output values generated by blocks. If a block modifies an `input`, it will be reflected in the `values` dict after calling `set_block_state`.
```py
PipelineState(
diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md
index 34cd8f72b5b7..e28e13ed5655 100644
--- a/docs/source/en/modular_diffusers/modular_pipeline.md
+++ b/docs/source/en/modular_diffusers/modular_pipeline.md
@@ -12,27 +12,28 @@ specific language governing permissions and limitations under the License.
# ModularPipeline
-[`ModularPipeline`] converts [`~modular_pipelines.ModularPipelineBlocks`]'s into an executable pipeline that loads models and performs the computation steps defined in the block. It is the main interface for running a pipeline and it is very similar to the [`DiffusionPipeline`] API.
+[`ModularPipeline`] converts [`~modular_pipelines.ModularPipelineBlocks`] into an executable pipeline that loads models and performs the computation steps defined in the blocks. It is the main interface for running a pipeline and the API is very similar to [`DiffusionPipeline`] but with a few key differences.
-The main difference is to include an expected `output` argument in the pipeline.
+- **Loading is lazy.** With [`DiffusionPipeline`], [`~DiffusionPipeline.from_pretrained`] creates the pipeline and loads all models at the same time. With [`ModularPipeline`], creating and loading are two separate steps: [`~ModularPipeline.from_pretrained`] reads the configuration and knows where to load each component from, but doesn't actually load the model weights. You load the models later with [`~ModularPipeline.load_components`], which is where you pass loading arguments like `torch_dtype` and `quantization_config`.
+
+- **Two ways to create a pipeline.** You can use [`~ModularPipeline.from_pretrained`] with an existing diffusers model repository — it automatically maps to the default pipeline blocks and then converts to a [`ModularPipeline`] with no extra setup. You can check the [modular_pipelines_directory](https://github.com/huggingface/diffusers/tree/main/src/diffusers/modular_pipelines) to see which models are currently supported. You can also assemble your own pipeline from [`ModularPipelineBlocks`] and convert it with the [`~ModularPipelineBlocks.init_pipeline`] method (see [Creating a pipeline](#creating-a-pipeline) for more details).
+
+- **Running the pipeline is the same.** Once loaded, you call the pipeline with the same arguments you're used to. A single [`ModularPipeline`] can support multiple workflows (text-to-image, image-to-image, inpainting, etc.) when the pipeline blocks use [`AutoPipelineBlocks`](./auto_pipeline_blocks) to automatically select the workflow based on your inputs.
+
+Below are complete examples for text-to-image, image-to-image, and inpainting with SDXL.
```py
import torch
-from diffusers.modular_pipelines import SequentialPipelineBlocks
-from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS
-
-blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS)
-
-modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
-pipeline = blocks.init_pipeline(modular_repo_id)
+from diffusers import ModularPipeline
+pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
pipeline.load_components(torch_dtype=torch.float16)
pipeline.to("cuda")
-image = pipeline(prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", output="images")[0]
+image = pipeline(prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k").images[0]
image.save("modular_t2i_out.png")
```
@@ -41,21 +42,17 @@ image.save("modular_t2i_out.png")
```py
import torch
-from diffusers.modular_pipelines import SequentialPipelineBlocks
-from diffusers.modular_pipelines.stable_diffusion_xl import IMAGE2IMAGE_BLOCKS
-
-blocks = SequentialPipelineBlocks.from_blocks_dict(IMAGE2IMAGE_BLOCKS)
-
-modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
-pipeline = blocks.init_pipeline(modular_repo_id)
+from diffusers import ModularPipeline
+from diffusers.utils import load_image
+pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
pipeline.load_components(torch_dtype=torch.float16)
pipeline.to("cuda")
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-text2img.png"
init_image = load_image(url)
prompt = "a dog catching a frisbee in the jungle"
-image = pipeline(prompt=prompt, image=init_image, strength=0.8, output="images")[0]
+image = pipeline(prompt=prompt, image=init_image, strength=0.8).images[0]
image.save("modular_i2i_out.png")
```
@@ -64,15 +61,10 @@ image.save("modular_i2i_out.png")
```py
import torch
-from diffusers.modular_pipelines import SequentialPipelineBlocks
-from diffusers.modular_pipelines.stable_diffusion_xl import INPAINT_BLOCKS
+from diffusers import ModularPipeline
from diffusers.utils import load_image
-blocks = SequentialPipelineBlocks.from_blocks_dict(INPAINT_BLOCKS)
-
-modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
-pipeline = blocks.init_pipeline(modular_repo_id)
-
+pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
pipeline.load_components(torch_dtype=torch.float16)
pipeline.to("cuda")
@@ -83,276 +75,353 @@ init_image = load_image(img_url)
mask_image = load_image(mask_url)
prompt = "A deep sea diver floating"
-image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=0.85, output="images")[0]
-image.save("moduar_inpaint_out.png")
+image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=0.85).images[0]
+image.save("modular_inpaint_out.png")
```
-This guide will show you how to create a [`ModularPipeline`] and manage the components in it.
+This guide will show you how to create a [`ModularPipeline`], manage its components, and run the pipeline.
-## Adding blocks
+## Creating a pipeline
-Blocks are [`InsertableDict`] objects that can be inserted at specific positions, providing a flexible way to mix-and-match blocks.
+There are two ways to create a [`ModularPipeline`]. Assemble and create a pipeline from [`ModularPipelineBlocks`] with [`~ModularPipelineBlocks.init_pipeline`], or load an existing pipeline with [`~ModularPipeline.from_pretrained`].
-Use [`~modular_pipelines.modular_pipeline_utils.InsertableDict.insert`] on either the block class or `sub_blocks` attribute to add a block.
+You can also initialize a [`ComponentsManager`](./components_manager) to handle device placement and memory management. If you don't need automatic offloading, you can skip this and move the pipeline to your device manually with `pipeline.to("cuda")`.
+
+> [!TIP]
+> Refer to the [ComponentsManager](./components_manager) doc for more details about how it can help manage components across different workflows.
+
+### init_pipeline
+
+[`~ModularPipelineBlocks.init_pipeline`] converts any [`ModularPipelineBlocks`] into a [`ModularPipeline`].
+
+Let's define a minimal block to see how it works:
```py
-# BLOCKS is dict of block classes, you need to add class to it
-BLOCKS.insert("block_name", BlockClass, index)
-# sub_blocks attribute contains instance, add a block instance to the attribute
-t2i_blocks.sub_blocks.insert("block_name", block_instance, index)
+from transformers import CLIPTextModel
+from diffusers.modular_pipelines import (
+ ComponentSpec,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+
+class MyBlock(ModularPipelineBlocks):
+ @property
+ def expected_components(self):
+ return [
+ ComponentSpec(
+ name="text_encoder",
+ type_hint=CLIPTextModel,
+ pretrained_model_name_or_path="openai/clip-vit-large-patch14",
+ ),
+ ]
+
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ return components, state
```
-Use [`~modular_pipelines.modular_pipeline_utils.InsertableDict.pop`] on either the block class or `sub_blocks` attribute to remove a block.
+Call [`~ModularPipelineBlocks.init_pipeline`] to convert it into a pipeline. The `blocks` attribute on the pipeline is the blocks it was created from — it determines the expected inputs, outputs, and computation logic.
```py
-# remove a block class from preset
-BLOCKS.pop("text_encoder")
-# split out a block instance on its own
-text_encoder_block = t2i_blocks.sub_blocks.pop("text_encoder")
+block = MyBlock()
+pipe = block.init_pipeline()
+pipe.blocks
```
-Swap blocks by setting the existing block to the new block.
+```
+MyBlock {
+ "_class_name": "MyBlock",
+ "_diffusers_version": "0.37.0.dev0"
+}
+```
+
+> [!WARNING]
+> Blocks are mutable — you can freely add, remove, or swap blocks before creating a pipeline. However, once a pipeline is created, modifying `pipeline.blocks` won't affect the pipeline because it returns a copy. If you want a different block structure, create a new pipeline after modifying the blocks.
+
+When you call [`~ModularPipelineBlocks.init_pipeline`] without a repository, it uses the `pretrained_model_name_or_path` defined in the block's [`ComponentSpec`] to determine where to load each component from. Printing the pipeline shows the component loading configuration.
```py
-# Replace block class in preset
-BLOCKS["prepare_latents"] = CustomPrepareLatents
-# Replace in sub_blocks attribute using an block instance
-t2i_blocks.sub_blocks["prepare_latents"] = CustomPrepareLatents()
+pipe
+ModularPipeline {
+ "_blocks_class_name": "MyBlock",
+ "_class_name": "ModularPipeline",
+ "_diffusers_version": "0.37.0.dev0",
+ "text_encoder": [
+ null,
+ null,
+ {
+ "pretrained_model_name_or_path": "openai/clip-vit-large-patch14",
+ "revision": null,
+ "subfolder": "",
+ "type_hint": [
+ "transformers",
+ "CLIPTextModel"
+ ],
+ "variant": null
+ }
+ ]
+}
```
-## Creating a pipeline
+If you pass a repository to [`~ModularPipelineBlocks.init_pipeline`], it overrides the loading path by matching your block's components against the pipeline config in that repository (`model_index.json` or `modular_model_index.json`).
-There are two ways to create a [`ModularPipeline`]. Assemble and create a pipeline from [`ModularPipelineBlocks`] or load an existing pipeline with [`~ModularPipeline.from_pretrained`].
+In the example below, the `pretrained_model_name_or_path` will be updated to `"stabilityai/stable-diffusion-xl-base-1.0"`.
-You should also initialize a [`ComponentsManager`] to handle device placement and memory and component management.
+```py
+pipe = block.init_pipeline("stabilityai/stable-diffusion-xl-base-1.0")
+pipe
+ModularPipeline {
+ "_blocks_class_name": "MyBlock",
+ "_class_name": "ModularPipeline",
+ "_diffusers_version": "0.37.0.dev0",
+ "text_encoder": [
+ null,
+ null,
+ {
+ "pretrained_model_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0",
+ "revision": null,
+ "subfolder": "text_encoder",
+ "type_hint": [
+ "transformers",
+ "CLIPTextModel"
+ ],
+ "variant": null
+ }
+ ]
+}
+```
-> [!TIP]
-> Refer to the [ComponentsManager](./components_manager) doc for more details about how it can help manage components across different workflows.
+If a component in your block doesn't exist in the repository, it remains `null` and is skipped during [`~ModularPipeline.load_components`].
-
-
+### from_pretrained
-Use the [`~ModularPipelineBlocks.init_pipeline`] method to create a [`ModularPipeline`] from the component and configuration specifications. This method loads the *specifications* from a `modular_model_index.json` file, but it doesn't load the *models* yet.
+[`~ModularPipeline.from_pretrained`] is a convenient way to create a [`ModularPipeline`] without defining blocks yourself.
-```py
-from diffusers import ComponentsManager
-from diffusers.modular_pipelines import SequentialPipelineBlocks
-from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS
+It works with three types of repositories.
+
+**A regular diffusers repository.** Pass any supported model repository and it automatically maps to the default pipeline blocks. Currently supported models include SDXL, Wan, Qwen, Z-Image, Flux, and Flux2.
-t2i_blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS)
+```py
+from diffusers import ModularPipeline, ComponentsManager
-modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
components = ComponentsManager()
-t2i_pipeline = t2i_blocks.init_pipeline(modular_repo_id, components_manager=components)
+pipeline = ModularPipeline.from_pretrained(
+ "stabilityai/stable-diffusion-xl-base-1.0", components_manager=components
+)
```
-
-
-
-The [`~ModularPipeline.from_pretrained`] method creates a [`ModularPipeline`] from a modular repository on the Hub.
+**A modular repository.** These repositories contain a `modular_model_index.json` that specifies where to load each component from — the components can come from different repositories and the modular repository itself may not contain any model weights. For example, [diffusers/flux2-bnb-4bit-modular](https://huggingface.co/diffusers/flux2-bnb-4bit-modular) loads a quantized transformer from one repository and the remaining components from another. See [Modular repository](#modular-repository) for more details on the format.
```py
from diffusers import ModularPipeline, ComponentsManager
components = ComponentsManager()
-pipeline = ModularPipeline.from_pretrained("YiYiXu/modular-loader-t2i-0704", components_manager=components)
+pipeline = ModularPipeline.from_pretrained(
+ "diffusers/flux2-bnb-4bit-modular", components_manager=components
+)
```
-Add the `trust_remote_code` argument to load a custom [`ModularPipeline`].
+**A modular repository with custom code.** Some repositories include custom pipeline blocks alongside the loading configuration. Add `trust_remote_code=True` to load them. See [Custom blocks](./custom_blocks) for how to create your own.
```py
from diffusers import ModularPipeline, ComponentsManager
components = ComponentsManager()
-modular_repo_id = "YiYiXu/modular-diffdiff-0704"
-diffdiff_pipeline = ModularPipeline.from_pretrained(modular_repo_id, trust_remote_code=True, components_manager=components)
+pipeline = ModularPipeline.from_pretrained(
+ "diffusers/Florence2-image-Annotator", trust_remote_code=True, components_manager=components
+)
```
-
-
-
## Loading components
-A [`ModularPipeline`] doesn't automatically instantiate with components. It only loads the configuration and component specifications. You can load all components with [`~ModularPipeline.load_components`] or only load specific components with [`~ModularPipeline.load_components`].
+A [`ModularPipeline`] doesn't automatically instantiate with components. It only loads the configuration and component specifications. You can load components with [`~ModularPipeline.load_components`].
-
-
+This will load all the components that have a valid loading spec.
```py
import torch
-t2i_pipeline.load_components(torch_dtype=torch.float16)
-t2i_pipeline.to("cuda")
+pipeline.load_components(torch_dtype=torch.float16)
```
-
-
-
-The example below only loads the UNet and VAE.
+You can also load specific components by name. The example below only loads the `text_encoder`.
```py
-import torch
-
-t2i_pipeline.load_components(names=["unet", "vae"], torch_dtype=torch.float16)
+pipeline.load_components(names=["text_encoder"], torch_dtype=torch.float16)
```
-
-
-
-Print the pipeline to inspect the loaded pretrained components.
+After loading, printing the pipeline shows which components are loaded — the first two fields change from `null` to the component's library and class.
```py
-t2i_pipeline
+pipeline
```
-This should match the `modular_model_index.json` file from the modular repository a pipeline is initialized from. If a pipeline doesn't need a component, it won't be included even if it exists in the modular repository.
-
-To modify where components are loaded from, edit the `modular_model_index.json` file in the repository and change it to your desired loading path. The example below loads a UNet from a different repository.
-
-```json
-# original
-"unet": [
- null, null,
- {
- "repo": "stabilityai/stable-diffusion-xl-base-1.0",
- "subfolder": "unet",
- "variant": "fp16"
- }
+```
+# text_encoder is loaded - shows library and class
+"text_encoder": [
+ "transformers",
+ "CLIPTextModel",
+ { ... }
]
-# modified
+# unet is not loaded yet - still null
"unet": [
- null, null,
- {
- "repo": "RunDiffusion/Juggernaut-XL-v9",
- "subfolder": "unet",
- "variant": "fp16"
- }
+ null,
+ null,
+ { ... }
]
```
-### Component loading status
-
-The pipeline properties below provide more information about which components are loaded.
-
-Use `component_names` to return all expected components.
+Loading keyword arguments like `torch_dtype`, `variant`, `revision`, and `quantization_config` are passed through to `from_pretrained()` for each component. You can pass a single value to apply to all components, or a dict to set per-component values.
```py
-t2i_pipeline.component_names
-['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'guider', 'scheduler', 'unet', 'vae', 'image_processor']
+# apply bfloat16 to all components
+pipeline.load_components(torch_dtype=torch.bfloat16)
+
+# different dtypes per component
+pipeline.load_components(torch_dtype={"transformer": torch.bfloat16, "default": torch.float32})
```
-Use `null_component_names` to return components that aren't loaded yet. Load these components with [`~ModularPipeline.from_pretrained`].
+[`~ModularPipeline.load_components`] only loads components that haven't been loaded yet and have a valid loading spec. This means if you've already set a component on the pipeline, calling [`~ModularPipeline.load_components`] again won't reload it.
-```py
-t2i_pipeline.null_component_names
-['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'scheduler']
-```
+## Updating components
-Use `pretrained_component_names` to return components that will be loaded from pretrained models.
+[`~ModularPipeline.update_components`] replaces a component on the pipeline with a new one. When a component is updated, the loading specifications are also updated in the pipeline config and [`~ModularPipeline.load_components`] will skip it on subsequent calls.
-```py
-t2i_pipeline.pretrained_component_names
-['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'scheduler', 'unet', 'vae']
-```
+### From AutoModel
-Use `config_component_names` to return components that are created with the default config (not loaded from a modular repository). Components from a config aren't included because they are already initialized during pipeline creation. This is why they aren't listed in `null_component_names`.
+You can pass a model object loaded with `AutoModel.from_pretrained()`. Models loaded this way are automatically tagged with their loading information.
```py
-t2i_pipeline.config_component_names
-['guider', 'image_processor']
+from diffusers import AutoModel
+
+unet = AutoModel.from_pretrained(
+ "RunDiffusion/Juggernaut-XL-v9", subfolder="unet", variant="fp16", torch_dtype=torch.float16
+)
+pipeline.update_components(unet=unet)
```
-## Updating components
+### From ComponentSpec
-Components may be updated depending on whether it is a *pretrained component* or a *config component*.
+Use [`~ModularPipeline.get_component_spec`] to get a copy of the current component specification, modify it, and load a new component.
-> [!WARNING]
-> A component may change from pretrained to config when updating a component. The component type is initially defined in a block's `expected_components` field.
+```py
+unet_spec = pipeline.get_component_spec("unet")
+
+# modify to load from a different repository
+unet_spec.pretrained_model_name_or_path = "RunDiffusion/Juggernaut-XL-v9"
-A pretrained component is updated with [`ComponentSpec`] whereas a config component is updated by eihter passing the object directly or with [`ComponentSpec`].
+# load and update
+unet = unet_spec.load(torch_dtype=torch.float16)
+pipeline.update_components(unet=unet)
+```
-The [`ComponentSpec`] shows `default_creation_method="from_pretrained"` for a pretrained component shows `default_creation_method="from_config` for a config component.
+You can also create a [`ComponentSpec`] from scratch.
-To update a pretrained component, create a [`ComponentSpec`] with the name of the component and where to load it from. Use the [`~ComponentSpec.load`] method to load the component.
+Not all components are loaded from pretrained weights — some are created from a config (listed under `pipeline.config_component_names`). For these, use [`~ComponentSpec.create`] instead of [`~ComponentSpec.load`].
```py
-from diffusers import ComponentSpec, UNet2DConditionModel
-
-unet_spec = ComponentSpec(name="unet",type_hint=UNet2DConditionModel, repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", variant="fp16")
-unet = unet_spec.load(torch_dtype=torch.float16)
+guider_spec = pipeline.get_component_spec("guider")
+guider_spec.config = {"guidance_scale": 5.0}
+guider = guider_spec.create()
+pipeline.update_components(guider=guider)
```
-The [`~ModularPipeline.update_components`] method replaces the component with a new one.
+Or simply pass the object directly.
```py
-t2i_pipeline.update_components(unet=unet2)
+from diffusers.guiders import ClassifierFreeGuidance
+
+guider = ClassifierFreeGuidance(guidance_scale=5.0)
+pipeline.update_components(guider=guider)
```
-When a component is updated, the loading specifications are also updated in the pipeline config.
+See the [Guiders](./guiders) guide for more details on available guiders and how to configure them.
-### Component extraction and modification
+## Splitting a pipeline into stages
-When you use [`~ComponentSpec.load`], the new component maintains its loading specifications. This makes it possible to extract the specification and recreate the component.
+Since blocks are composable, you can take a pipeline apart and reconstruct it into separate pipelines for each stage. The example below shows how we can separate the text encoder block from the rest of the pipeline, so you can encode the prompt independently and pass the embeddings to the main pipeline.
```py
-spec = ComponentSpec.from_component("unet", unet2)
-spec
-ComponentSpec(name='unet', type_hint=, description=None, config=None, repo='stabilityai/stable-diffusion-xl-base-1.0', subfolder='unet', variant='fp16', revision=None, default_creation_method='from_pretrained')
-unet2_recreated = spec.load(torch_dtype=torch.float16)
-```
+from diffusers import ModularPipeline, ComponentsManager
+import torch
-The [`~ModularPipeline.get_component_spec`] method gets a copy of the current component specification to modify or update.
+device = "cuda"
+dtype = torch.bfloat16
+repo_id = "black-forest-labs/FLUX.2-klein-4B"
-```py
-unet_spec = t2i_pipeline.get_component_spec("unet")
-unet_spec
-ComponentSpec(
- name='unet',
- type_hint=,
- pretrained_model_name_or_path='RunDiffusion/Juggernaut-XL-v9',
- subfolder='unet',
- variant='fp16',
- default_creation_method='from_pretrained'
-)
+# get the blocks and separate out the text encoder
+blocks = ModularPipeline.from_pretrained(repo_id).blocks
+text_block = blocks.sub_blocks.pop("text_encoder")
-# modify to load from a different repository
-unet_spec.pretrained_model_name_or_path = "stabilityai/stable-diffusion-xl-base-1.0"
+# use ComponentsManager to handle offloading across multiple pipelines
+manager = ComponentsManager()
+manager.enable_auto_cpu_offload(device=device)
-# load component with modified spec
-unet = unet_spec.load(torch_dtype=torch.float16)
+# create separate pipelines for each stage
+text_encoder_pipeline = text_block.init_pipeline(repo_id, components_manager=manager)
+pipeline = blocks.init_pipeline(repo_id, components_manager=manager)
+
+# encode text
+text_encoder_pipeline.load_components(torch_dtype=dtype)
+text_embeddings = text_encoder_pipeline(prompt="a cat").get_by_kwargs("denoiser_input_fields")
+
+# denoise and decode
+pipeline.load_components(torch_dtype=dtype)
+output = pipeline(
+ **text_embeddings,
+ num_inference_steps=4,
+).images[0]
```
+[`ComponentsManager`] handles memory across multiple pipelines. Unlike the offloading strategies in [`DiffusionPipeline`] that follow a fixed order, [`ComponentsManager`] makes offloading decisions dynamically each time a model forward pass runs, based on the current memory situation. This means it works regardless of how many pipelines you create or what order you run them in. See the [ComponentsManager](./components_manager) guide for more details.
+
+If pipeline stages share components (e.g., the same VAE used for encoding and decoding), you can use [`~ModularPipeline.update_components`] to pass an already-loaded component to another pipeline instead of loading it again.
+
## Modular repository
A repository is required if the pipeline blocks use *pretrained components*. The repository supplies loading specifications and metadata.
-[`ModularPipeline`] specifically requires *modular repositories* (see [example repository](https://huggingface.co/YiYiXu/modular-diffdiff)) which are more flexible than a typical repository. It contains a `modular_model_index.json` file containing the following 3 elements.
+[`ModularPipeline`] works with regular diffusers repositories out of the box. However, you can also create a *modular repository* for more flexibility. A modular repository contains a `modular_model_index.json` file containing the following 3 elements.
-- `library` and `class` shows which library the component was loaded from and it's class. If `null`, the component hasn't been loaded yet.
+- `library` and `class` shows which library the component was loaded from and its class. If `null`, the component hasn't been loaded yet.
- `loading_specs_dict` contains the information required to load the component such as the repository and subfolder it is loaded from.
-Unlike standard repositories, a modular repository can fetch components from different repositories based on the `loading_specs_dict`. Components don't need to exist in the same repository.
+The key advantage of a modular repository is that components can be loaded from different repositories. For example, [diffusers/flux2-bnb-4bit-modular](https://huggingface.co/diffusers/flux2-bnb-4bit-modular) loads a quantized transformer from `diffusers/FLUX.2-dev-bnb-4bit` while loading the remaining components from `black-forest-labs/FLUX.2-dev`.
+
+To convert a regular diffusers repository into a modular one, create the pipeline using the regular repository, and then push to the Hub. The saved repository will contain a `modular_model_index.json` with all the loading specifications.
-A modular repository may contain custom code for loading a [`ModularPipeline`]. This allows you to use specialized blocks that aren't native to Diffusers.
+```py
+from diffusers import ModularPipeline
+
+# load from a regular repo
+pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
+
+# push as a modular repository
+pipeline.save_pretrained("local/path", repo_id="my-username/sdxl-modular", push_to_hub=True)
+```
+
+A modular repository can also include custom pipeline blocks as Python code. This allows you to share specialized blocks that aren't native to Diffusers. For example, [diffusers/Florence2-image-Annotator](https://huggingface.co/diffusers/Florence2-image-Annotator) contains custom blocks alongside the loading configuration:
```
-modular-diffdiff-0704/
+Florence2-image-Annotator/
├── block.py # Custom pipeline blocks implementation
├── config.json # Pipeline configuration and auto_map
+├── mellon_config.json # UI configuration for Mellon
└── modular_model_index.json # Component loading specifications
```
-The [config.json](https://huggingface.co/YiYiXu/modular-diffdiff-0704/blob/main/config.json) file contains an `auto_map` key that points to where a custom block is defined in `block.py`.
+The `config.json` file contains an `auto_map` key that tells [`ModularPipeline`] where to find the custom blocks:
```json
{
- "_class_name": "DiffDiffBlocks",
+ "_class_name": "Florence2AnnotatorBlocks",
"auto_map": {
- "ModularPipelineBlocks": "block.DiffDiffBlocks"
+ "ModularPipelineBlocks": "block.Florence2AnnotatorBlocks"
}
}
```
+
+Load custom code repositories with `trust_remote_code=True` as shown in [from_pretrained](#from_pretrained). See [Custom blocks](./custom_blocks) for how to create and share your own.
\ No newline at end of file
diff --git a/docs/source/en/modular_diffusers/pipeline_block.md b/docs/source/en/modular_diffusers/pipeline_block.md
index 06c115e1fb52..612736f3d307 100644
--- a/docs/source/en/modular_diffusers/pipeline_block.md
+++ b/docs/source/en/modular_diffusers/pipeline_block.md
@@ -25,81 +25,151 @@ This guide will show you how to create a [`~modular_pipelines.ModularPipelineBlo
A [`~modular_pipelines.ModularPipelineBlocks`] requires `inputs`, and `intermediate_outputs`.
-- `inputs` are values provided by a user and retrieved from the [`~modular_pipelines.PipelineState`]. This is useful because some workflows resize an image, but the original image is still required. The [`~modular_pipelines.PipelineState`] maintains the original image.
+- `inputs` are values a block reads from the [`~modular_pipelines.PipelineState`] to perform its computation. These can be values provided by a user (like a prompt or image) or values produced by a previous block (like encoded `image_latents`).
Use `InputParam` to define `inputs`.
- ```py
- from diffusers.modular_pipelines import InputParam
-
- user_inputs = [
- InputParam(name="image", type_hint="PIL.Image", description="raw input image to process")
- ]
- ```
+```py
+class ImageEncodeStep(ModularPipelineBlocks):
+ ...
+
+ @property
+ def inputs(self):
+ return [
+ InputParam(name="image", type_hint="PIL.Image", required=True, description="raw input image to process"),
+ ]
+ ...
+```
- `intermediate_outputs` are new values created by a block and added to the [`~modular_pipelines.PipelineState`]. The `intermediate_outputs` are available as `inputs` for subsequent blocks or available as the final output from running the pipeline.
Use `OutputParam` to define `intermediate_outputs`.
- ```py
- from diffusers.modular_pipelines import OutputParam
+```py
+class ImageEncodeStep(ModularPipelineBlocks):
+ ...
- user_intermediate_outputs = [
- OutputParam(name="image_latents", description="latents representing the image")
- ]
- ```
+ @property
+ def intermediate_outputs(self):
+ return [
+ OutputParam(name="image_latents", description="latents representing the image"),
+ ]
+
+ ...
+```
The intermediate inputs and outputs share data to connect blocks. They are accessible at any point, allowing you to track the workflow's progress.
+## Components and configs
+
+The components and pipeline-level configs a block needs are specified in [`ComponentSpec`] and [`~modular_pipelines.ConfigSpec`].
+
+- [`ComponentSpec`] contains the expected components used by a block. You need the `name` of the component and ideally a `type_hint` that specifies exactly what the component is.
+- [`~modular_pipelines.ConfigSpec`] contains pipeline-level settings that control behavior across all blocks.
+
+```py
+class ImageEncodeStep(ModularPipelineBlocks):
+ ...
+
+ @property
+ def expected_components(self):
+ return [
+ ComponentSpec(name="vae", type_hint=AutoencoderKL),
+ ]
+
+ @property
+ def expected_configs(self):
+ return [
+ ConfigSpec("force_zeros_for_empty_prompt", True),
+ ]
+
+ ...
+```
+
+When the blocks are converted into a pipeline, the components become available to the block as the first argument in `__call__`.
+
## Computation logic
The computation a block performs is defined in the `__call__` method and it follows a specific structure.
-1. Retrieve the [`~modular_pipelines.BlockState`] to get a local view of the `inputs`
+1. Retrieve the [`~modular_pipelines.BlockState`] to get a local view of the `inputs`.
2. Implement the computation logic on the `inputs`.
3. Update [`~modular_pipelines.PipelineState`] to push changes from the local [`~modular_pipelines.BlockState`] back to the global [`~modular_pipelines.PipelineState`].
4. Return the components and state which becomes available to the next block.
```py
-def __call__(self, components, state):
- # Get a local view of the state variables this block needs
- block_state = self.get_block_state(state)
+class ImageEncodeStep(ModularPipelineBlocks):
+
+ def __call__(self, components, state):
+ # Get a local view of the state variables this block needs
+ block_state = self.get_block_state(state)
- # Your computation logic here
- # block_state contains all your inputs
- # Access them like: block_state.image, block_state.processed_image
+ # Your computation logic here
+ # block_state contains all your inputs
+ # Access them like: block_state.image, block_state.processed_image
- # Update the pipeline state with your updated block_states
- self.set_block_state(state, block_state)
- return components, state
+ # Update the pipeline state with your updated block_states
+ self.set_block_state(state, block_state)
+ return components, state
```
-### Components and configs
+## Putting it all together
-The components and pipeline-level configs a block needs are specified in [`ComponentSpec`] and [`~modular_pipelines.ConfigSpec`].
+Here is the complete block with all the pieces connected.
-- [`ComponentSpec`] contains the expected components used by a block. You need the `name` of the component and ideally a `type_hint` that specifies exactly what the component is.
-- [`~modular_pipelines.ConfigSpec`] contains pipeline-level settings that control behavior across all blocks.
+```py
+from diffusers import ComponentSpec, AutoencoderKL
+from diffusers.modular_pipelines import InputParam, ModularPipelineBlocks, OutputParam
+
+
+class ImageEncodeStep(ModularPipelineBlocks):
+
+ @property
+ def description(self):
+ return "Encode an image into latent space."
+
+ @property
+ def expected_components(self):
+ return [
+ ComponentSpec(name="vae", type_hint=AutoencoderKL),
+ ]
+
+ @property
+ def inputs(self):
+ return [
+ InputParam(name="image", type_hint="PIL.Image", required=True, description="raw input image to process"),
+ ]
+
+ @property
+ def intermediate_outputs(self):
+ return [
+ OutputParam(name="image_latents", type_hint="torch.Tensor", description="latents representing the image"),
+ ]
+
+ def __call__(self, components, state):
+ block_state = self.get_block_state(state)
+ block_state.image_latents = components.vae.encode(block_state.image)
+ self.set_block_state(state, block_state)
+ return components, state
+```
+
+Every block has a `doc` property that is automatically generated from the properties you defined above. It provides a summary of the block's description, components, inputs, and outputs.
```py
-from diffusers import ComponentSpec, ConfigSpec
+block = ImageEncoderStep()
+print(block.doc)
+class ImageEncodeStep
-expected_components = [
- ComponentSpec(name="unet", type_hint=UNet2DConditionModel),
- ComponentSpec(name="scheduler", type_hint=EulerDiscreteScheduler)
-]
+ Encode an image into latent space.
-expected_config = [
- ConfigSpec("force_zeros_for_empty_prompt", True)
-]
-```
+ Components:
+ vae (`AutoencoderKL`)
-When the blocks are converted into a pipeline, the components become available to the block as the first argument in `__call__`.
+ Inputs:
+ image (`PIL.Image`):
+ raw input image to process
-```py
-def __call__(self, components, state):
- # Access components using dot notation
- unet = components.unet
- vae = components.vae
- scheduler = components.scheduler
-```
+ Outputs:
+ image_latents (`torch.Tensor`):
+ latents representing the image
+```
\ No newline at end of file
diff --git a/docs/source/en/modular_diffusers/quickstart.md b/docs/source/en/modular_diffusers/quickstart.md
index 5a455f0b3093..884495e107b4 100644
--- a/docs/source/en/modular_diffusers/quickstart.md
+++ b/docs/source/en/modular_diffusers/quickstart.md
@@ -39,17 +39,44 @@ image
[`~ModularPipeline.from_pretrained`] uses lazy loading - it reads the configuration to learn where to load each component from, but doesn't actually load the model weights until you call [`~ModularPipeline.load_components`]. This gives you control over when and how components are loaded.
> [!TIP]
-> [`ComponentsManager`] with `enable_auto_cpu_offload` automatically moves models between CPU and GPU as needed, reducing memory usage for large models like Qwen-Image. Learn more in the [ComponentsManager](./components_manager) guide.
+> `ComponentsManager` with `enable_auto_cpu_offload` automatically moves models between CPU and GPU as needed, reducing memory usage for large models like Qwen-Image. Learn more in the [ComponentsManager](./components_manager) guide.
+>
+> If you don't need offloading, remove the `components_manager` argument and move the pipeline to your device manually with `to("cuda")`.
Learn more about creating and loading pipelines in the [Creating a pipeline](https://huggingface.co/docs/diffusers/modular_diffusers/modular_pipeline#creating-a-pipeline) and [Loading components](https://huggingface.co/docs/diffusers/modular_diffusers/modular_pipeline#loading-components) guides.
## Understand the structure
-A [`ModularPipeline`] has two parts:
-- **State**: the loaded components (models, schedulers, processors) and configuration
-- **Definition**: the [`ModularPipelineBlocks`] that specify inputs, outputs, expected components and computation logic
+A [`ModularPipeline`] has two parts: a **definition** (the blocks) and a **state** (the loaded components and configs).
-The blocks define *what* the pipeline does. Access them through `pipe.blocks`.
+Print the pipeline to see its state — the components and their loading status and configuration.
+```py
+print(pipe)
+```
+```
+QwenImageModularPipeline {
+ "_blocks_class_name": "QwenImageAutoBlocks",
+ "_class_name": "QwenImageModularPipeline",
+ "_diffusers_version": "0.37.0.dev0",
+ "transformer": [
+ "diffusers",
+ "QwenImageTransformer2DModel",
+ {
+ "pretrained_model_name_or_path": "Qwen/Qwen-Image",
+ "revision": null,
+ "subfolder": "transformer",
+ "type_hint": [
+ "diffusers",
+ "QwenImageTransformer2DModel"
+ ],
+ "variant": null
+ }
+ ],
+ ...
+}
+```
+
+Access the definition through `pipe.blocks` — this is the [`~modular_pipelines.ModularPipelineBlocks`] that defines the pipeline's workflows, inputs, outputs, and computation logic.
```py
print(pipe.blocks)
```
@@ -87,7 +114,8 @@ The output returns:
### Workflows
-`QwenImageAutoBlocks` is a [`ConditionalPipelineBlocks`], so this pipeline supports multiple workflows and adapts its behavior based on the inputs you provide. For example, if you pass `image` to the pipeline, it runs an image-to-image workflow instead of text-to-image. Let's see this in action with an example.
+This pipeline supports multiple workflows and adapts its behavior based on the inputs you provide. For example, if you pass `image` to the pipeline, it runs an image-to-image workflow instead of text-to-image. Learn more about how this works under the hood in the [AutoPipelineBlocks](https://huggingface.co/docs/diffusers/modular_diffusers/auto_pipeline_blocks) guide.
+
```py
from diffusers.utils import load_image
@@ -99,20 +127,21 @@ image = pipe(
).images[0]
```
-Use `get_workflow()` to extract the blocks for a specific workflow. Pass the workflow name (e.g., `"image2image"`, `"inpainting"`, `"controlnet_text2image"`) to get only the blocks relevant to that workflow.
+Use `get_workflow()` to extract the blocks for a specific workflow. Pass the workflow name (e.g., `"image2image"`, `"inpainting"`, `"controlnet_text2image"`) to get only the blocks relevant to that workflow. This is useful when you want to customize or debug a specific workflow. You can check `pipe.blocks.available_workflows` to see all available workflows.
```py
img2img_blocks = pipe.blocks.get_workflow("image2image")
```
-Conditional blocks are convenient for users, but their conditional logic adds complexity when customizing or debugging. Extracting a workflow gives you the specific blocks relevant to your workflow, making it easier to work with. Learn more in the [AutoPipelineBlocks](https://huggingface.co/docs/diffusers/modular_diffusers/auto_pipeline_blocks) guide.
### Sub-blocks
Blocks can contain other blocks. `pipe.blocks` gives you the top-level block definition (here, `QwenImageAutoBlocks`), while `sub_blocks` lets you access the smaller blocks inside it.
-`QwenImageAutoBlocks` is composed of: `text_encoder`, `vae_encoder`, `controlnet_vae_encoder`, `denoise`, and `decode`. Access them through the `sub_blocks` property.
+`QwenImageAutoBlocks` is composed of: `text_encoder`, `vae_encoder`, `controlnet_vae_encoder`, `denoise`, and `decode`.
-The `doc` property is useful for seeing the full documentation of any block, including its inputs, outputs, and components.
+These sub-blocks run one after another and data flows linearly from one block to the next — each block's `intermediate_outputs` become available as `inputs` to the next block. This is how [`SequentialPipelineBlocks`](./sequential_pipeline_blocks) work.
+
+You can access them through the `sub_blocks` property. The `doc` property is useful for seeing the full documentation of any block, including its inputs, outputs, and components.
```py
vae_encoder_block = pipe.blocks.sub_blocks["vae_encoder"]
print(vae_encoder_block.doc)
@@ -165,7 +194,7 @@ class CannyBlock
Canny map for input image
```
-UUse `get_workflow` to extract the ControlNet workflow from [`QwenImageAutoBlocks`].
+Use `get_workflow` to extract the ControlNet workflow from [`QwenImageAutoBlocks`].
```py
# Get the controlnet workflow that we want to work with
blocks = pipe.blocks.get_workflow("controlnet_text2image")
@@ -182,9 +211,8 @@ class SequentialPipelineBlocks
...
```
-The extracted workflow is a [`SequentialPipelineBlocks`](./sequential_pipeline_blocks) - a multi-block type where blocks run one after another and data flows linearly from one block to the next. Each block's `intermediate_outputs` become available as `inputs` to subsequent blocks.
-Currently this workflow requires `control_image` as input. Let's insert the canny block at the beginning so the pipeline accepts a regular image instead.
+The extracted workflow is a [`SequentialPipelineBlocks`](./sequential_pipeline_blocks) and it currently requires `control_image` as input. Insert the canny block at the beginning so the pipeline accepts a regular image instead.
```py
# Insert canny at the beginning
blocks.sub_blocks.insert("canny", canny_block, 0)
@@ -211,7 +239,7 @@ class SequentialPipelineBlocks
Now the pipeline takes `image` as input instead of `control_image`. Because blocks in a sequence share data automatically, the canny block's output (`control_image`) flows to the denoise block that needs it, and the canny block's input (`image`) becomes a pipeline input since no earlier block provides it.
-Create a pipeline from the modified blocks and load a ControlNet model.
+Create a pipeline from the modified blocks and load a ControlNet model. The ControlNet isn't part of the original model repository, so load it separately and add it with [`~ModularPipeline.update_components`].
```py
pipeline = blocks.init_pipeline("Qwen/Qwen-Image", components_manager=manager)
@@ -241,6 +269,16 @@ output
## Next steps
+
+
+Understand the core building blocks of Modular Diffusers:
+
+- [ModularPipelineBlocks](./pipeline_block): The basic unit for defining a step in a pipeline.
+- [SequentialPipelineBlocks](./sequential_pipeline_blocks): Chain blocks to run in sequence.
+- [AutoPipelineBlocks](./auto_pipeline_blocks): Create pipelines that support multiple workflows.
+- [States](./modular_diffusers_states): How data is shared between blocks.
+
+
Learn how to create your own blocks with custom logic in the [Building Custom Blocks](./custom_blocks) guide.
diff --git a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md
index f1549a26b86f..1bd67e17b8bf 100644
--- a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md
+++ b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md
@@ -91,23 +91,42 @@ class ImageEncoderBlock(ModularPipelineBlocks):
-Connect the two blocks by defining an [`InsertableDict`] to map the block names to the block instances. Blocks are executed in the order they're registered in `blocks_dict`.
-
-Use [`~modular_pipelines.SequentialPipelineBlocks.from_blocks_dict`] to create a [`~modular_pipelines.SequentialPipelineBlocks`].
+Connect the two blocks by defining a [`~modular_pipelines.SequentialPipelineBlocks`]. List the block instances in `block_classes` and their corresponding names in `block_names`. The blocks are executed in the order they appear in `block_classes`, and data flows from one block to the next through [`~modular_pipelines.PipelineState`].
```py
-from diffusers.modular_pipelines import SequentialPipelineBlocks, InsertableDict
-
-blocks_dict = InsertableDict()
-blocks_dict["input"] = input_block
-blocks_dict["image_encoder"] = image_encoder_block
+class ImageProcessingStep(SequentialPipelineBlocks):
+ """
+ # auto_docstring
+ """
+ model_name = "my_model"
+ block_classes = [InputBlock(), ImageEncoderBlock()]
+ block_names = ["input", "image_encoder"]
-blocks = SequentialPipelineBlocks.from_blocks_dict(blocks_dict)
+ @property
+ def description(self):
+ return (
+ "Process text prompts and images for the pipeline. It:\n"
+ " - Determines the batch size from the prompts.\n"
+ " - Encodes the image into latent space."
+ )
```
-Inspect the sub-blocks in [`~modular_pipelines.SequentialPipelineBlocks`] by calling `blocks`, and for more details about the inputs and outputs, access the `docs` attribute.
+When you create a [`~modular_pipelines.SequentialPipelineBlocks`], properties like `inputs`, `intermediate_outputs`, and `expected_components` are automatically aggregated from the sub-blocks, so there is no need to define them again.
+
+There are a few properties you should set:
+
+- `description`: We recommend adding a description for the assembled block to explain what the combined step does.
+- `model_name`: This is automatically derived from the sub-blocks but isn't always correct, so you may need to override it.
+- `outputs`: By default this is the same as `intermediate_outputs`, but you can manually set it to control which values appear in the doc. This is useful for showing only the final outputs instead of all intermediate values.
+
+These properties, together with the aggregated `inputs`, `intermediate_outputs`, and `expected_components`, are used to automatically generate the `doc` property.
+
+
+Print the `ImageProcessingStep` block to inspect its sub-blocks, and use `doc` for a full summary of the block's inputs, outputs, and components.
+
```py
+blocks = ImageProcessingStep()
print(blocks)
print(blocks.doc)
-```
+```
\ No newline at end of file
diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md
index bdaa2ae8ffff..c9a341df40c5 100644
--- a/docs/source/en/training/distributed_inference.md
+++ b/docs/source/en/training/distributed_inference.md
@@ -111,7 +111,7 @@ if __name__ == "__main__":
Call `torchrun` to run the inference script and use the `--nproc_per_node` argument to set the number of GPUs to use.
```bash
-torchrun run_distributed.py --nproc_per_node=2
+torchrun --nproc_per_node=2 run_distributed.py
```
## device_map
diff --git a/docs/source/en/using-diffusers/automodel.md b/docs/source/en/using-diffusers/automodel.md
index 957cbd17e3f7..d8cea79a10c9 100644
--- a/docs/source/en/using-diffusers/automodel.md
+++ b/docs/source/en/using-diffusers/automodel.md
@@ -29,8 +29,31 @@ text_encoder = AutoModel.from_pretrained(
)
```
+## Custom models
+
[`AutoModel`] also loads models from the [Hub](https://huggingface.co/models) that aren't included in Diffusers. Set `trust_remote_code=True` in [`AutoModel.from_pretrained`] to load custom models.
+A custom model repository needs a Python module with the model class, and a `config.json` with an `auto_map` entry that maps `"AutoModel"` to `"module_file.ClassName"`.
+
+```
+custom/custom-transformer-model/
+├── config.json
+├── my_model.py
+└── diffusion_pytorch_model.safetensors
+```
+
+The `config.json` includes the `auto_map` field pointing to the custom class.
+
+```json
+{
+ "auto_map": {
+ "AutoModel": "my_model.MyCustomModel"
+ }
+}
+```
+
+Then load it with `trust_remote_code=True`.
+
```py
import torch
from diffusers import AutoModel
@@ -40,7 +63,39 @@ transformer = AutoModel.from_pretrained(
)
```
+For a real-world example, [Overworld/Waypoint-1-Small](https://huggingface.co/Overworld/Waypoint-1-Small/tree/main/transformer) hosts a custom `WorldModel` class across several modules in its `transformer` subfolder.
+
+```
+transformer/
+├── config.json # auto_map: "model.WorldModel"
+├── model.py
+├── attn.py
+├── nn.py
+├── cache.py
+├── quantize.py
+├── __init__.py
+└── diffusion_pytorch_model.safetensors
+```
+
+```py
+import torch
+from diffusers import AutoModel
+
+transformer = AutoModel.from_pretrained(
+ "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda"
+)
+```
+
If the custom model inherits from the [`ModelMixin`] class, it gets access to the same features as Diffusers model classes, like [regional compilation](../optimization/fp16#regional-compilation) and [group offloading](../optimization/memory#group-offloading).
+> [!WARNING]
+> As a precaution with `trust_remote_code=True`, pass a commit hash to the `revision` argument in [`AutoModel.from_pretrained`] to make sure the code hasn't been updated with new malicious code (unless you fully trust the model owners).
+>
+> ```py
+> transformer = AutoModel.from_pretrained(
+> "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, revision="a3d8cb2"
+> )
+> ```
+
> [!NOTE]
> Learn more about implementing custom models in the [Community components](../using-diffusers/custom_pipeline_overview#community-components) guide.
\ No newline at end of file
diff --git a/docs/source/zh/modular_diffusers/guiders.md b/docs/source/zh/modular_diffusers/guiders.md
index 50436f90c4a5..2315625a197a 100644
--- a/docs/source/zh/modular_diffusers/guiders.md
+++ b/docs/source/zh/modular_diffusers/guiders.md
@@ -86,10 +86,7 @@ t2i_pipeline.guider
## 更改引导器参数
-引导器参数可以通过 [`~ComponentSpec.create`] 方法或 [`~ModularPipeline.update_components`] 方法进行调整。下面的示例更改了 `guidance_scale` 值。
-
-
-
+引导器参数可以通过 [`~ComponentSpec.create`] 方法以及 [`~ModularPipeline.update_components`] 方法进行调整。下面的示例更改了 `guidance_scale` 值。
```py
guider_spec = t2i_pipeline.get_component_spec("guider")
@@ -97,18 +94,6 @@ guider = guider_spec.create(guidance_scale=10)
t2i_pipeline.update_components(guider=guider)
```
-
-
-
-```py
-guider_spec = t2i_pipeline.get_component_spec("guider")
-guider_spec.config["guidance_scale"] = 10
-t2i_pipeline.update_components(guider=guider_spec)
-```
-
-
-
-
## 上传自定义引导器
在自定义引导器上调用 [`~utils.PushToHubMixin.push_to_hub`] 方法，将其分享到 Hub。
diff --git a/examples/cogvideo/train_cogvideox_image_to_video_lora.py b/examples/cogvideo/train_cogvideox_image_to_video_lora.py
index 001934298abe..c59986d2fde7 100644
--- a/examples/cogvideo/train_cogvideox_image_to_video_lora.py
+++ b/examples/cogvideo/train_cogvideox_image_to_video_lora.py
@@ -432,9 +432,9 @@ def get_args():
class VideoDataset(Dataset):
def __init__(
self,
- instance_data_root: Optional[str] = None,
- dataset_name: Optional[str] = None,
- dataset_config_name: Optional[str] = None,
+ instance_data_root: str | None = None,
+ dataset_name: str | None = None,
+ dataset_config_name: str | None = None,
caption_column: str = "text",
video_column: str = "video",
height: int = 480,
@@ -443,8 +443,8 @@ def __init__(
max_num_frames: int = 49,
skip_frames_start: int = 0,
skip_frames_end: int = 0,
- cache_dir: Optional[str] = None,
- id_token: Optional[str] = None,
+ cache_dir: str | None = None,
+ id_token: str | None = None,
) -> None:
super().__init__()
diff --git a/examples/cogvideo/train_cogvideox_lora.py b/examples/cogvideo/train_cogvideox_lora.py
index f6f2dc83a3f9..4e22d3f8727d 100644
--- a/examples/cogvideo/train_cogvideox_lora.py
+++ b/examples/cogvideo/train_cogvideox_lora.py
@@ -416,9 +416,9 @@ def get_args():
class VideoDataset(Dataset):
def __init__(
self,
- instance_data_root: Optional[str] = None,
- dataset_name: Optional[str] = None,
- dataset_config_name: Optional[str] = None,
+ instance_data_root: str | None = None,
+ dataset_name: str | None = None,
+ dataset_config_name: str | None = None,
caption_column: str = "text",
video_column: str = "video",
height: int = 480,
@@ -428,8 +428,8 @@ def __init__(
max_num_frames: int = 49,
skip_frames_start: int = 0,
skip_frames_end: int = 0,
- cache_dir: Optional[str] = None,
- id_token: Optional[str] = None,
+ cache_dir: str | None = None,
+ id_token: str | None = None,
) -> None:
super().__init__()
diff --git a/examples/community/README_community_scripts.md b/examples/community/README_community_scripts.md
index 3c9ad0d89bb4..d790b67a44c5 100644
--- a/examples/community/README_community_scripts.md
+++ b/examples/community/README_community_scripts.md
@@ -260,7 +260,7 @@ class SDPromptSchedulingCallback(PipelineCallback):
def callback_fn(
self, pipeline, step_index, timestep, callback_kwargs
- ) -> Dict[str, Any]:
+ ) -> dict[str, Any]:
cutoff_step_ratio = self.config.cutoff_step_ratio
cutoff_step_index = self.config.cutoff_step_index
if isinstance(self.config.encoded_prompt, tuple):
@@ -343,7 +343,7 @@ class SDXLPromptSchedulingCallback(PipelineCallback):
def callback_fn(
self, pipeline, step_index, timestep, callback_kwargs
- ) -> Dict[str, Any]:
+ ) -> dict[str, Any]:
cutoff_step_ratio = self.config.cutoff_step_ratio
cutoff_step_index = self.config.cutoff_step_index
if isinstance(self.config.encoded_prompt, tuple):
diff --git a/examples/community/adaptive_mask_inpainting.py b/examples/community/adaptive_mask_inpainting.py
index da67debe72ec..0378ffcdf985 100644
--- a/examples/community/adaptive_mask_inpainting.py
+++ b/examples/community/adaptive_mask_inpainting.py
@@ -871,7 +871,7 @@ def __call__(
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/bit_diffusion.py b/examples/community/bit_diffusion.py
index 67f4cd3fe199..c19bc95eefe0 100644
--- a/examples/community/bit_diffusion.py
+++ b/examples/community/bit_diffusion.py
@@ -231,9 +231,9 @@ def __call__(
height: Optional[int] = 256,
width: Optional[int] = 256,
num_inference_steps: Optional[int] = 50,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
batch_size: Optional[int] = 1,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
**kwargs,
) -> Union[Tuple, ImagePipelineOutput]:
diff --git a/examples/community/clip_guided_images_mixing_stable_diffusion.py b/examples/community/clip_guided_images_mixing_stable_diffusion.py
index 2cd3daf68c24..bacf9f0ec8b4 100644
--- a/examples/community/clip_guided_images_mixing_stable_diffusion.py
+++ b/examples/community/clip_guided_images_mixing_stable_diffusion.py
@@ -235,8 +235,8 @@ def __call__(
self,
style_image: Union[torch.Tensor, PIL.Image.Image],
content_image: Union[torch.Tensor, PIL.Image.Image],
- style_prompt: Optional[str] = None,
- content_prompt: Optional[str] = None,
+ style_prompt: str | None = None,
+ content_prompt: str | None = None,
height: Optional[int] = 512,
width: Optional[int] = 512,
noise_strength: float = 0.6,
@@ -245,8 +245,8 @@ def __call__(
batch_size: Optional[int] = 1,
eta: float = 0.0,
clip_guidance_scale: Optional[float] = 100,
- generator: Optional[torch.Generator] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
slerp_latent_style_strength: float = 0.8,
slerp_prompt_style_strength: float = 0.1,
diff --git a/examples/community/clip_guided_stable_diffusion.py b/examples/community/clip_guided_stable_diffusion.py
index bfd0858d245e..013df7acc7c7 100644
--- a/examples/community/clip_guided_stable_diffusion.py
+++ b/examples/community/clip_guided_stable_diffusion.py
@@ -179,9 +179,9 @@ def __call__(
clip_prompt: Optional[Union[str, List[str]]] = None,
num_cutouts: Optional[int] = 4,
use_cutouts: Optional[bool] = True,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
if isinstance(prompt, str):
diff --git a/examples/community/clip_guided_stable_diffusion_img2img.py b/examples/community/clip_guided_stable_diffusion_img2img.py
index f3dd4903f851..739d6dafff73 100644
--- a/examples/community/clip_guided_stable_diffusion_img2img.py
+++ b/examples/community/clip_guided_stable_diffusion_img2img.py
@@ -316,9 +316,9 @@ def __call__(
clip_prompt: Optional[Union[str, List[str]]] = None,
num_cutouts: Optional[int] = 4,
use_cutouts: Optional[bool] = True,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
if isinstance(prompt, str):
diff --git a/examples/community/composable_stable_diffusion.py b/examples/community/composable_stable_diffusion.py
index a7c610ad4355..466502e8394c 100644
--- a/examples/community/composable_stable_diffusion.py
+++ b/examples/community/composable_stable_diffusion.py
@@ -357,13 +357,13 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
- weights: Optional[str] = "",
+ weights: str | None = "",
):
r"""
Function invoked when calling the pipeline for generation.
diff --git a/examples/community/ddim_noise_comparative_analysis.py b/examples/community/ddim_noise_comparative_analysis.py
index 829106c47f65..453f8525a552 100644
--- a/examples/community/ddim_noise_comparative_analysis.py
+++ b/examples/community/ddim_noise_comparative_analysis.py
@@ -110,7 +110,7 @@ def __call__(
eta: float = 0.0,
num_inference_steps: int = 50,
use_clipped_model_output: Optional[bool] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
) -> Union[ImagePipelineOutput, Tuple]:
r"""
diff --git a/examples/community/dps_pipeline.py b/examples/community/dps_pipeline.py
index b29b06365bf3..6942d8ae608f 100755
--- a/examples/community/dps_pipeline.py
+++ b/examples/community/dps_pipeline.py
@@ -54,7 +54,7 @@ def __call__(
batch_size: int = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
num_inference_steps: int = 1000,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
zeta: float = 0.3,
) -> Union[ImagePipelineOutput, Tuple]:
diff --git a/examples/community/edict_pipeline.py b/examples/community/edict_pipeline.py
index a7bc892ddf93..bd96cdd6b2fc 100644
--- a/examples/community/edict_pipeline.py
+++ b/examples/community/edict_pipeline.py
@@ -1,5 +1,3 @@
-from typing import Optional
-
import torch
from PIL import Image
from tqdm.auto import tqdm
@@ -39,7 +37,7 @@ def __init__(
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
def _encode_prompt(
- self, prompt: str, negative_prompt: Optional[str] = None, do_classifier_free_guidance: bool = False
+ self, prompt: str, negative_prompt: str | None = None, do_classifier_free_guidance: bool = False
):
text_inputs = self.tokenizer(
prompt,
@@ -141,7 +139,7 @@ def prepare_latents(
text_embeds: torch.Tensor,
timesteps: torch.Tensor,
guidance_scale: float,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
):
do_classifier_free_guidance = guidance_scale > 1.0
@@ -194,9 +192,9 @@ def __call__(
guidance_scale: float = 3.0,
num_inference_steps: int = 50,
strength: float = 0.8,
- negative_prompt: Optional[str] = None,
- generator: Optional[torch.Generator] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | None = None,
+ generator: torch.Generator | None = None,
+ output_type: str | None = "pil",
):
do_classifier_free_guidance = guidance_scale > 1.0
diff --git a/examples/community/fresco_v2v.py b/examples/community/fresco_v2v.py
index b79834db5eed..974f71805baa 100644
--- a/examples/community/fresco_v2v.py
+++ b/examples/community/fresco_v2v.py
@@ -1208,7 +1208,7 @@ def apply_FRESCO_attn(pipe):
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -2064,7 +2064,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
diff --git a/examples/community/gluegen.py b/examples/community/gluegen.py
index 86813b63eca5..b98e0465561d 100644
--- a/examples/community/gluegen.py
+++ b/examples/community/gluegen.py
@@ -597,7 +597,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/hd_painter.py b/examples/community/hd_painter.py
index 70e5656855ff..f412437b30ff 100644
--- a/examples/community/hd_painter.py
+++ b/examples/community/hd_painter.py
@@ -462,7 +462,7 @@ def __call__(
num_inference_steps: int = 50,
timesteps: List[int] = None,
guidance_scale: float = 7.5,
- positive_prompt: Optional[str] = "",
+ positive_prompt: str | None = "",
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.01,
@@ -471,7 +471,7 @@ def __call__(
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
clip_skip: int = None,
diff --git a/examples/community/iadb.py b/examples/community/iadb.py
index 6262c3cb15fc..d61e2573b551 100644
--- a/examples/community/iadb.py
+++ b/examples/community/iadb.py
@@ -86,7 +86,7 @@ def __call__(
batch_size: int = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
num_inference_steps: int = 50,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
) -> Union[ImagePipelineOutput, Tuple]:
r"""
diff --git a/examples/community/imagic_stable_diffusion.py b/examples/community/imagic_stable_diffusion.py
index 091d0fbf8d3a..2efea594aff2 100644
--- a/examples/community/imagic_stable_diffusion.py
+++ b/examples/community/imagic_stable_diffusion.py
@@ -113,7 +113,7 @@ def train(
image: Union[torch.Tensor, PIL.Image.Image],
height: Optional[int] = 512,
width: Optional[int] = 512,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
embedding_learning_rate: float = 0.001,
diffusion_model_learning_rate: float = 2e-6,
text_embedding_optimization_steps: int = 500,
@@ -314,8 +314,8 @@ def __call__(
height: Optional[int] = 512,
width: Optional[int] = 512,
num_inference_steps: Optional[int] = 50,
- generator: Optional[torch.Generator] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
guidance_scale: float = 7.5,
eta: float = 0.0,
diff --git a/examples/community/img2img_inpainting.py b/examples/community/img2img_inpainting.py
index bef682425a2c..efc24a339d4f 100644
--- a/examples/community/img2img_inpainting.py
+++ b/examples/community/img2img_inpainting.py
@@ -143,9 +143,9 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/instaflow_one_step.py b/examples/community/instaflow_one_step.py
index 0f16707eadf5..d5c304b48d81 100644
--- a/examples/community/instaflow_one_step.py
+++ b/examples/community/instaflow_one_step.py
@@ -512,7 +512,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/interpolate_stable_diffusion.py b/examples/community/interpolate_stable_diffusion.py
index 5b96c14d6367..448cf94216a8 100644
--- a/examples/community/interpolate_stable_diffusion.py
+++ b/examples/community/interpolate_stable_diffusion.py
@@ -131,9 +131,9 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
@@ -401,8 +401,8 @@ def walk(
prompts: List[str],
seeds: List[int],
num_interpolation_steps: Optional[int] = 6,
- output_dir: Optional[str] = "./dreams",
- name: Optional[str] = None,
+ output_dir: str | None = "./dreams",
+ name: str | None = None,
batch_size: Optional[int] = 1,
height: Optional[int] = 512,
width: Optional[int] = 512,
diff --git a/examples/community/ip_adapter_face_id.py b/examples/community/ip_adapter_face_id.py
index d16aaf5a54c6..817d31e51af5 100644
--- a/examples/community/ip_adapter_face_id.py
+++ b/examples/community/ip_adapter_face_id.py
@@ -855,7 +855,7 @@ def __call__(
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/latent_consistency_img2img.py b/examples/community/latent_consistency_img2img.py
index 0f5711f34b62..a3381dd9bd6f 100644
--- a/examples/community/latent_consistency_img2img.py
+++ b/examples/community/latent_consistency_img2img.py
@@ -286,7 +286,7 @@ def __call__(
num_inference_steps: int = 4,
lcm_origin_steps: int = 50,
prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
):
diff --git a/examples/community/latent_consistency_interpolate.py b/examples/community/latent_consistency_interpolate.py
index e8349ba317e9..ce87f61b13b0 100644
--- a/examples/community/latent_consistency_interpolate.py
+++ b/examples/community/latent_consistency_interpolate.py
@@ -669,7 +669,7 @@ def __call__(
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
clip_skip: Optional[int] = None,
diff --git a/examples/community/latent_consistency_txt2img.py b/examples/community/latent_consistency_txt2img.py
index 0ce982065619..80dc7e2fc2c3 100755
--- a/examples/community/latent_consistency_txt2img.py
+++ b/examples/community/latent_consistency_txt2img.py
@@ -212,7 +212,7 @@ def __call__(
num_inference_steps: int = 4,
lcm_origin_steps: int = 50,
prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
):
diff --git a/examples/community/llm_grounded_diffusion.py b/examples/community/llm_grounded_diffusion.py
index 5bf6674a43e3..1ec565bac00e 100644
--- a/examples/community/llm_grounded_diffusion.py
+++ b/examples/community/llm_grounded_diffusion.py
@@ -769,7 +769,7 @@ def __call__(
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/lpw_stable_diffusion.py b/examples/community/lpw_stable_diffusion.py
index 58e932bbcf74..3407c01e1edd 100644
--- a/examples/community/lpw_stable_diffusion.py
+++ b/examples/community/lpw_stable_diffusion.py
@@ -830,7 +830,7 @@ def __call__(
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
max_embeddings_multiples: Optional[int] = 3,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
is_cancelled_callback: Optional[Callable[[], bool]] = None,
@@ -1091,7 +1091,7 @@ def text2img(
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
max_embeddings_multiples: Optional[int] = 3,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
is_cancelled_callback: Optional[Callable[[], bool]] = None,
@@ -1209,7 +1209,7 @@ def img2img(
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
max_embeddings_multiples: Optional[int] = 3,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
is_cancelled_callback: Optional[Callable[[], bool]] = None,
@@ -1323,7 +1323,7 @@ def inpaint(
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
max_embeddings_multiples: Optional[int] = 3,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
is_cancelled_callback: Optional[Callable[[], bool]] = None,
diff --git a/examples/community/lpw_stable_diffusion_onnx.py b/examples/community/lpw_stable_diffusion_onnx.py
index 92effc193329..eb7d86d09d3f 100644
--- a/examples/community/lpw_stable_diffusion_onnx.py
+++ b/examples/community/lpw_stable_diffusion_onnx.py
@@ -664,10 +664,10 @@ def __call__(
strength: float = 0.8,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[np.ndarray] = None,
max_embeddings_multiples: Optional[int] = 3,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
is_cancelled_callback: Optional[Callable[[], bool]] = None,
@@ -877,10 +877,10 @@ def text2img(
guidance_scale: float = 7.5,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[np.ndarray] = None,
max_embeddings_multiples: Optional[int] = 3,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
callback_steps: int = 1,
@@ -969,9 +969,9 @@ def img2img(
guidance_scale: Optional[float] = 7.5,
num_images_per_prompt: Optional[int] = 1,
eta: Optional[float] = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
max_embeddings_multiples: Optional[int] = 3,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
callback_steps: int = 1,
@@ -1061,9 +1061,9 @@ def inpaint(
guidance_scale: Optional[float] = 7.5,
num_images_per_prompt: Optional[int] = 1,
eta: Optional[float] = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
max_embeddings_multiples: Optional[int] = 3,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/lpw_stable_diffusion_xl.py b/examples/community/lpw_stable_diffusion_xl.py
index 1aba1f270198..b7a3623bf6ce 100644
--- a/examples/community/lpw_stable_diffusion_xl.py
+++ b/examples/community/lpw_stable_diffusion_xl.py
@@ -519,7 +519,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -724,12 +724,12 @@ def enable_model_cpu_offload(self, gpu_id=0):
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -1399,7 +1399,7 @@ def num_timesteps(self):
def __call__(
self,
prompt: str = None,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
image: Optional[PipelineImageInput] = None,
mask_image: Optional[PipelineImageInput] = None,
masked_image_latents: Optional[torch.Tensor] = None,
@@ -1411,8 +1411,8 @@ def __call__(
denoising_start: Optional[float] = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
@@ -1422,7 +1422,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
@@ -1955,7 +1955,7 @@ def denoising_value_valid(dnv):
def text2img(
self,
prompt: str = None,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
@@ -1963,8 +1963,8 @@ def text2img(
denoising_start: Optional[float] = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
@@ -1974,7 +1974,7 @@ def text2img(
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
@@ -2028,7 +2028,7 @@ def text2img(
def img2img(
self,
prompt: str = None,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
image: Optional[PipelineImageInput] = None,
height: Optional[int] = None,
width: Optional[int] = None,
@@ -2038,8 +2038,8 @@ def img2img(
denoising_start: Optional[float] = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
@@ -2049,7 +2049,7 @@ def img2img(
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
@@ -2105,7 +2105,7 @@ def img2img(
def inpaint(
self,
prompt: str = None,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
image: Optional[PipelineImageInput] = None,
mask_image: Optional[PipelineImageInput] = None,
masked_image_latents: Optional[torch.Tensor] = None,
@@ -2117,8 +2117,8 @@ def inpaint(
denoising_start: Optional[float] = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
@@ -2128,7 +2128,7 @@ def inpaint(
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/masked_stable_diffusion_img2img.py b/examples/community/masked_stable_diffusion_img2img.py
index 570bd0963e28..ac1612527d6a 100644
--- a/examples/community/masked_stable_diffusion_img2img.py
+++ b/examples/community/masked_stable_diffusion_img2img.py
@@ -32,7 +32,7 @@ def __call__(
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/masked_stable_diffusion_xl_img2img.py b/examples/community/masked_stable_diffusion_xl_img2img.py
index 14d8c7c2da78..9e47b79e18be 100644
--- a/examples/community/masked_stable_diffusion_xl_img2img.py
+++ b/examples/community/masked_stable_diffusion_xl_img2img.py
@@ -59,7 +59,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/matryoshka.py b/examples/community/matryoshka.py
index 09b1d1b24465..a4971f94829d 100644
--- a/examples/community/matryoshka.py
+++ b/examples/community/matryoshka.py
@@ -783,7 +783,7 @@ def __init__(
norm_type: str = "layer_norm",
num_attention_heads: int = 1,
cross_attention_dim: int = 1280,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
output_scale_factor: float = 1.0,
downsample_padding: int = 1,
add_downsample: bool = True,
@@ -922,7 +922,7 @@ def __init__(
num_attention_heads: int = 1,
output_scale_factor: float = 1.0,
cross_attention_dim: int = 1280,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
upcast_attention: bool = False,
@@ -1055,7 +1055,7 @@ def __init__(
norm_type: str = "layer_norm",
num_attention_heads: int = 1,
cross_attention_dim: int = 1280,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
output_scale_factor: float = 1.0,
add_upsample: bool = True,
dual_cross_attention: bool = False,
@@ -1617,10 +1617,10 @@ def get_down_block(
attention_pre_only: bool = False,
resnet_skip_time_act: bool = False,
resnet_out_scale_factor: float = 1.0,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
attention_head_dim: Optional[int] = None,
use_attention_ffn: bool = True,
- downsample_type: Optional[str] = None,
+ downsample_type: str | None = None,
dropout: float = 0.0,
):
# If attn head dim is not defined, we default it to the number of heads
@@ -1695,7 +1695,7 @@ def get_mid_block(
attention_type: str = "default",
attention_pre_only: bool = False,
resnet_skip_time_act: bool = False,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
attention_head_dim: Optional[int] = 1,
dropout: float = 0.0,
):
@@ -1747,10 +1747,10 @@ def get_up_block(
attention_pre_only: bool = False,
resnet_skip_time_act: bool = False,
resnet_out_scale_factor: float = 1.0,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
attention_head_dim: Optional[int] = None,
use_attention_ffn: bool = True,
- upsample_type: Optional[str] = None,
+ upsample_type: str | None = None,
dropout: float = 0.0,
) -> nn.Module:
# If attn head dim is not defined, we default it to the number of heads
@@ -1972,7 +1972,7 @@ def __init__(
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+ mid_block_type: str | None = "UNetMidBlock2DCrossAttn",
up_block_types: Tuple[str, ...] = (
"UpBlock2D",
"CrossAttnUpBlock2D",
@@ -1993,14 +1993,14 @@ def __init__(
transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
+ encoder_hid_dim_type: str | None = None,
attention_head_dim: Union[int, Tuple[int]] = 8,
num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
dual_cross_attention: bool = False,
use_attention_ffn: bool = True,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
+ class_embed_type: str | None = None,
+ addition_embed_type: str | None = None,
addition_time_embed_dim: Optional[int] = None,
num_class_embeds: Optional[int] = None,
upcast_attention: bool = False,
@@ -2009,8 +2009,8 @@ def __init__(
resnet_out_scale_factor: float = 1.0,
time_embedding_type: str = "positional",
time_embedding_dim: Optional[int] = None,
- time_embedding_act_fn: Optional[str] = None,
- timestep_post_act: Optional[str] = None,
+ time_embedding_act_fn: str | None = None,
+ timestep_post_act: str | None = None,
time_cond_proj_dim: Optional[int] = None,
conv_in_kernel: int = 3,
conv_out_kernel: int = 3,
@@ -2021,7 +2021,7 @@ def __init__(
micro_conditioning_scale: int = None,
class_embeddings_concat: bool = False,
mid_block_only_cross_attention: Optional[bool] = None,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
addition_embed_type_num_heads: int = 64,
temporal_mode: bool = False,
temporal_spatial_ds: bool = False,
@@ -2384,7 +2384,7 @@ def _set_time_proj(
def _set_encoder_hid_proj(
self,
- encoder_hid_dim_type: Optional[str],
+ encoder_hid_dim_type: str | None,
cross_attention_dim: Union[int, Tuple[int]],
encoder_hid_dim: Optional[int],
):
@@ -2424,7 +2424,7 @@ def _set_encoder_hid_proj(
def _set_class_embedding(
self,
- class_embed_type: Optional[str],
+ class_embed_type: str | None,
act_fn: str,
num_class_embeds: Optional[int],
projection_class_embeddings_input_dim: Optional[int],
@@ -2524,7 +2524,7 @@ def _set_pos_net_if_use_gligen(self, attention_type: str, cross_attention_dim: i
)
@property
- def attn_processors(self) -> Dict[str, AttentionProcessor]:
+ def attn_processors(self) -> dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
@@ -4264,7 +4264,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/mixture_tiling_sdxl.py b/examples/community/mixture_tiling_sdxl.py
index e09f5a25db73..dee1c3ef055e 100644
--- a/examples/community/mixture_tiling_sdxl.py
+++ b/examples/community/mixture_tiling_sdxl.py
@@ -388,12 +388,12 @@ class SeedTilesMode(Enum):
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -780,7 +780,7 @@ def __call__(
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
original_size: Optional[Tuple[int, int]] = None,
diff --git a/examples/community/mod_controlnet_tile_sr_sdxl.py b/examples/community/mod_controlnet_tile_sr_sdxl.py
index fe8bd73eabbe..959ef6453de0 100644
--- a/examples/community/mod_controlnet_tile_sr_sdxl.py
+++ b/examples/community/mod_controlnet_tile_sr_sdxl.py
@@ -243,7 +243,7 @@ def _tile2latent_indices(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -395,12 +395,12 @@ class TileWeightingMethod(Enum):
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -1265,7 +1265,7 @@ def __call__(
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
diff --git a/examples/community/multilingual_stable_diffusion.py b/examples/community/multilingual_stable_diffusion.py
index 436803f201fe..444012acacec 100644
--- a/examples/community/multilingual_stable_diffusion.py
+++ b/examples/community/multilingual_stable_diffusion.py
@@ -146,9 +146,9 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/pipeline_animatediff_controlnet.py b/examples/community/pipeline_animatediff_controlnet.py
index 903bfd4fd57b..b195ba0f8a20 100644
--- a/examples/community/pipeline_animatediff_controlnet.py
+++ b/examples/community/pipeline_animatediff_controlnet.py
@@ -762,7 +762,7 @@ def __call__(
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[PipelineImageInput] = None,
conditioning_frames: Optional[List[PipelineImageInput]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
diff --git a/examples/community/pipeline_animatediff_img2video.py b/examples/community/pipeline_animatediff_img2video.py
index feba19f70fc6..5deadf6784c0 100644
--- a/examples/community/pipeline_animatediff_img2video.py
+++ b/examples/community/pipeline_animatediff_img2video.py
@@ -182,7 +182,7 @@ def tensor2vid(video: torch.Tensor, processor, output_type="np"):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -755,7 +755,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: Optional[int] = 1,
diff --git a/examples/community/pipeline_animatediff_ipex.py b/examples/community/pipeline_animatediff_ipex.py
index 409ab9d6ad73..43d93bd3af2a 100644
--- a/examples/community/pipeline_animatediff_ipex.py
+++ b/examples/community/pipeline_animatediff_ipex.py
@@ -588,7 +588,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
clip_skip: Optional[int] = None,
@@ -844,7 +844,7 @@ def prepare_for_ipex(
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
clip_skip: Optional[int] = None,
diff --git a/examples/community/pipeline_controlnet_xl_kolors.py b/examples/community/pipeline_controlnet_xl_kolors.py
index 96c801bab379..d35176ba4f45 100644
--- a/examples/community/pipeline_controlnet_xl_kolors.py
+++ b/examples/community/pipeline_controlnet_xl_kolors.py
@@ -111,7 +111,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -800,7 +800,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
diff --git a/examples/community/pipeline_controlnet_xl_kolors_img2img.py b/examples/community/pipeline_controlnet_xl_kolors_img2img.py
index 41f2550138b3..077241d9da22 100644
--- a/examples/community/pipeline_controlnet_xl_kolors_img2img.py
+++ b/examples/community/pipeline_controlnet_xl_kolors_img2img.py
@@ -131,7 +131,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -972,7 +972,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
diff --git a/examples/community/pipeline_controlnet_xl_kolors_inpaint.py b/examples/community/pipeline_controlnet_xl_kolors_inpaint.py
index 07bd552a29e7..387a42ba22b3 100644
--- a/examples/community/pipeline_controlnet_xl_kolors_inpaint.py
+++ b/examples/community/pipeline_controlnet_xl_kolors_inpaint.py
@@ -118,7 +118,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -1139,7 +1139,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
diff --git a/examples/community/pipeline_demofusion_sdxl.py b/examples/community/pipeline_demofusion_sdxl.py
index b5b68577982b..b0a3cab347a0 100644
--- a/examples/community/pipeline_demofusion_sdxl.py
+++ b/examples/community/pipeline_demofusion_sdxl.py
@@ -184,12 +184,12 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -637,7 +637,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = False,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/pipeline_fabric.py b/examples/community/pipeline_fabric.py
index d29e98df5e50..fe7de1dd94a7 100644
--- a/examples/community/pipeline_fabric.py
+++ b/examples/community/pipeline_fabric.py
@@ -508,7 +508,7 @@ def __call__(
num_images: int = 4,
guidance_scale: float = 7.0,
num_inference_steps: int = 20,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
feedback_start_ratio: float = 0.33,
feedback_end_ratio: float = 0.66,
min_weight: float = 0.05,
diff --git a/examples/community/pipeline_faithdiff_stable_diffusion_xl.py b/examples/community/pipeline_faithdiff_stable_diffusion_xl.py
index 4b564d5ee5c8..796a900b2f75 100644
--- a/examples/community/pipeline_faithdiff_stable_diffusion_xl.py
+++ b/examples/community/pipeline_faithdiff_stable_diffusion_xl.py
@@ -437,7 +437,7 @@ def __init__(
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+ mid_block_type: str | None = "UNetMidBlock2DCrossAttn",
up_block_types: Tuple[str, ...] = (
"UpBlock2D",
"CrossAttnUpBlock2D",
@@ -457,13 +457,13 @@ def __init__(
transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
+ encoder_hid_dim_type: str | None = None,
attention_head_dim: Union[int, Tuple[int]] = 8,
num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
+ class_embed_type: str | None = None,
+ addition_embed_type: str | None = None,
addition_time_embed_dim: Optional[int] = None,
num_class_embeds: Optional[int] = None,
upcast_attention: bool = False,
@@ -472,8 +472,8 @@ def __init__(
resnet_out_scale_factor: float = 1.0,
time_embedding_type: str = "positional",
time_embedding_dim: Optional[int] = None,
- time_embedding_act_fn: Optional[str] = None,
- timestep_post_act: Optional[str] = None,
+ time_embedding_act_fn: str | None = None,
+ timestep_post_act: str | None = None,
time_cond_proj_dim: Optional[int] = None,
conv_in_kernel: int = 3,
conv_out_kernel: int = 3,
@@ -481,7 +481,7 @@ def __init__(
attention_type: str = "default",
class_embeddings_concat: bool = False,
mid_block_only_cross_attention: Optional[bool] = None,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
addition_embed_type_num_heads: int = 64,
):
"""Initialize the UnifiedUNet2DConditionModel."""
@@ -565,7 +565,7 @@ def init_extra_weights(self):
self.agg_net = nn.ModuleList()
def load_additional_layers(
- self, dtype: Optional[torch.dtype] = torch.float16, channel: int = 512, weight_path: Optional[str] = None
+ self, dtype: Optional[torch.dtype] = torch.float16, channel: int = 512, weight_path: str | None = None
):
"""Load additional layers and weights from a file.
@@ -1096,7 +1096,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
"""Retrieve latents from an encoder output.
@@ -1267,12 +1267,12 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
@@ -1808,7 +1808,7 @@ def __call__(
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
- start_point: Optional[str] = "noise",
+ start_point: str | None = "noise",
timesteps: List[int] = None,
denoising_end: Optional[float] = None,
overlap: float = 0.5,
@@ -1823,7 +1823,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/pipeline_flux_differential_img2img.py b/examples/community/pipeline_flux_differential_img2img.py
index 3677e73136f7..7cfb816d7310 100644
--- a/examples/community/pipeline_flux_differential_img2img.py
+++ b/examples/community/pipeline_flux_differential_img2img.py
@@ -97,7 +97,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -698,7 +698,7 @@ def __call__(
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_flux_kontext_multiple_images.py b/examples/community/pipeline_flux_kontext_multiple_images.py
index 9e6ae427dbfa..fb344859e99e 100644
--- a/examples/community/pipeline_flux_kontext_multiple_images.py
+++ b/examples/community/pipeline_flux_kontext_multiple_images.py
@@ -188,7 +188,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -838,7 +838,7 @@ def __call__(
negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_flux_rf_inversion.py b/examples/community/pipeline_flux_rf_inversion.py
index 2cd6eb088cd8..16bb83bbcf34 100644
--- a/examples/community/pipeline_flux_rf_inversion.py
+++ b/examples/community/pipeline_flux_rf_inversion.py
@@ -685,7 +685,7 @@ def __call__(
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_flux_semantic_guidance.py b/examples/community/pipeline_flux_semantic_guidance.py
index 74cd5c6981b0..a8d64b9caca6 100644
--- a/examples/community/pipeline_flux_semantic_guidance.py
+++ b/examples/community/pipeline_flux_semantic_guidance.py
@@ -802,7 +802,7 @@ def __call__(
negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_flux_with_cfg.py b/examples/community/pipeline_flux_with_cfg.py
index 5bc13f7e5e11..107901a76f57 100644
--- a/examples/community/pipeline_flux_with_cfg.py
+++ b/examples/community/pipeline_flux_with_cfg.py
@@ -622,7 +622,7 @@ def __call__(
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_hunyuandit_differential_img2img.py b/examples/community/pipeline_hunyuandit_differential_img2img.py
index bc6841525b49..123e46e53bae 100644
--- a/examples/community/pipeline_hunyuandit_differential_img2img.py
+++ b/examples/community/pipeline_hunyuandit_differential_img2img.py
@@ -164,7 +164,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
encoder_output: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
sample_mode: str = "sample",
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
@@ -349,7 +349,7 @@ def encode_prompt(
dtype: torch.dtype = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
+ negative_prompt: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
prompt_attention_mask: Optional[torch.Tensor] = None,
@@ -749,7 +749,7 @@ def __call__(
prompt_attention_mask_2: Optional[torch.Tensor] = None,
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
negative_prompt_attention_mask_2: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback_on_step_end: Optional[
Union[
diff --git a/examples/community/pipeline_kolors_differential_img2img.py b/examples/community/pipeline_kolors_differential_img2img.py
index d456af8b3385..507f71ad61ef 100644
--- a/examples/community/pipeline_kolors_differential_img2img.py
+++ b/examples/community/pipeline_kolors_differential_img2img.py
@@ -67,7 +67,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -800,7 +800,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
original_size: Optional[Tuple[int, int]] = None,
diff --git a/examples/community/pipeline_kolors_inpainting.py b/examples/community/pipeline_kolors_inpainting.py
index 85a3d698efa1..d5cb57580f58 100644
--- a/examples/community/pipeline_kolors_inpainting.py
+++ b/examples/community/pipeline_kolors_inpainting.py
@@ -239,7 +239,7 @@ def prepare_mask_and_masked_image(image, mask, height, width, return_image: bool
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -1100,7 +1100,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/pipeline_prompt2prompt.py b/examples/community/pipeline_prompt2prompt.py
index eb19667970b0..9ca7058892c4 100644
--- a/examples/community/pipeline_prompt2prompt.py
+++ b/examples/community/pipeline_prompt2prompt.py
@@ -571,7 +571,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: Optional[int] = 1,
diff --git a/examples/community/pipeline_sdxl_style_aligned.py b/examples/community/pipeline_sdxl_style_aligned.py
index ac7ad33b68c5..75f3cb9ead76 100644
--- a/examples/community/pipeline_sdxl_style_aligned.py
+++ b/examples/community/pipeline_sdxl_style_aligned.py
@@ -371,7 +371,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -508,12 +508,12 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -1394,7 +1394,7 @@ def __call__(
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/pipeline_stable_diffusion_3_differential_img2img.py b/examples/community/pipeline_stable_diffusion_3_differential_img2img.py
index 1803cf60cc4b..df5628e55eb6 100644
--- a/examples/community/pipeline_stable_diffusion_3_differential_img2img.py
+++ b/examples/community/pipeline_stable_diffusion_3_differential_img2img.py
@@ -68,7 +68,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -688,7 +688,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
clip_skip: Optional[int] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_stable_diffusion_3_instruct_pix2pix.py b/examples/community/pipeline_stable_diffusion_3_instruct_pix2pix.py
index d9cee800e8ad..d4cb2924b9dc 100644
--- a/examples/community/pipeline_stable_diffusion_3_instruct_pix2pix.py
+++ b/examples/community/pipeline_stable_diffusion_3_instruct_pix2pix.py
@@ -102,7 +102,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -880,7 +880,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
clip_skip: Optional[int] = None,
diff --git a/examples/community/pipeline_stable_diffusion_boxdiff.py b/examples/community/pipeline_stable_diffusion_boxdiff.py
index c05a6143132c..69cff867ed6e 100644
--- a/examples/community/pipeline_stable_diffusion_boxdiff.py
+++ b/examples/community/pipeline_stable_diffusion_boxdiff.py
@@ -1341,7 +1341,7 @@ def __call__(
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/pipeline_stable_diffusion_pag.py b/examples/community/pipeline_stable_diffusion_pag.py
index 3f98dca0b986..f8061e11aa99 100644
--- a/examples/community/pipeline_stable_diffusion_pag.py
+++ b/examples/community/pipeline_stable_diffusion_pag.py
@@ -1114,7 +1114,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py b/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py
index ba89556a5e4d..d537ef879711 100644
--- a/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py
+++ b/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py
@@ -523,7 +523,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py b/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py
index 822c804ec8a9..a001e458392c 100644
--- a/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py
+++ b/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py
@@ -435,7 +435,7 @@ def prepare_mask_and_masked_image(image, mask, height, width, return_image: bool
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -694,12 +694,12 @@ def prepare_ip_adapter_image_embeds(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
@@ -1640,7 +1640,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py
index d435a6e146b2..e38801cd7647 100644
--- a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py
+++ b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py
@@ -241,12 +241,12 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -878,7 +878,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py
index 8d70e8f3c2a1..2e05e3380316 100644
--- a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py
+++ b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py
@@ -388,12 +388,12 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -1219,7 +1219,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/pipeline_stable_diffusion_xl_differential_img2img.py b/examples/community/pipeline_stable_diffusion_xl_differential_img2img.py
index 99657facddf6..16144671d892 100644
--- a/examples/community/pipeline_stable_diffusion_xl_differential_img2img.py
+++ b/examples/community/pipeline_stable_diffusion_xl_differential_img2img.py
@@ -101,7 +101,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -268,12 +268,12 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -949,7 +949,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/pipeline_stable_diffusion_xl_instandid_img2img.py b/examples/community/pipeline_stable_diffusion_xl_instandid_img2img.py
index 33bf5ad346f4..1710f682d0ed 100644
--- a/examples/community/pipeline_stable_diffusion_xl_instandid_img2img.py
+++ b/examples/community/pipeline_stable_diffusion_xl_instandid_img2img.py
@@ -567,7 +567,7 @@ def __call__(
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
diff --git a/examples/community/pipeline_stable_diffusion_xl_instantid.py b/examples/community/pipeline_stable_diffusion_xl_instantid.py
index 45876b91f7d8..4dfbcc194dd8 100644
--- a/examples/community/pipeline_stable_diffusion_xl_instantid.py
+++ b/examples/community/pipeline_stable_diffusion_xl_instantid.py
@@ -565,7 +565,7 @@ def __call__(
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
diff --git a/examples/community/pipeline_stable_diffusion_xl_ipex.py b/examples/community/pipeline_stable_diffusion_xl_ipex.py
index 8e8357db3c7a..a0c07ccaeefb 100644
--- a/examples/community/pipeline_stable_diffusion_xl_ipex.py
+++ b/examples/community/pipeline_stable_diffusion_xl_ipex.py
@@ -268,12 +268,12 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
+ prompt_2: str | None = None,
device: Optional[torch.device] = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
@@ -741,7 +741,7 @@ def __call__(
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
@@ -1181,7 +1181,7 @@ def prepare_for_ipex(
pooled_prompt_embeds: Optional[torch.Tensor] = None,
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/pipeline_stg_hunyuan_video.py b/examples/community/pipeline_stg_hunyuan_video.py
index 028d54d047e4..489125ff6a62 100644
--- a/examples/community/pipeline_stg_hunyuan_video.py
+++ b/examples/community/pipeline_stg_hunyuan_video.py
@@ -571,7 +571,7 @@ def __call__(
prompt_embeds: Optional[torch.Tensor] = None,
pooled_prompt_embeds: Optional[torch.Tensor] = None,
prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[
diff --git a/examples/community/pipeline_stg_ltx.py b/examples/community/pipeline_stg_ltx.py
index 70069a33f5d9..5bf98f97fa7d 100644
--- a/examples/community/pipeline_stg_ltx.py
+++ b/examples/community/pipeline_stg_ltx.py
@@ -506,7 +506,7 @@ def prepare_latents(
num_frames: int = 161,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if latents is not None:
@@ -576,7 +576,7 @@ def __call__(
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
decode_timestep: Union[float, List[float]] = 0.0,
decode_noise_scale: Optional[Union[float, List[float]]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_stg_ltx_image2video.py b/examples/community/pipeline_stg_ltx_image2video.py
index c32805e1419f..4eec632d7853 100644
--- a/examples/community/pipeline_stg_ltx_image2video.py
+++ b/examples/community/pipeline_stg_ltx_image2video.py
@@ -200,7 +200,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -535,7 +535,7 @@ def prepare_latents(
num_frames: int = 161,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
) -> torch.Tensor:
height = height // self.vae_spatial_compression_ratio
@@ -636,7 +636,7 @@ def __call__(
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
decode_timestep: Union[float, List[float]] = 0.0,
decode_noise_scale: Optional[Union[float, List[float]]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_stg_mochi.py b/examples/community/pipeline_stg_mochi.py
index ad9317f6bc9d..350fba3bcf80 100644
--- a/examples/community/pipeline_stg_mochi.py
+++ b/examples/community/pipeline_stg_mochi.py
@@ -579,7 +579,7 @@ def __call__(
prompt_attention_mask: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_stg_wan.py b/examples/community/pipeline_stg_wan.py
index 39f208bad7c5..0833869bd35c 100644
--- a/examples/community/pipeline_stg_wan.py
+++ b/examples/community/pipeline_stg_wan.py
@@ -423,7 +423,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ output_type: str | None = "np",
return_dict: bool = True,
attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[
diff --git a/examples/community/pipeline_z_image_differential_img2img.py b/examples/community/pipeline_z_image_differential_img2img.py
index 8bde065c4013..6309e91abbdd 100644
--- a/examples/community/pipeline_z_image_differential_img2img.py
+++ b/examples/community/pipeline_z_image_differential_img2img.py
@@ -81,7 +81,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -467,7 +467,7 @@ def __call__(
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[List[torch.FloatTensor]] = None,
negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/pipeline_zero1to3.py b/examples/community/pipeline_zero1to3.py
index 1be59fd8323e..125566d7cf2d 100644
--- a/examples/community/pipeline_zero1to3.py
+++ b/examples/community/pipeline_zero1to3.py
@@ -614,7 +614,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/pipline_flux_fill_controlnet_Inpaint.py b/examples/community/pipline_flux_fill_controlnet_Inpaint.py
index cc642a767f87..88dcbecf8de0 100644
--- a/examples/community/pipline_flux_fill_controlnet_Inpaint.py
+++ b/examples/community/pipline_flux_fill_controlnet_Inpaint.py
@@ -99,7 +99,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -112,7 +112,7 @@ def retrieve_latents(
def retrieve_latents_fill(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -855,7 +855,7 @@ def __call__(
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/examples/community/regional_prompting_stable_diffusion.py b/examples/community/regional_prompting_stable_diffusion.py
index 3bc780cfcf7a..daea47b3c3dd 100644
--- a/examples/community/regional_prompting_stable_diffusion.py
+++ b/examples/community/regional_prompting_stable_diffusion.py
@@ -148,9 +148,9 @@ def __call__(
negative_prompt: str = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
rp_args: Dict[str, str] = None,
):
@@ -800,7 +800,7 @@ def stable_diffusion_call(
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/rerender_a_video.py b/examples/community/rerender_a_video.py
index 840f9e206d4d..b2680ebb2d4f 100644
--- a/examples/community/rerender_a_video.py
+++ b/examples/community/rerender_a_video.py
@@ -607,7 +607,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/run_onnx_controlnet.py b/examples/community/run_onnx_controlnet.py
index 2b56e8a1e51b..2c631ae37ed7 100644
--- a/examples/community/run_onnx_controlnet.py
+++ b/examples/community/run_onnx_controlnet.py
@@ -136,7 +136,7 @@ def _encode_prompt(
prompt: Union[str, List[str]],
num_images_per_prompt: Optional[int],
do_classifier_free_guidance: bool,
- negative_prompt: Optional[str],
+ negative_prompt: str | None,
prompt_embeds: Optional[np.ndarray] = None,
negative_prompt_embeds: Optional[np.ndarray] = None,
):
@@ -534,7 +534,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/run_tensorrt_controlnet.py b/examples/community/run_tensorrt_controlnet.py
index b62eb4f58e8f..23c13e6a501c 100644
--- a/examples/community/run_tensorrt_controlnet.py
+++ b/examples/community/run_tensorrt_controlnet.py
@@ -240,7 +240,7 @@ def _encode_prompt(
prompt: Union[str, List[str]],
num_images_per_prompt: Optional[int],
do_classifier_free_guidance: bool,
- negative_prompt: Optional[str],
+ negative_prompt: str | None,
prompt_embeds: Optional[np.ndarray] = None,
negative_prompt_embeds: Optional[np.ndarray] = None,
):
@@ -638,7 +638,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/scheduling_ufogen.py b/examples/community/scheduling_ufogen.py
index fada2636e98d..2d109ef36909 100644
--- a/examples/community/scheduling_ufogen.py
+++ b/examples/community/scheduling_ufogen.py
@@ -377,7 +377,7 @@ def step(
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
) -> Union[UFOGenSchedulerOutput, Tuple]:
"""
diff --git a/examples/community/sd_text2img_k_diffusion.py b/examples/community/sd_text2img_k_diffusion.py
index e351420f786a..690182714cde 100755
--- a/examples/community/sd_text2img_k_diffusion.py
+++ b/examples/community/sd_text2img_k_diffusion.py
@@ -285,9 +285,9 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/sde_drag.py b/examples/community/sde_drag.py
index 63899ce738bb..0b654730fcc4 100644
--- a/examples/community/sde_drag.py
+++ b/examples/community/sde_drag.py
@@ -73,7 +73,7 @@ def __call__(
image_scale: Optional[float] = 0.3,
adapt_radius: Optional[int] = 5,
min_lora_scale: Optional[float] = 0.5,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
):
r"""
Function invoked when calling the pipeline for image editing.
diff --git a/examples/community/seed_resize_stable_diffusion.py b/examples/community/seed_resize_stable_diffusion.py
index eafe7572aab5..373095fa2ca2 100644
--- a/examples/community/seed_resize_stable_diffusion.py
+++ b/examples/community/seed_resize_stable_diffusion.py
@@ -80,9 +80,9 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/speech_to_image_diffusion.py b/examples/community/speech_to_image_diffusion.py
index a8ec1620a2eb..695dca2465b4 100644
--- a/examples/community/speech_to_image_diffusion.py
+++ b/examples/community/speech_to_image_diffusion.py
@@ -75,9 +75,9 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_comparison.py b/examples/community/stable_diffusion_comparison.py
index ce6e77c87ff0..8ce9c39227e7 100644
--- a/examples/community/stable_diffusion_comparison.py
+++ b/examples/community/stable_diffusion_comparison.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, List, Optional, Union
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
@@ -81,7 +81,7 @@ def __init__(
self.register_modules(pipeline1=self.pipe1, pipeline2=self.pipe2, pipeline3=self.pipe3, pipeline4=self.pipe4)
@property
- def layers(self) -> Dict[str, Any]:
+ def layers(self) -> dict[str, Any]:
return {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")}
@torch.no_grad()
@@ -95,9 +95,9 @@ def text2img_sd1_1(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
@@ -132,9 +132,9 @@ def text2img_sd1_2(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
@@ -169,9 +169,9 @@ def text2img_sd1_3(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
@@ -206,9 +206,9 @@ def text2img_sd1_4(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
@@ -243,9 +243,9 @@ def _call_(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_controlnet_img2img.py b/examples/community/stable_diffusion_controlnet_img2img.py
index aa116112be1c..03c6fe7f6466 100644
--- a/examples/community/stable_diffusion_controlnet_img2img.py
+++ b/examples/community/stable_diffusion_controlnet_img2img.py
@@ -614,7 +614,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_controlnet_inpaint.py b/examples/community/stable_diffusion_controlnet_inpaint.py
index 6d710e0d73c7..9b76faf56a8a 100644
--- a/examples/community/stable_diffusion_controlnet_inpaint.py
+++ b/examples/community/stable_diffusion_controlnet_inpaint.py
@@ -757,7 +757,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_controlnet_inpaint_img2img.py b/examples/community/stable_diffusion_controlnet_inpaint_img2img.py
index fcb5ed059bb4..299dad58707d 100644
--- a/examples/community/stable_diffusion_controlnet_inpaint_img2img.py
+++ b/examples/community/stable_diffusion_controlnet_inpaint_img2img.py
@@ -745,7 +745,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_controlnet_reference.py b/examples/community/stable_diffusion_controlnet_reference.py
index 74c81b63626d..18c79a0853f9 100644
--- a/examples/community/stable_diffusion_controlnet_reference.py
+++ b/examples/community/stable_diffusion_controlnet_reference.py
@@ -119,7 +119,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_ipex.py b/examples/community/stable_diffusion_ipex.py
index 4f545aa09ded..ddc39b5b7b1a 100644
--- a/examples/community/stable_diffusion_ipex.py
+++ b/examples/community/stable_diffusion_ipex.py
@@ -574,7 +574,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_mega.py b/examples/community/stable_diffusion_mega.py
index c67ebc80b006..a22ecb209ada 100644
--- a/examples/community/stable_diffusion_mega.py
+++ b/examples/community/stable_diffusion_mega.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, List, Optional, Union
import PIL.Image
import torch
@@ -92,7 +92,7 @@ def __init__(
self.register_to_config(requires_safety_checker=requires_safety_checker)
@property
- def components(self) -> Dict[str, Any]:
+ def components(self) -> dict[str, Any]:
return {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")}
@torch.no_grad()
@@ -107,8 +107,8 @@ def inpaint(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: Optional[float] = 0.0,
- generator: Optional[torch.Generator] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
@@ -141,8 +141,8 @@ def img2img(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: Optional[float] = 0.0,
- generator: Optional[torch.Generator] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
@@ -176,9 +176,9 @@ def text2img(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_reference.py b/examples/community/stable_diffusion_reference.py
index d0372bbeba65..7a32827a9350 100644
--- a/examples/community/stable_diffusion_reference.py
+++ b/examples/community/stable_diffusion_reference.py
@@ -276,7 +276,7 @@ def check_inputs(
height: int,
width: int,
callback_steps: Optional[int],
- negative_prompt: Optional[str] = None,
+ negative_prompt: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[torch.Tensor] = None,
@@ -291,7 +291,7 @@ def check_inputs(
height (int): The height of the input image.
width (int): The width of the input image.
callback_steps (Optional[int]): The number of steps to perform the callback on.
- negative_prompt (Optional[str]): The negative prompt text.
+ negative_prompt (str | None): The negative prompt text.
prompt_embeds (Optional[torch.Tensor]): The prompt embeddings.
negative_prompt_embeds (Optional[torch.Tensor]): The negative prompt embeddings.
ip_adapter_image (Optional[torch.Tensor]): The input adapter image.
@@ -411,11 +411,11 @@ def _encode_prompt(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Optional[str],
+ prompt: str | None,
device: torch.device,
num_images_per_prompt: int,
do_classifier_free_guidance: bool,
- negative_prompt: Optional[str] = None,
+ negative_prompt: str | None = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
lora_scale: Optional[float] = None,
@@ -642,7 +642,7 @@ def prepare_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
def prepare_extra_step_kwargs(
self, generator: Union[torch.Generator, List[torch.Generator]], eta: float
- ) -> Dict[str, Any]:
+ ) -> dict[str, Any]:
r"""
Prepare extra keyword arguments for the scheduler step.
@@ -834,7 +834,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_repaint.py b/examples/community/stable_diffusion_repaint.py
index b974e3c7ae74..7248e4101206 100644
--- a/examples/community/stable_diffusion_repaint.py
+++ b/examples/community/stable_diffusion_repaint.py
@@ -625,7 +625,7 @@ def __call__(
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/stable_diffusion_xl_controlnet_reference.py b/examples/community/stable_diffusion_xl_controlnet_reference.py
index 421e67f5bba6..a458ee7c6506 100644
--- a/examples/community/stable_diffusion_xl_controlnet_reference.py
+++ b/examples/community/stable_diffusion_xl_controlnet_reference.py
@@ -345,7 +345,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
diff --git a/examples/community/stable_diffusion_xl_reference.py b/examples/community/stable_diffusion_xl_reference.py
index eb055574966d..c6be397144c5 100644
--- a/examples/community/stable_diffusion_xl_reference.py
+++ b/examples/community/stable_diffusion_xl_reference.py
@@ -290,7 +290,7 @@ def __call__(
negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
diff --git a/examples/community/stable_unclip.py b/examples/community/stable_unclip.py
index f13c4e0a490b..22c3b8c6fe9c 100644
--- a/examples/community/stable_unclip.py
+++ b/examples/community/stable_unclip.py
@@ -190,7 +190,7 @@ def __call__(
width: Optional[int] = None,
num_images_per_prompt: int = 1,
prior_num_inference_steps: int = 25,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
prior_latents: Optional[torch.Tensor] = None,
text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None,
text_attention_mask: Optional[torch.Tensor] = None,
@@ -199,7 +199,7 @@ def __call__(
decoder_num_inference_steps: int = 50,
decoder_num_images_per_prompt: Optional[int] = 1,
decoder_eta: float = 0.0,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
if prompt is not None:
diff --git a/examples/community/text_inpainting.py b/examples/community/text_inpainting.py
index bdf9eca498e7..6df957089883 100644
--- a/examples/community/text_inpainting.py
+++ b/examples/community/text_inpainting.py
@@ -134,9 +134,9 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/tiled_upscaling.py b/examples/community/tiled_upscaling.py
index 7a5e77155cd0..5fe67dce43d8 100644
--- a/examples/community/tiled_upscaling.py
+++ b/examples/community/tiled_upscaling.py
@@ -192,7 +192,7 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/community/unclip_image_interpolation.py b/examples/community/unclip_image_interpolation.py
index 65b52578601e..54e3a9813908 100644
--- a/examples/community/unclip_image_interpolation.py
+++ b/examples/community/unclip_image_interpolation.py
@@ -216,7 +216,7 @@ def __call__(
decoder_latents: Optional[torch.Tensor] = None,
super_res_latents: Optional[torch.Tensor] = None,
decoder_guidance_scale: float = 8.0,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
"""
diff --git a/examples/community/unclip_text_interpolation.py b/examples/community/unclip_text_interpolation.py
index 6fd4f348f48d..e2478811a059 100644
--- a/examples/community/unclip_text_interpolation.py
+++ b/examples/community/unclip_text_interpolation.py
@@ -225,7 +225,7 @@ def __call__(
decoder_guidance_scale: float = 8.0,
enable_sequential_cpu_offload=True,
gpu_id=0,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
"""
diff --git a/examples/community/wildcard_stable_diffusion.py b/examples/community/wildcard_stable_diffusion.py
index d40221e5b1cf..342913b1508e 100644
--- a/examples/community/wildcard_stable_diffusion.py
+++ b/examples/community/wildcard_stable_diffusion.py
@@ -165,9 +165,9 @@ def __call__(
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/custom_diffusion/test_custom_diffusion.py b/examples/custom_diffusion/test_custom_diffusion.py
index 9af84ec7598f..ad18eb246777 100644
--- a/examples/custom_diffusion/test_custom_diffusion.py
+++ b/examples/custom_diffusion/test_custom_diffusion.py
@@ -17,6 +17,9 @@
import os
import sys
import tempfile
+import unittest
+
+from diffusers.utils import is_transformers_version
sys.path.append("..")
@@ -30,6 +33,7 @@
logger.addHandler(stream_handler)
+@unittest.skipIf(is_transformers_version(">=", "4.57.5"), "Size mismatch")
class CustomDiffusion(ExamplesTestsAccelerate):
def test_custom_diffusion(self):
with tempfile.TemporaryDirectory() as tmpdir:
diff --git a/examples/research_projects/anytext/anytext.py b/examples/research_projects/anytext/anytext.py
index 7ae6ae57c22a..1e925e4fafaa 100644
--- a/examples/research_projects/anytext/anytext.py
+++ b/examples/research_projects/anytext/anytext.py
@@ -540,7 +540,7 @@ def __init__(
max_length=77,
freeze=True,
use_fp16=False,
- variant: Optional[str] = None,
+ variant: str | None = None,
):
super().__init__()
self.tokenizer = CLIPTokenizer.from_pretrained("tolgacangoz/anytext", subfolder="tokenizer")
@@ -1018,7 +1018,7 @@ def insert_spaces(self, string, nSpace):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -1938,7 +1938,7 @@ def __call__(
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
- mode: Optional[str] = "generate",
+ mode: str | None = "generate",
draw_pos: Optional[Union[str, torch.Tensor]] = None,
ori_image: Optional[Union[str, torch.Tensor]] = None,
timesteps: List[int] = None,
@@ -1953,7 +1953,7 @@ def __call__(
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
diff --git a/examples/research_projects/anytext/anytext_controlnet.py b/examples/research_projects/anytext/anytext_controlnet.py
index 879d48fc8496..60a998dd27df 100644
--- a/examples/research_projects/anytext/anytext_controlnet.py
+++ b/examples/research_projects/anytext/anytext_controlnet.py
@@ -185,7 +185,7 @@ def __init__(
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+ mid_block_type: str | None = "UNetMidBlock2DCrossAttn",
only_cross_attention: Union[bool, Tuple[bool]] = False,
block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
layers_per_block: int = 2,
@@ -197,12 +197,12 @@ def __init__(
cross_attention_dim: int = 1280,
transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
+ encoder_hid_dim_type: str | None = None,
attention_head_dim: Union[int, Tuple[int, ...]] = 8,
num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
+ class_embed_type: str | None = None,
+ addition_embed_type: str | None = None,
addition_time_embed_dim: Optional[int] = None,
num_class_embeds: Optional[int] = None,
upcast_attention: bool = False,
diff --git a/examples/research_projects/pixart/controlnet_pixart_alpha.py b/examples/research_projects/pixart/controlnet_pixart_alpha.py
index 8f2eb974398d..faf118798ecc 100644
--- a/examples/research_projects/pixart/controlnet_pixart_alpha.py
+++ b/examples/research_projects/pixart/controlnet_pixart_alpha.py
@@ -26,7 +26,7 @@ def __init__(
norm_type: str = "ada_norm_single",
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-6,
- attention_type: Optional[str] = "default",
+ attention_type: str | None = "default",
):
super().__init__()
diff --git a/examples/research_projects/pixart/pipeline_pixart_alpha_controlnet.py b/examples/research_projects/pixart/pipeline_pixart_alpha_controlnet.py
index 89228983d4d8..27273594f474 100644
--- a/examples/research_projects/pixart/pipeline_pixart_alpha_controlnet.py
+++ b/examples/research_projects/pixart/pipeline_pixart_alpha_controlnet.py
@@ -808,7 +808,7 @@ def __call__(
# rc todo: controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
# rc todo: control_guidance_start = 0.0,
# rc todo: control_guidance_end = 1.0,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: int = 1,
diff --git a/examples/research_projects/promptdiffusion/convert_original_promptdiffusion_to_diffusers.py b/examples/research_projects/promptdiffusion/convert_original_promptdiffusion_to_diffusers.py
index c9efcffa5bb8..31105b6a5890 100644
--- a/examples/research_projects/promptdiffusion/convert_original_promptdiffusion_to_diffusers.py
+++ b/examples/research_projects/promptdiffusion/convert_original_promptdiffusion_to_diffusers.py
@@ -1110,7 +1110,7 @@ def stable_unclip_image_encoder(original_config, local_files_only=False):
def stable_unclip_image_noising_components(
- original_config, clip_stats_path: Optional[str] = None, device: Optional[str] = None
+ original_config, clip_stats_path: str | None = None, device: str | None = None
):
"""
Returns the noising components for the img2img and txt2img unclip pipelines.
@@ -1267,9 +1267,9 @@ def download_from_original_stable_diffusion_ckpt(
upcast_attention: Optional[bool] = None,
device: str = None,
from_safetensors: bool = False,
- stable_unclip: Optional[str] = None,
- stable_unclip_prior: Optional[str] = None,
- clip_stats_path: Optional[str] = None,
+ stable_unclip: str | None = None,
+ stable_unclip_prior: str | None = None,
+ clip_stats_path: str | None = None,
controlnet: Optional[bool] = None,
adapter: Optional[bool] = None,
load_safety_checker: bool = True,
diff --git a/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py b/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py
index 233df1276563..8b23570aea77 100644
--- a/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py
+++ b/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py
@@ -923,7 +923,7 @@ def __call__(
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
diff --git a/examples/research_projects/promptdiffusion/promptdiffusioncontrolnet.py b/examples/research_projects/promptdiffusion/promptdiffusioncontrolnet.py
index 7853695f0566..aff0302ac202 100644
--- a/examples/research_projects/promptdiffusion/promptdiffusioncontrolnet.py
+++ b/examples/research_projects/promptdiffusion/promptdiffusioncontrolnet.py
@@ -111,7 +111,7 @@ def __init__(
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+ mid_block_type: str | None = "UNetMidBlock2DCrossAttn",
only_cross_attention: Union[bool, Tuple[bool]] = False,
block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
layers_per_block: int = 2,
@@ -123,12 +123,12 @@ def __init__(
cross_attention_dim: int = 1280,
transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
+ encoder_hid_dim_type: str | None = None,
attention_head_dim: Union[int, Tuple[int, ...]] = 8,
num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
+ class_embed_type: str | None = None,
+ addition_embed_type: str | None = None,
addition_time_embed_dim: Optional[int] = None,
num_class_embeds: Optional[int] = None,
upcast_attention: bool = False,
diff --git a/examples/research_projects/rdm/pipeline_rdm.py b/examples/research_projects/rdm/pipeline_rdm.py
index 9b696874c5d1..08a941438ac4 100644
--- a/examples/research_projects/rdm/pipeline_rdm.py
+++ b/examples/research_projects/rdm/pipeline_rdm.py
@@ -162,10 +162,10 @@ def __call__(
guidance_scale: float = 7.5,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
callback_steps: Optional[int] = 1,
diff --git a/examples/research_projects/sdxl_flax/README.md b/examples/research_projects/sdxl_flax/README.md
index dfbe90e63bde..e2c02d7bb617 100644
--- a/examples/research_projects/sdxl_flax/README.md
+++ b/examples/research_projects/sdxl_flax/README.md
@@ -231,7 +231,7 @@ images = generate(prompt, neg_prompt)
print(f"First inference in {time.time() - start}")
```
-From this point forward, any calls to generate should result in a faster inference
+From this point forward, Any calls to generate should result in a faster inference
time and it won't change.
```python
diff --git a/examples/research_projects/sdxl_flax/sdxl_single_aot.py b/examples/research_projects/sdxl_flax/sdxl_single_aot.py
index 08bd13902aa9..a0f6b243b5df 100644
--- a/examples/research_projects/sdxl_flax/sdxl_single_aot.py
+++ b/examples/research_projects/sdxl_flax/sdxl_single_aot.py
@@ -131,7 +131,7 @@ def generate(prompt, negative_prompt, seed=default_seed, guidance_scale=default_
images = generate(prompt, neg_prompt)
print(f"First inference in {time.time() - start}")
-# 9. From this point forward, any calls to generate should result in a faster inference
+# 9. From this point forward, Any calls to generate should result in a faster inference
# time and it won't change.
start = time.time()
prompt = "photo of a rhino dressed suit and tie sitting at a table in a bar with a bar stools, award winning photography, Elke vogelsang"
diff --git a/examples/research_projects/vae/vae_roundtrip.py b/examples/research_projects/vae/vae_roundtrip.py
index 922cb42615e9..bfb32e484381 100644
--- a/examples/research_projects/vae/vae_roundtrip.py
+++ b/examples/research_projects/vae/vae_roundtrip.py
@@ -41,10 +41,10 @@ def load_vae_model(
*,
device: torch.device,
model_name_or_path: str,
- revision: Optional[str],
- variant: Optional[str],
+ revision: str | None,
+ variant: str | None,
# NOTE: use subfolder="vae" if the pointed model is for stable diffusion as a whole instead of just the VAE
- subfolder: Optional[str],
+ subfolder: str | None,
use_tiny_nn: bool,
) -> SupportedAutoencoder:
if use_tiny_nn:
@@ -156,9 +156,9 @@ def main_kwargs(
device: torch.device,
input_image_path: str,
pretrained_model_name_or_path: str,
- revision: Optional[str],
- variant: Optional[str],
- subfolder: Optional[str],
+ revision: str | None,
+ variant: str | None,
+ subfolder: str | None,
use_tiny_nn: bool,
) -> None:
vae = load_vae_model(
diff --git a/examples/server-async/utils/requestscopedpipeline.py b/examples/server-async/utils/requestscopedpipeline.py
index 9c3276c31c69..d8d2ecd77e2c 100644
--- a/examples/server-async/utils/requestscopedpipeline.py
+++ b/examples/server-async/utils/requestscopedpipeline.py
@@ -68,7 +68,7 @@ def _detect_kernel_pipeline(self, pipeline) -> bool:
return any(hasattr(pipeline, attr) for attr in kernel_indicators)
- def _make_local_scheduler(self, num_inference_steps: int, device: Optional[str] = None, **clone_kwargs):
+ def _make_local_scheduler(self, num_inference_steps: int, device: str | None = None, **clone_kwargs):
base_sched = getattr(self._base, "scheduler", None)
if base_sched is None:
return None
@@ -231,7 +231,7 @@ def _is_tokenizer_component(self, component) -> bool:
def _should_wrap_tokenizers(self) -> bool:
return True
- def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] = None, **kwargs):
+ def generate(self, *args, num_inference_steps: int = 50, device: str | None = None, **kwargs):
local_scheduler = self._make_local_scheduler(num_inference_steps=num_inference_steps, device=device)
try:
diff --git a/scripts/convert_animatediff_sparsectrl_to_diffusers.py b/scripts/convert_animatediff_sparsectrl_to_diffusers.py
index f246dceb97f8..b96b42d60f44 100644
--- a/scripts/convert_animatediff_sparsectrl_to_diffusers.py
+++ b/scripts/convert_animatediff_sparsectrl_to_diffusers.py
@@ -18,7 +18,7 @@
}
-def convert(original_state_dict: Dict[str, nn.Module]) -> Dict[str, nn.Module]:
+def convert(original_state_dict: Dict[str, nn.Module]) -> dict[str, nn.Module]:
converted_state_dict = {}
for key in list(original_state_dict.keys()):
diff --git a/scripts/convert_asymmetric_vqgan_to_diffusers.py b/scripts/convert_asymmetric_vqgan_to_diffusers.py
index ffb735e18224..a9f83a46c2d0 100644
--- a/scripts/convert_asymmetric_vqgan_to_diffusers.py
+++ b/scripts/convert_asymmetric_vqgan_to_diffusers.py
@@ -61,7 +61,7 @@
}
-def convert_asymmetric_autoencoder_kl_state_dict(original_state_dict: Dict[str, Any]) -> Dict[str, Any]:
+def convert_asymmetric_autoencoder_kl_state_dict(original_state_dict: Dict[str, Any]) -> dict[str, Any]:
converted_state_dict = {}
for k, v in original_state_dict.items():
if k.startswith("encoder."):
diff --git a/scripts/convert_cogvideox_to_diffusers.py b/scripts/convert_cogvideox_to_diffusers.py
index 7eeed240c4de..5220aa0bde9d 100644
--- a/scripts/convert_cogvideox_to_diffusers.py
+++ b/scripts/convert_cogvideox_to_diffusers.py
@@ -120,7 +120,7 @@ def replace_up_keys_inplace(key: str, state_dict: Dict[str, Any]):
TOKENIZER_MAX_LENGTH = 226
-def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
+def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]:
state_dict = saved_dict
if "model" in saved_dict.keys():
state_dict = state_dict["model"]
@@ -131,7 +131,7 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
return state_dict
-def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
+def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]:
state_dict[new_key] = state_dict.pop(old_key)
diff --git a/scripts/convert_cosmos_to_diffusers.py b/scripts/convert_cosmos_to_diffusers.py
index bc6014068e87..b6f321b5b2fb 100644
--- a/scripts/convert_cosmos_to_diffusers.py
+++ b/scripts/convert_cosmos_to_diffusers.py
@@ -78,12 +78,85 @@
--save_pipeline
```
+# Cosmos 2.5 Transfer
+
+Download checkpoint
+```bash
+hf download nvidia/Cosmos-Transfer2.5-2B
+```
+
+Convert checkpoint
+```bash
+# depth
+transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/depth/626e6618-bfcd-4d9a-a077-1409e2ce353f_ema_bf16.pt
+
+python scripts/convert_cosmos_to_diffusers.py \
+ --transformer_type Cosmos-2.5-Transfer-General-2B \
+ --transformer_ckpt_path $transformer_ckpt_path \
+ --vae_type wan2.1 \
+ --output_path converted/transfer/2b/general/depth/pipeline \
+ --save_pipeline
+
+python scripts/convert_cosmos_to_diffusers.py \
+ --transformer_type Cosmos-2.5-Transfer-General-2B \
+ --transformer_ckpt_path $transformer_ckpt_path \
+ --vae_type wan2.1 \
+ --output_path converted/transfer/2b/general/depth/models
+
+# edge
+transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/edge/61f5694b-0ad5-4ecd-8ad7-c8545627d125_ema_bf16.pt
+
+python scripts/convert_cosmos_to_diffusers.py \
+ --transformer_type Cosmos-2.5-Transfer-General-2B \
+ --transformer_ckpt_path $transformer_ckpt_path \
+ --vae_type wan2.1 \
+ --output_path converted/transfer/2b/general/edge/pipeline \
+ --save_pipeline
+
+python scripts/convert_cosmos_to_diffusers.py \
+ --transformer_type Cosmos-2.5-Transfer-General-2B \
+ --transformer_ckpt_path $transformer_ckpt_path \
+ --vae_type wan2.1 \
+ --output_path converted/transfer/2b/general/edge/models
+
+# blur
+transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/blur/ba2f44f2-c726-4fe7-949f-597069d9b91c_ema_bf16.pt
+
+python scripts/convert_cosmos_to_diffusers.py \
+ --transformer_type Cosmos-2.5-Transfer-General-2B \
+ --transformer_ckpt_path $transformer_ckpt_path \
+ --vae_type wan2.1 \
+ --output_path converted/transfer/2b/general/blur/pipeline \
+ --save_pipeline
+
+python scripts/convert_cosmos_to_diffusers.py \
+ --transformer_type Cosmos-2.5-Transfer-General-2B \
+ --transformer_ckpt_path $transformer_ckpt_path \
+ --vae_type wan2.1 \
+ --output_path converted/transfer/2b/general/blur/models
+
+# seg
+transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/seg/5136ef49-6d8d-42e8-8abf-7dac722a304a_ema_bf16.pt
+
+python scripts/convert_cosmos_to_diffusers.py \
+ --transformer_type Cosmos-2.5-Transfer-General-2B \
+ --transformer_ckpt_path $transformer_ckpt_path \
+ --vae_type wan2.1 \
+ --output_path converted/transfer/2b/general/seg/pipeline \
+ --save_pipeline
+
+python scripts/convert_cosmos_to_diffusers.py \
+ --transformer_type Cosmos-2.5-Transfer-General-2B \
+ --transformer_ckpt_path $transformer_ckpt_path \
+ --vae_type wan2.1 \
+ --output_path converted/transfer/2b/general/seg/models
+```
"""
import argparse
import pathlib
import sys
-from typing import Any, Dict
+from typing import Any, Dict, Optional
import torch
from accelerate import init_empty_weights
@@ -95,6 +168,7 @@
AutoencoderKLWan,
Cosmos2TextToImagePipeline,
Cosmos2VideoToWorldPipeline,
+ CosmosControlNetModel,
CosmosTextToWorldPipeline,
CosmosTransformer3DModel,
CosmosVideoToWorldPipeline,
@@ -103,13 +177,14 @@
UniPCMultistepScheduler,
)
from diffusers.pipelines.cosmos.pipeline_cosmos2_5_predict import Cosmos2_5_PredictBasePipeline
+from diffusers.pipelines.cosmos.pipeline_cosmos2_5_transfer import Cosmos2_5_TransferPipeline
def remove_keys_(key: str, state_dict: Dict[str, Any]):
state_dict.pop(key)
-def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
+def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]:
state_dict[new_key] = state_dict.pop(old_key)
@@ -356,8 +431,62 @@ def rename_transformer_blocks_(key: str, state_dict: Dict[str, Any]):
"crossattn_proj_in_channels": 100352,
"encoder_hidden_states_channels": 1024,
},
+ "Cosmos-2.5-Transfer-General-2B": {
+ "in_channels": 16 + 1,
+ "out_channels": 16,
+ "num_attention_heads": 16,
+ "attention_head_dim": 128,
+ "num_layers": 28,
+ "mlp_ratio": 4.0,
+ "text_embed_dim": 1024,
+ "adaln_lora_dim": 256,
+ "max_size": (128, 240, 240),
+ "patch_size": (1, 2, 2),
+ "rope_scale": (1.0, 3.0, 3.0),
+ "concat_padding_mask": True,
+ "extra_pos_embed_type": None,
+ "use_crossattn_projection": True,
+ "crossattn_proj_in_channels": 100352,
+ "encoder_hidden_states_channels": 1024,
+ "controlnet_block_every_n": 7,
+ "img_context_dim_in": 1152,
+ "img_context_dim_out": 2048,
+ "img_context_num_tokens": 256,
+ },
+}
+
+CONTROLNET_CONFIGS = {
+ "Cosmos-2.5-Transfer-General-2B": {
+ "n_controlnet_blocks": 4,
+ "model_channels": 2048,
+ "in_channels": 130,
+ "latent_channels": 18, # (16 latent + 1 condition_mask) + 1 padding_mask = 18
+ "num_attention_heads": 16,
+ "attention_head_dim": 128,
+ "mlp_ratio": 4.0,
+ "text_embed_dim": 1024,
+ "adaln_lora_dim": 256,
+ "patch_size": (1, 2, 2),
+ "max_size": (128, 240, 240),
+ "rope_scale": (1.0, 3.0, 3.0),
+ "extra_pos_embed_type": None,
+ "img_context_dim_in": 1152,
+ "img_context_dim_out": 2048,
+ "use_crossattn_projection": True,
+ "crossattn_proj_in_channels": 100352,
+ "encoder_hidden_states_channels": 1024,
+ },
+}
+
+CONTROLNET_KEYS_RENAME_DICT = {
+ **TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0,
+ "blocks": "blocks",
+ "control_embedder.proj.1": "patch_embed.proj",
}
+
+CONTROLNET_SPECIAL_KEYS_REMAP = {**TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0}
+
VAE_KEYS_RENAME_DICT = {
"down.0": "down_blocks.0",
"down.1": "down_blocks.1",
@@ -436,7 +565,7 @@ def rename_transformer_blocks_(key: str, state_dict: Dict[str, Any]):
}
-def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
+def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]:
state_dict = saved_dict
if "model" in saved_dict.keys():
state_dict = state_dict["model"]
@@ -447,9 +576,12 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
return state_dict
-def convert_transformer(transformer_type: str, ckpt_path: str, weights_only: bool = True):
+def convert_transformer(
+ transformer_type: str,
+ state_dict: Optional[Dict[str, Any]] = None,
+ weights_only: bool = True,
+):
PREFIX_KEY = "net."
- original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=weights_only))
if "Cosmos-1.0" in transformer_type:
TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0
@@ -467,23 +599,29 @@ def convert_transformer(transformer_type: str, ckpt_path: str, weights_only: boo
config = TRANSFORMER_CONFIGS[transformer_type]
transformer = CosmosTransformer3DModel(**config)
- for key in list(original_state_dict.keys()):
+ old2new = {}
+ new2old = {}
+ for key in list(state_dict.keys()):
new_key = key[:]
if new_key.startswith(PREFIX_KEY):
new_key = new_key.removeprefix(PREFIX_KEY)
for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():
new_key = new_key.replace(replace_key, rename_key)
print(key, "->", new_key, flush=True)
- update_state_dict_(original_state_dict, key, new_key)
+ assert new_key not in new2old, f"new key {new_key} already mapped"
+ assert key not in old2new, f"old key {key} already mapped"
+ old2new[key] = new_key
+ new2old[new_key] = key
+ update_state_dict_(state_dict, key, new_key)
- for key in list(original_state_dict.keys()):
+ for key in list(state_dict.keys()):
for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items():
if special_key not in key:
continue
- handler_fn_inplace(key, original_state_dict)
+ handler_fn_inplace(key, state_dict)
expected_keys = set(transformer.state_dict().keys())
- mapped_keys = set(original_state_dict.keys())
+ mapped_keys = set(state_dict.keys())
missing_keys = expected_keys - mapped_keys
unexpected_keys = mapped_keys - expected_keys
if missing_keys:
@@ -497,10 +635,86 @@ def convert_transformer(transformer_type: str, ckpt_path: str, weights_only: boo
print(k)
sys.exit(2)
- transformer.load_state_dict(original_state_dict, strict=True, assign=True)
+ transformer.load_state_dict(state_dict, strict=True, assign=True)
return transformer
+def convert_controlnet(
+ transformer_type: str,
+ control_state_dict: Dict[str, Any],
+ base_state_dict: Dict[str, Any],
+ weights_only: bool = True,
+):
+ """
+ Convert controlnet weights.
+
+ Args:
+ transformer_type: The type of transformer/controlnet
+ control_state_dict: State dict containing controlnet-specific weights
+ base_state_dict: State dict containing base transformer weights (for shared modules)
+ weights_only: Whether to use weights_only loading
+ """
+ if transformer_type not in CONTROLNET_CONFIGS:
+ raise AssertionError(f"{transformer_type} does not define a ControlNet config")
+
+ PREFIX_KEY = "net."
+
+ # Process control-specific keys
+ for key in list(control_state_dict.keys()):
+ new_key = key[:]
+ if new_key.startswith(PREFIX_KEY):
+ new_key = new_key.removeprefix(PREFIX_KEY)
+ for replace_key, rename_key in CONTROLNET_KEYS_RENAME_DICT.items():
+ new_key = new_key.replace(replace_key, rename_key)
+ update_state_dict_(control_state_dict, key, new_key)
+
+ for key in list(control_state_dict.keys()):
+ for special_key, handler_fn_inplace in CONTROLNET_SPECIAL_KEYS_REMAP.items():
+ if special_key not in key:
+ continue
+ handler_fn_inplace(key, control_state_dict)
+
+ # Copy shared weights from base transformer to controlnet
+ # These are the duplicated modules: patch_embed_base, time_embed, learnable_pos_embed, img_context_proj, crossattn_proj
+ shared_module_mappings = {
+ # transformer key prefix -> controlnet key prefix
+ "patch_embed.": "patch_embed_base.",
+ "time_embed.": "time_embed.",
+ "learnable_pos_embed.": "learnable_pos_embed.",
+ "img_context_proj.": "img_context_proj.",
+ "crossattn_proj.": "crossattn_proj.",
+ }
+
+ for key in list(base_state_dict.keys()):
+ for transformer_prefix, controlnet_prefix in shared_module_mappings.items():
+ if key.startswith(transformer_prefix):
+ controlnet_key = controlnet_prefix + key[len(transformer_prefix) :]
+ control_state_dict[controlnet_key] = base_state_dict[key].clone()
+ print(f"Copied shared weight: {key} -> {controlnet_key}", flush=True)
+ break
+
+ cfg = CONTROLNET_CONFIGS[transformer_type]
+ controlnet = CosmosControlNetModel(**cfg)
+
+ expected_keys = set(controlnet.state_dict().keys())
+ mapped_keys = set(control_state_dict.keys())
+ missing_keys = expected_keys - mapped_keys
+ unexpected_keys = mapped_keys - expected_keys
+ if missing_keys:
+ print(f"WARNING: missing controlnet keys ({len(missing_keys)}):", file=sys.stderr, flush=True)
+ for k in sorted(missing_keys):
+ print(k, file=sys.stderr)
+ sys.exit(3)
+ if unexpected_keys:
+ print(f"WARNING: unexpected controlnet keys ({len(unexpected_keys)}):", file=sys.stderr, flush=True)
+ for k in sorted(unexpected_keys):
+ print(k, file=sys.stderr)
+ sys.exit(4)
+
+ controlnet.load_state_dict(control_state_dict, strict=True, assign=True)
+ return controlnet
+
+
def convert_vae(vae_type: str):
model_name = VAE_CONFIGS[vae_type]["name"]
snapshot_directory = snapshot_download(model_name, repo_type="model")
@@ -586,7 +800,7 @@ def save_pipeline_cosmos_2_0(args, transformer, vae):
pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
-def save_pipeline_cosmos2_5(args, transformer, vae):
+def save_pipeline_cosmos2_5_predict(args, transformer, vae):
text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B"
tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct"
@@ -614,6 +828,35 @@ def save_pipeline_cosmos2_5(args, transformer, vae):
pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
+def save_pipeline_cosmos2_5_transfer(args, transformer, controlnet, vae):
+ text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B"
+ tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct"
+
+ text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+ text_encoder_path, torch_dtype="auto", device_map="cpu"
+ )
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
+ scheduler = UniPCMultistepScheduler(
+ use_karras_sigmas=True,
+ use_flow_sigmas=True,
+ prediction_type="flow_prediction",
+ sigma_max=200.0,
+ sigma_min=0.01,
+ )
+
+ pipe = Cosmos2_5_TransferPipeline(
+ text_encoder=text_encoder,
+ tokenizer=tokenizer,
+ transformer=transformer,
+ controlnet=controlnet,
+ vae=vae,
+ scheduler=scheduler,
+ safety_checker=lambda *args, **kwargs: None,
+ )
+ pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
+
+
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--transformer_type", type=str, default=None, choices=list(TRANSFORMER_CONFIGS.keys()))
@@ -642,18 +885,61 @@ def get_args():
args = get_args()
transformer = None
+ controlnet = None
dtype = DTYPE_MAPPING[args.dtype]
if args.save_pipeline:
assert args.transformer_ckpt_path is not None
assert args.vae_type is not None
+ raw_state_dict = None
if args.transformer_ckpt_path is not None:
weights_only = "Cosmos-1.0" in args.transformer_type
- transformer = convert_transformer(args.transformer_type, args.transformer_ckpt_path, weights_only)
- transformer = transformer.to(dtype=dtype)
- if not args.save_pipeline:
- transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
+ raw_state_dict = get_state_dict(
+ torch.load(args.transformer_ckpt_path, map_location="cpu", weights_only=weights_only)
+ )
+
+ if raw_state_dict is not None:
+ if "Transfer" in args.transformer_type:
+ base_state_dict = {}
+ control_state_dict = {}
+ for k, v in raw_state_dict.items():
+ plain_key = k.removeprefix("net.") if k.startswith("net.") else k
+ if "control" in plain_key.lower():
+ control_state_dict[k] = v
+ else:
+ base_state_dict[k] = v
+ assert len(base_state_dict.keys() & control_state_dict.keys()) == 0
+
+ # Convert transformer first to get the processed base state dict
+ transformer = convert_transformer(
+ args.transformer_type, state_dict=base_state_dict, weights_only=weights_only
+ )
+ transformer = transformer.to(dtype=dtype)
+
+ # Get converted transformer state dict to copy shared weights to controlnet
+ converted_base_state_dict = transformer.state_dict()
+
+ # Convert controlnet with both control-specific and shared weights from transformer
+ controlnet = convert_controlnet(
+ args.transformer_type, control_state_dict, converted_base_state_dict, weights_only=weights_only
+ )
+ controlnet = controlnet.to(dtype=dtype)
+
+ if not args.save_pipeline:
+ transformer.save_pretrained(
+ pathlib.Path(args.output_path) / "transformer", safe_serialization=True, max_shard_size="5GB"
+ )
+ controlnet.save_pretrained(
+ pathlib.Path(args.output_path) / "controlnet", safe_serialization=True, max_shard_size="5GB"
+ )
+ else:
+ transformer = convert_transformer(
+ args.transformer_type, state_dict=raw_state_dict, weights_only=weights_only
+ )
+ transformer = transformer.to(dtype=dtype)
+ if not args.save_pipeline:
+ transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
if args.vae_type is not None:
if "Cosmos-1.0" in args.transformer_type:
@@ -667,6 +953,8 @@ def get_args():
if not args.save_pipeline:
vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
+ else:
+ vae = None
if args.save_pipeline:
if "Cosmos-1.0" in args.transformer_type:
@@ -678,6 +966,11 @@ def get_args():
assert args.tokenizer_path is not None
save_pipeline_cosmos_2_0(args, transformer, vae)
elif "Cosmos-2.5" in args.transformer_type:
- save_pipeline_cosmos2_5(args, transformer, vae)
+ if "Predict" in args.transformer_type:
+ save_pipeline_cosmos2_5_predict(args, transformer, vae)
+ elif "Transfer" in args.transformer_type:
+ save_pipeline_cosmos2_5_transfer(args, transformer, None, vae)
+ else:
+ raise AssertionError(f"{args.transformer_type} not supported")
else:
raise AssertionError(f"{args.transformer_type} not supported")
diff --git a/scripts/convert_dcae_to_diffusers.py b/scripts/convert_dcae_to_diffusers.py
index 15f79a8154e6..1a94c4b06640 100644
--- a/scripts/convert_dcae_to_diffusers.py
+++ b/scripts/convert_dcae_to_diffusers.py
@@ -79,7 +79,7 @@ def remap_proj_conv_(key: str, state_dict: Dict[str, Any]):
}
-def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
+def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]:
state_dict = saved_dict
if "model" in saved_dict.keys():
state_dict = state_dict["model"]
@@ -90,7 +90,7 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
return state_dict
-def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
+def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]:
state_dict[new_key] = state_dict.pop(old_key)
diff --git a/scripts/convert_hunyuan_video_to_diffusers.py b/scripts/convert_hunyuan_video_to_diffusers.py
index c84809d7f68a..fb604c98b65b 100644
--- a/scripts/convert_hunyuan_video_to_diffusers.py
+++ b/scripts/convert_hunyuan_video_to_diffusers.py
@@ -203,11 +203,11 @@ def remap_single_transformer_blocks_(key, state_dict):
}
-def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
+def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]:
state_dict[new_key] = state_dict.pop(old_key)
-def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
+def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]:
state_dict = saved_dict
if "model" in saved_dict.keys():
state_dict = state_dict["model"]
diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 88a9c8700a38..9dee954af6d0 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -1,7 +1,7 @@
import argparse
import os
from contextlib import nullcontext
-from typing import Any, Dict, Optional, Tuple
+from typing import Any
import safetensors.torch
import torch
@@ -94,15 +94,15 @@
}
-def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> None:
+def update_state_dict_inplace(state_dict: dict[str, Any], old_key: str, new_key: str) -> None:
state_dict[new_key] = state_dict.pop(old_key)
-def remove_keys_inplace(key: str, state_dict: Dict[str, Any]) -> None:
+def remove_keys_inplace(key: str, state_dict: dict[str, Any]) -> None:
state_dict.pop(key)
-def convert_ltx2_transformer_adaln_single(key: str, state_dict: Dict[str, Any]) -> None:
+def convert_ltx2_transformer_adaln_single(key: str, state_dict: dict[str, Any]) -> None:
# Skip if not a weight, bias
if ".weight" not in key and ".bias" not in key:
return
@@ -120,7 +120,7 @@ def convert_ltx2_transformer_adaln_single(key: str, state_dict: Dict[str, Any])
return
-def convert_ltx2_audio_vae_per_channel_statistics(key: str, state_dict: Dict[str, Any]) -> None:
+def convert_ltx2_audio_vae_per_channel_statistics(key: str, state_dict: dict[str, Any]) -> None:
if key.startswith("per_channel_statistics"):
new_key = ".".join(["decoder", key])
param = state_dict.pop(key)
@@ -156,7 +156,7 @@ def convert_ltx2_audio_vae_per_channel_statistics(key: str, state_dict: Dict[str
LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP = {}
-def split_transformer_and_connector_state_dict(state_dict: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+def split_transformer_and_connector_state_dict(state_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
connector_prefixes = (
"video_embeddings_connector",
"audio_embeddings_connector",
@@ -178,7 +178,7 @@ def split_transformer_and_connector_state_dict(state_dict: Dict[str, Any]) -> Tu
return transformer_state_dict, connector_state_dict
-def get_ltx2_transformer_config(version: str) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
+def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
if version == "test":
# Produces a transformer of the same size as used in test_models_transformer_ltx2.py
config = {
@@ -270,7 +270,7 @@ def get_ltx2_transformer_config(version: str) -> Tuple[Dict[str, Any], Dict[str,
return config, rename_dict, special_keys_remap
-def get_ltx2_connectors_config(version: str) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
+def get_ltx2_connectors_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
if version == "test":
config = {
"model_id": "diffusers-internal-dev/dummy-ltx2",
@@ -319,7 +319,7 @@ def get_ltx2_connectors_config(version: str) -> Tuple[Dict[str, Any], Dict[str,
return config, rename_dict, special_keys_remap
-def convert_ltx2_transformer(original_state_dict: Dict[str, Any], version: str) -> Dict[str, Any]:
+def convert_ltx2_transformer(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_transformer_config(version)
diffusers_config = config["diffusers_config"]
@@ -347,7 +347,7 @@ def convert_ltx2_transformer(original_state_dict: Dict[str, Any], version: str)
return transformer
-def convert_ltx2_connectors(original_state_dict: Dict[str, Any], version: str) -> LTX2TextConnectors:
+def convert_ltx2_connectors(original_state_dict: dict[str, Any], version: str) -> LTX2TextConnectors:
config, rename_dict, special_keys_remap = get_ltx2_connectors_config(version)
diffusers_config = config["diffusers_config"]
@@ -376,7 +376,7 @@ def convert_ltx2_connectors(original_state_dict: Dict[str, Any], version: str) -
def get_ltx2_video_vae_config(
version: str, timestep_conditioning: bool = False
-) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
+) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
if version == "test":
config = {
"model_id": "diffusers-internal-dev/dummy-ltx2",
@@ -455,8 +455,8 @@ def get_ltx2_video_vae_config(
def convert_ltx2_video_vae(
- original_state_dict: Dict[str, Any], version: str, timestep_conditioning: bool
-) -> Dict[str, Any]:
+ original_state_dict: dict[str, Any], version: str, timestep_conditioning: bool
+) -> dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_video_vae_config(version, timestep_conditioning)
diffusers_config = config["diffusers_config"]
@@ -482,7 +482,7 @@ def convert_ltx2_video_vae(
return vae
-def get_ltx2_audio_vae_config(version: str) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
+def get_ltx2_audio_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
if version == "2.0":
config = {
"model_id": "diffusers-internal-dev/new-ltx-model",
@@ -511,7 +511,7 @@ def get_ltx2_audio_vae_config(version: str) -> Tuple[Dict[str, Any], Dict[str, A
return config, rename_dict, special_keys_remap
-def convert_ltx2_audio_vae(original_state_dict: Dict[str, Any], version: str) -> Dict[str, Any]:
+def convert_ltx2_audio_vae(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_audio_vae_config(version)
diffusers_config = config["diffusers_config"]
@@ -537,7 +537,7 @@ def convert_ltx2_audio_vae(original_state_dict: Dict[str, Any], version: str) ->
return vae
-def get_ltx2_vocoder_config(version: str) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
+def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
if version == "2.0":
config = {
"model_id": "diffusers-internal-dev/new-ltx-model",
@@ -558,7 +558,7 @@ def get_ltx2_vocoder_config(version: str) -> Tuple[Dict[str, Any], Dict[str, Any
return config, rename_dict, special_keys_remap
-def convert_ltx2_vocoder(original_state_dict: Dict[str, Any], version: str) -> Dict[str, Any]:
+def convert_ltx2_vocoder(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_vocoder_config(version)
diffusers_config = config["diffusers_config"]
@@ -601,7 +601,7 @@ def get_ltx2_spatial_latent_upsampler_config(version: str):
def convert_ltx2_spatial_latent_upsampler(
- original_state_dict: Dict[str, Any], config: Dict[str, Any], dtype: torch.dtype
+ original_state_dict: dict[str, Any], config: dict[str, Any], dtype: torch.dtype
):
with init_empty_weights():
latent_upsampler = LTX2LatentUpsamplerModel(**config)
@@ -611,7 +611,7 @@ def convert_ltx2_spatial_latent_upsampler(
return latent_upsampler
-def load_original_checkpoint(args, filename: Optional[str]) -> Dict[str, Any]:
+def load_original_checkpoint(args, filename: str | None) -> dict[str, Any]:
if args.original_state_dict_repo_id is not None:
ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=filename)
elif args.checkpoint_path is not None:
@@ -623,7 +623,7 @@ def load_original_checkpoint(args, filename: Optional[str]) -> Dict[str, Any]:
return original_state_dict
-def load_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Optional[str] = None) -> Dict[str, Any]:
+def load_hub_or_local_checkpoint(repo_id: str | None = None, filename: str | None = None) -> dict[str, Any]:
if repo_id is None and filename is None:
raise ValueError("Please supply at least one of `repo_id` or `filename`")
@@ -643,7 +643,7 @@ def load_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Option
return state_dict
-def get_model_state_dict_from_combined_ckpt(combined_ckpt: Dict[str, Any], prefix: str) -> Dict[str, Any]:
+def get_model_state_dict_from_combined_ckpt(combined_ckpt: dict[str, Any], prefix: str) -> dict[str, Any]:
# Ensure that the key prefix ends with a dot (.)
if not prefix.endswith("."):
prefix = prefix + "."
diff --git a/scripts/convert_ltx_to_diffusers.py b/scripts/convert_ltx_to_diffusers.py
index 19e5602039e5..b4ae57f9c8f0 100644
--- a/scripts/convert_ltx_to_diffusers.py
+++ b/scripts/convert_ltx_to_diffusers.py
@@ -116,7 +116,7 @@ def remove_keys_(key: str, state_dict: Dict[str, Any]):
}
-def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
+def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]:
state_dict = saved_dict
if "model" in saved_dict.keys():
state_dict = state_dict["model"]
@@ -127,7 +127,7 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
return state_dict
-def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
+def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]:
state_dict[new_key] = state_dict.pop(old_key)
@@ -192,7 +192,7 @@ def convert_spatial_latent_upsampler(ckpt_path: str, config, dtype: torch.dtype)
return latent_upsampler
-def get_transformer_config(version: str) -> Dict[str, Any]:
+def get_transformer_config(version: str) -> dict[str, Any]:
if version == "0.9.7":
config = {
"in_channels": 128,
@@ -232,7 +232,7 @@ def get_transformer_config(version: str) -> Dict[str, Any]:
return config
-def get_vae_config(version: str) -> Dict[str, Any]:
+def get_vae_config(version: str) -> dict[str, Any]:
if version in ["0.9.0"]:
config = {
"in_channels": 3,
@@ -359,7 +359,7 @@ def get_vae_config(version: str) -> Dict[str, Any]:
return config
-def get_spatial_latent_upsampler_config(version: str) -> Dict[str, Any]:
+def get_spatial_latent_upsampler_config(version: str) -> dict[str, Any]:
if version == "0.9.7":
config = {
"in_channels": 128,
diff --git a/scripts/convert_prx_to_diffusers.py b/scripts/convert_prx_to_diffusers.py
index d9bde2f34d56..00bb3f6fe99e 100644
--- a/scripts/convert_prx_to_diffusers.py
+++ b/scripts/convert_prx_to_diffusers.py
@@ -86,7 +86,7 @@ def create_parameter_mapping(depth: int) -> dict:
return mapping
-def convert_checkpoint_parameters(old_state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
+def convert_checkpoint_parameters(old_state_dict: Dict[str, torch.Tensor], depth: int) -> dict[str, torch.Tensor]:
"""Convert old checkpoint parameters to new diffusers format."""
print("Converting checkpoint parameters...")
diff --git a/scripts/convert_skyreelsv2_to_diffusers.py b/scripts/convert_skyreelsv2_to_diffusers.py
index 3bc3c435685b..29d5a7666a6e 100644
--- a/scripts/convert_skyreelsv2_to_diffusers.py
+++ b/scripts/convert_skyreelsv2_to_diffusers.py
@@ -66,7 +66,7 @@
TRANSFORMER_SPECIAL_KEYS_REMAP = {}
-def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
+def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]:
state_dict[new_key] = state_dict.pop(old_key)
@@ -81,7 +81,7 @@ def load_sharded_safetensors(dir: pathlib.Path):
return state_dict
-def get_transformer_config(model_type: str) -> Dict[str, Any]:
+def get_transformer_config(model_type: str) -> dict[str, Any]:
if model_type == "SkyReels-V2-DF-1.3B-540P":
config = {
"model_id": "Skywork/SkyReels-V2-DF-1.3B-540P",
diff --git a/scripts/convert_wan_to_diffusers.py b/scripts/convert_wan_to_diffusers.py
index 06f87409262a..eca1af61ae69 100644
--- a/scripts/convert_wan_to_diffusers.py
+++ b/scripts/convert_wan_to_diffusers.py
@@ -313,7 +313,7 @@ def convert_animate_face_adapter_weights(key: str, state_dict: Dict[str, Any]) -
}
-def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
+def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]:
state_dict[new_key] = state_dict.pop(old_key)
diff --git a/setup.py b/setup.py
index d52d37787fdb..45bffd376351 100644
--- a/setup.py
+++ b/setup.py
@@ -101,6 +101,7 @@
"datasets",
"filelock",
"flax>=0.4.1",
+ "ftfy",
"hf-doc-builder>=0.3.0",
"httpx<1.0.0",
"huggingface-hub>=0.34.0,<2.0",
@@ -111,7 +112,6 @@
"jax>=0.4.1",
"jaxlib>=0.4.1",
"Jinja2",
- "k-diffusion==0.0.12",
"torchsde",
"note_seq",
"librosa",
@@ -122,7 +122,7 @@
"pytest",
"pytest-timeout",
"pytest-xdist",
- "python>=3.8.0",
+ "python>=3.10.0",
"ruff==0.9.10",
"safetensors>=0.3.1",
"sentencepiece>=0.1.91,!=0.1.92",
@@ -222,13 +222,14 @@ def run(self):
extras["training"] = deps_list("accelerate", "datasets", "protobuf", "tensorboard", "Jinja2", "peft", "timm")
extras["test"] = deps_list(
"compel",
+ "ftfy",
"GitPython",
"datasets",
"Jinja2",
"invisible-watermark",
- "k-diffusion",
"librosa",
"parameterized",
+ "protobuf",
"pytest",
"pytest-timeout",
"pytest-xdist",
@@ -237,6 +238,7 @@ def run(self):
"sentencepiece",
"scipy",
"tiktoken",
+ "torchsde",
"torchvision",
"transformers",
"phonemizer",
@@ -287,7 +289,7 @@ def run(self):
packages=find_packages("src"),
package_data={"diffusers": ["py.typed"]},
include_package_data=True,
- python_requires=">=3.8.0",
+ python_requires=">=3.10.0",
install_requires=list(install_requires),
extras_require=extras,
entry_points={"console_scripts": ["diffusers-cli=diffusers.commands.diffusers_cli:main"]},
diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
index 8be8f472591f..1fc0914fe09e 100644
--- a/src/diffusers/__init__.py
+++ b/src/diffusers/__init__.py
@@ -10,7 +10,6 @@
is_bitsandbytes_available,
is_flax_available,
is_gguf_available,
- is_k_diffusion_available,
is_librosa_available,
is_note_seq_available,
is_nvidia_modelopt_available,
@@ -50,8 +49,6 @@
"is_flax_available",
"is_inflect_available",
"is_invisible_watermark_available",
- "is_k_diffusion_available",
- "is_k_diffusion_version",
"is_librosa_available",
"is_note_seq_available",
"is_onnx_available",
@@ -221,6 +218,7 @@
"ControlNetModel",
"ControlNetUnionModel",
"ControlNetXSAdapter",
+ "CosmosControlNetModel",
"CosmosTransformer3DModel",
"DiTTransformer2DModel",
"EasyAnimateTransformer3DModel",
@@ -293,10 +291,17 @@
)
_import_structure["modular_pipelines"].extend(
[
+ "AutoPipelineBlocks",
"ComponentsManager",
"ComponentSpec",
+ "ConditionalPipelineBlocks",
+ "ConfigSpec",
+ "InputParam",
+ "LoopSequentialPipelineBlocks",
"ModularPipeline",
"ModularPipelineBlocks",
+ "OutputParam",
+ "SequentialPipelineBlocks",
]
)
_import_structure["optimization"] = [
@@ -485,6 +490,7 @@
"CogView4Pipeline",
"ConsisIDPipeline",
"Cosmos2_5_PredictBasePipeline",
+ "Cosmos2_5_TransferPipeline",
"Cosmos2TextToImagePipeline",
"Cosmos2VideoToWorldPipeline",
"CosmosTextToWorldPipeline",
@@ -722,19 +728,6 @@
else:
_import_structure["pipelines"].extend(["ConsisIDPipeline"])
-try:
- if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- from .utils import dummy_torch_and_transformers_and_k_diffusion_objects # noqa F403
-
- _import_structure["utils.dummy_torch_and_transformers_and_k_diffusion_objects"] = [
- name for name in dir(dummy_torch_and_transformers_and_k_diffusion_objects) if not name.startswith("_")
- ]
-
-else:
- _import_structure["pipelines"].extend(["StableDiffusionKDiffusionPipeline", "StableDiffusionXLKDiffusionPipeline"])
-
try:
if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
raise OptionalDependencyNotAvailable()
@@ -992,6 +985,7 @@
ControlNetModel,
ControlNetUnionModel,
ControlNetXSAdapter,
+ CosmosControlNetModel,
CosmosTransformer3DModel,
DiTTransformer2DModel,
EasyAnimateTransformer3DModel,
@@ -1060,7 +1054,19 @@
ZImageTransformer2DModel,
attention_backend,
)
- from .modular_pipelines import ComponentsManager, ComponentSpec, ModularPipeline, ModularPipelineBlocks
+ from .modular_pipelines import (
+ AutoPipelineBlocks,
+ ComponentsManager,
+ ComponentSpec,
+ ConditionalPipelineBlocks,
+ ConfigSpec,
+ InputParam,
+ LoopSequentialPipelineBlocks,
+ ModularPipeline,
+ ModularPipelineBlocks,
+ OutputParam,
+ SequentialPipelineBlocks,
+ )
from .optimization import (
get_constant_schedule,
get_constant_schedule_with_warmup,
@@ -1226,6 +1232,7 @@
CogView4Pipeline,
ConsisIDPipeline,
Cosmos2_5_PredictBasePipeline,
+ Cosmos2_5_TransferPipeline,
Cosmos2TextToImagePipeline,
Cosmos2VideoToWorldPipeline,
CosmosTextToWorldPipeline,
@@ -1446,14 +1453,6 @@
ZImagePipeline,
)
- try:
- if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import * # noqa F403
- else:
- from .pipelines import StableDiffusionKDiffusionPipeline, StableDiffusionXLKDiffusionPipeline
-
try:
if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
raise OptionalDependencyNotAvailable()
diff --git a/src/diffusers/callbacks.py b/src/diffusers/callbacks.py
index 2a08f091d9f3..087a6b7fee56 100644
--- a/src/diffusers/callbacks.py
+++ b/src/diffusers/callbacks.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any
from .configuration_utils import ConfigMixin, register_to_config
from .utils import CONFIG_NAME
@@ -33,13 +33,13 @@ def __init__(self, cutoff_step_ratio=1.0, cutoff_step_index=None):
raise ValueError("cutoff_step_ratio must be a float between 0.0 and 1.0.")
@property
- def tensor_inputs(self) -> List[str]:
+ def tensor_inputs(self) -> list[str]:
raise NotImplementedError(f"You need to set the attribute `tensor_inputs` for {self.__class__}")
- def callback_fn(self, pipeline, step_index, timesteps, callback_kwargs) -> Dict[str, Any]:
+ def callback_fn(self, pipeline, step_index, timesteps, callback_kwargs) -> dict[str, Any]:
raise NotImplementedError(f"You need to implement the method `callback_fn` for {self.__class__}")
- def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
+ def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
return self.callback_fn(pipeline, step_index, timestep, callback_kwargs)
@@ -49,14 +49,14 @@ class MultiPipelineCallbacks:
provides a unified interface for calling all of them.
"""
- def __init__(self, callbacks: List[PipelineCallback]):
+ def __init__(self, callbacks: list[PipelineCallback]):
self.callbacks = callbacks
@property
- def tensor_inputs(self) -> List[str]:
+ def tensor_inputs(self) -> list[str]:
return [input for callback in self.callbacks for input in callback.tensor_inputs]
- def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
+ def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
"""
Calls all the callbacks in order with the given arguments and returns the final callback_kwargs.
"""
@@ -76,7 +76,7 @@ class SDCFGCutoffCallback(PipelineCallback):
tensor_inputs = ["prompt_embeds"]
- def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
cutoff_step_ratio = self.config.cutoff_step_ratio
cutoff_step_index = self.config.cutoff_step_index
@@ -109,7 +109,7 @@ class SDXLCFGCutoffCallback(PipelineCallback):
"add_time_ids",
]
- def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
cutoff_step_ratio = self.config.cutoff_step_ratio
cutoff_step_index = self.config.cutoff_step_index
@@ -152,7 +152,7 @@ class SDXLControlnetCFGCutoffCallback(PipelineCallback):
"image",
]
- def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
cutoff_step_ratio = self.config.cutoff_step_ratio
cutoff_step_index = self.config.cutoff_step_index
@@ -195,7 +195,7 @@ class IPAdapterScaleCutoffCallback(PipelineCallback):
tensor_inputs = []
- def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
cutoff_step_ratio = self.config.cutoff_step_ratio
cutoff_step_index = self.config.cutoff_step_index
@@ -219,7 +219,7 @@ class SD3CFGCutoffCallback(PipelineCallback):
tensor_inputs = ["prompt_embeds", "pooled_prompt_embeds"]
- def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
cutoff_step_ratio = self.config.cutoff_step_ratio
cutoff_step_index = self.config.cutoff_step_index
diff --git a/src/diffusers/configuration_utils.py b/src/diffusers/configuration_utils.py
index 1c4ee33acbfd..2545c9db35b2 100644
--- a/src/diffusers/configuration_utils.py
+++ b/src/diffusers/configuration_utils.py
@@ -24,7 +24,7 @@
import re
from collections import OrderedDict
from pathlib import Path
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import numpy as np
from huggingface_hub import DDUFEntry, create_repo, hf_hub_download
@@ -94,10 +94,10 @@ class ConfigMixin:
Class attributes:
- **config_name** (`str`) -- A filename under which the config should stored when calling
[`~ConfigMixin.save_config`] (should be overridden by parent class).
- - **ignore_for_config** (`List[str]`) -- A list of attributes that should not be saved in the config (should be
+ - **ignore_for_config** (`list[str]`) -- A list of attributes that should not be saved in the config (should be
overridden by subclass).
- **has_compatibles** (`bool`) -- Whether the class has compatible classes (should be overridden by subclass).
- - **_deprecated_kwargs** (`List[str]`) -- Keyword arguments that are deprecated. Note that the `init` function
+ - **_deprecated_kwargs** (`list[str]`) -- Keyword arguments that are deprecated. Note that the `init` function
should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by
subclass).
"""
@@ -143,7 +143,7 @@ def __getattr__(self, name: str) -> Any:
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
- def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
+ def save_config(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
"""
Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the
[`~ConfigMixin.from_config`] class method.
@@ -155,7 +155,7 @@ def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool
Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
namespace).
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
"""
if os.path.isfile(save_directory):
@@ -189,13 +189,13 @@ def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool
@classmethod
def from_config(
- cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs
- ) -> Union[Self, Tuple[Self, Dict[str, Any]]]:
+ cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs
+ ) -> Self | tuple[Self, dict[str, Any]]:
r"""
Instantiate a Python class from a config dictionary.
Parameters:
- config (`Dict[str, Any]`):
+ config (`dict[str, Any]`):
A config dictionary from which the Python class is instantiated. Make sure to only load configuration
files of compatible classes.
return_unused_kwargs (`bool`, *optional*, defaults to `False`):
@@ -292,11 +292,11 @@ def get_config_dict(cls, *args, **kwargs):
@validate_hf_hub_args
def load_config(
cls,
- pretrained_model_name_or_path: Union[str, os.PathLike],
+ pretrained_model_name_or_path: str | os.PathLike,
return_unused_kwargs=False,
return_commit_hash=False,
**kwargs,
- ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
r"""
Load a model or scheduler configuration.
@@ -309,13 +309,13 @@ def load_config(
- A path to a *directory* (for example `./my_model_directory`) containing model weights saved with
[`~ConfigMixin.save_config`].
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
@@ -352,7 +352,7 @@ def load_config(
_ = kwargs.pop("mirror", None)
subfolder = kwargs.pop("subfolder", None)
user_agent = kwargs.pop("user_agent", {})
- dduf_entries: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_entries", None)
+ dduf_entries: dict[str, DDUFEntry] | None = kwargs.pop("dduf_entries", None)
user_agent = {**user_agent, "file_type": "config"}
user_agent = http_user_agent(user_agent)
@@ -563,9 +563,7 @@ def extract_init_dict(cls, config_dict, **kwargs):
return init_dict, unused_kwargs, hidden_config_dict
@classmethod
- def _dict_from_json_file(
- cls, json_file: Union[str, os.PathLike], dduf_entries: Optional[Dict[str, DDUFEntry]] = None
- ):
+ def _dict_from_json_file(cls, json_file: str | os.PathLike, dduf_entries: dict[str, DDUFEntry] | None = None):
if dduf_entries:
text = dduf_entries[json_file].read_text()
else:
@@ -577,12 +575,12 @@ def __repr__(self):
return f"{self.__class__.__name__} {self.to_json_string()}"
@property
- def config(self) -> Dict[str, Any]:
+ def config(self) -> dict[str, Any]:
"""
Returns the config of the class as a frozen dictionary
Returns:
- `Dict[str, Any]`: Config of the class.
+ `dict[str, Any]`: Config of the class.
"""
return self._internal_dict
@@ -625,7 +623,7 @@ def to_json_saveable(value):
return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
- def to_json_file(self, json_file_path: Union[str, os.PathLike]):
+ def to_json_file(self, json_file_path: str | os.PathLike):
"""
Save the configuration instance's parameters to a JSON file.
@@ -637,7 +635,7 @@ def to_json_file(self, json_file_path: Union[str, os.PathLike]):
writer.write(self.to_json_string())
@classmethod
- def _get_config_file_from_dduf(cls, pretrained_model_name_or_path: str, dduf_entries: Dict[str, DDUFEntry]):
+ def _get_config_file_from_dduf(cls, pretrained_model_name_or_path: str, dduf_entries: dict[str, DDUFEntry]):
# paths inside a DDUF file must always be "/"
config_file = (
cls.config_name
@@ -756,7 +754,7 @@ class LegacyConfigMixin(ConfigMixin):
"""
@classmethod
- def from_config(cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs):
+ def from_config(cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs):
# To prevent dependency import problem.
from .models.model_loading_utils import _fetch_remapped_cls_from_config
diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py
index 6e5ac630ab08..5b3f84213a2e 100644
--- a/src/diffusers/dependency_versions_table.py
+++ b/src/diffusers/dependency_versions_table.py
@@ -8,6 +8,7 @@
"datasets": "datasets",
"filelock": "filelock",
"flax": "flax>=0.4.1",
+ "ftfy": "ftfy",
"hf-doc-builder": "hf-doc-builder>=0.3.0",
"httpx": "httpx<1.0.0",
"huggingface-hub": "huggingface-hub>=0.34.0,<2.0",
@@ -18,7 +19,6 @@
"jax": "jax>=0.4.1",
"jaxlib": "jaxlib>=0.4.1",
"Jinja2": "Jinja2",
- "k-diffusion": "k-diffusion==0.0.12",
"torchsde": "torchsde",
"note_seq": "note_seq",
"librosa": "librosa",
@@ -29,7 +29,7 @@
"pytest": "pytest",
"pytest-timeout": "pytest-timeout",
"pytest-xdist": "pytest-xdist",
- "python": "python>=3.8.0",
+ "python": "python>=3.10.0",
"ruff": "ruff==0.9.10",
"safetensors": "safetensors>=0.3.1",
"sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
diff --git a/src/diffusers/guiders/__init__.py b/src/diffusers/guiders/__init__.py
index 58ad0c211b64..b6653817dc95 100644
--- a/src/diffusers/guiders/__init__.py
+++ b/src/diffusers/guiders/__init__.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Union
from ..utils import is_torch_available, logging
diff --git a/src/diffusers/guiders/adaptive_projected_guidance.py b/src/diffusers/guiders/adaptive_projected_guidance.py
index 8ec30d02d758..b210cb3e67aa 100644
--- a/src/diffusers/guiders/adaptive_projected_guidance.py
+++ b/src/diffusers/guiders/adaptive_projected_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
import torch
@@ -58,7 +60,7 @@ class AdaptiveProjectedGuidance(BaseGuidance):
def __init__(
self,
guidance_scale: float = 7.5,
- adaptive_projected_guidance_momentum: Optional[float] = None,
+ adaptive_projected_guidance_momentum: float | None = None,
adaptive_projected_guidance_rescale: float = 15.0,
eta: float = 1.0,
guidance_rescale: float = 0.0,
@@ -77,7 +79,7 @@ def __init__(
self.use_original_formulation = use_original_formulation
self.momentum_buffer = None
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
if self._step == 0:
if self.adaptive_projected_guidance_momentum is not None:
self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
@@ -89,8 +91,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
if self._step == 0:
if self.adaptive_projected_guidance_momentum is not None:
self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
@@ -101,7 +103,7 @@ def prepare_inputs_from_block_state(
data_batches.append(data_batch)
return data_batches
- def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput:
+ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
pred = None
if not self._is_apg_enabled():
@@ -204,7 +206,7 @@ def normalized_guidance(
pred_cond: torch.Tensor,
pred_uncond: torch.Tensor,
guidance_scale: float,
- momentum_buffer: Optional[MomentumBuffer] = None,
+ momentum_buffer: MomentumBuffer | None = None,
eta: float = 1.0,
norm_threshold: float = 0.0,
use_original_formulation: bool = False,
diff --git a/src/diffusers/guiders/adaptive_projected_guidance_mix.py b/src/diffusers/guiders/adaptive_projected_guidance_mix.py
index bdc97bcf6269..559e30d2aabe 100644
--- a/src/diffusers/guiders/adaptive_projected_guidance_mix.py
+++ b/src/diffusers/guiders/adaptive_projected_guidance_mix.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
import torch
@@ -88,7 +88,7 @@ def __init__(
self.use_original_formulation = use_original_formulation
self.momentum_buffer = None
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
if self._step == 0:
if self.adaptive_projected_guidance_momentum is not None:
self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
@@ -100,8 +100,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
if self._step == 0:
if self.adaptive_projected_guidance_momentum is not None:
self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
@@ -112,7 +112,7 @@ def prepare_inputs_from_block_state(
data_batches.append(data_batch)
return data_batches
- def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput:
+ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
pred = None
# no guidance
@@ -254,7 +254,7 @@ def __repr__(self) -> str:
def update_momentum_buffer(
pred_cond: torch.Tensor,
pred_uncond: torch.Tensor,
- momentum_buffer: Optional[MomentumBuffer] = None,
+ momentum_buffer: MomentumBuffer | None = None,
):
diff = pred_cond - pred_uncond
if momentum_buffer is not None:
@@ -265,7 +265,7 @@ def normalized_guidance(
pred_cond: torch.Tensor,
pred_uncond: torch.Tensor,
guidance_scale: float,
- momentum_buffer: Optional[MomentumBuffer] = None,
+ momentum_buffer: MomentumBuffer | None = None,
eta: float = 1.0,
norm_threshold: float = 0.0,
use_original_formulation: bool = False,
diff --git a/src/diffusers/guiders/auto_guidance.py b/src/diffusers/guiders/auto_guidance.py
index b7f62e2f4a6e..d6b6d3c492f0 100644
--- a/src/diffusers/guiders/auto_guidance.py
+++ b/src/diffusers/guiders/auto_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any
import torch
@@ -36,10 +38,10 @@ class AutoGuidance(BaseGuidance):
The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
deterioration of image quality.
- auto_guidance_layers (`int` or `List[int]`, *optional*):
+ auto_guidance_layers (`int` or `list[int]`, *optional*):
The layer indices to apply skip layer guidance to. Can be a single integer or a list of integers. If not
provided, `skip_layer_config` must be provided.
- auto_guidance_config (`LayerSkipConfig` or `List[LayerSkipConfig]`, *optional*):
+ auto_guidance_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*):
The configuration for the skip layer guidance. Can be a single `LayerSkipConfig` or a list of
`LayerSkipConfig`. If not provided, `skip_layer_guidance_layers` must be provided.
dropout (`float`, *optional*):
@@ -65,9 +67,9 @@ class AutoGuidance(BaseGuidance):
def __init__(
self,
guidance_scale: float = 7.5,
- auto_guidance_layers: Optional[Union[int, List[int]]] = None,
- auto_guidance_config: Union[LayerSkipConfig, List[LayerSkipConfig], Dict[str, Any]] = None,
- dropout: Optional[float] = None,
+ auto_guidance_layers: int | list[int] | None = None,
+ auto_guidance_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None,
+ dropout: float | None = None,
guidance_rescale: float = 0.0,
use_original_formulation: bool = False,
start: float = 0.0,
@@ -133,7 +135,7 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None:
registry = HookRegistry.check_if_exists_or_initialize(denoiser)
registry.remove_hook(name, recurse=True)
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -142,8 +144,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -151,7 +153,7 @@ def prepare_inputs_from_block_state(
data_batches.append(data_batch)
return data_batches
- def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput:
+ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
pred = None
if not self._is_ag_enabled():
diff --git a/src/diffusers/guiders/classifier_free_guidance.py b/src/diffusers/guiders/classifier_free_guidance.py
index 5e55d4d869c1..a2180a626bfc 100644
--- a/src/diffusers/guiders/classifier_free_guidance.py
+++ b/src/diffusers/guiders/classifier_free_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
import torch
@@ -91,7 +93,7 @@ def __init__(
self.guidance_rescale = guidance_rescale
self.use_original_formulation = use_original_formulation
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -100,8 +102,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -109,7 +111,7 @@ def prepare_inputs_from_block_state(
data_batches.append(data_batch)
return data_batches
- def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput:
+ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
pred = None
if not self._is_cfg_enabled():
diff --git a/src/diffusers/guiders/classifier_free_zero_star_guidance.py b/src/diffusers/guiders/classifier_free_zero_star_guidance.py
index 23b492e51b02..dd71c7537cac 100644
--- a/src/diffusers/guiders/classifier_free_zero_star_guidance.py
+++ b/src/diffusers/guiders/classifier_free_zero_star_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
import torch
@@ -77,7 +79,7 @@ def __init__(
self.guidance_rescale = guidance_rescale
self.use_original_formulation = use_original_formulation
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -86,8 +88,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -95,7 +97,7 @@ def prepare_inputs_from_block_state(
data_batches.append(data_batch)
return data_batches
- def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput:
+ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
pred = None
# YiYi Notes: add default behavior for self._enabled == False
diff --git a/src/diffusers/guiders/frequency_decoupled_guidance.py b/src/diffusers/guiders/frequency_decoupled_guidance.py
index 4ec6e2d36da9..b92ddf2c03f9 100644
--- a/src/diffusers/guiders/frequency_decoupled_guidance.py
+++ b/src/diffusers/guiders/frequency_decoupled_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
import torch
@@ -37,7 +39,7 @@
build_laplacian_pyramid_func = None
-def project(v0: torch.Tensor, v1: torch.Tensor, upcast_to_double: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
+def project(v0: torch.Tensor, v1: torch.Tensor, upcast_to_double: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
"""
Project vector v0 onto vector v1, returning the parallel and orthogonal components of v0. Implementation from paper
(Algorithm 2).
@@ -58,7 +60,7 @@ def project(v0: torch.Tensor, v1: torch.Tensor, upcast_to_double: bool = True) -
return v0_parallel, v0_orthogonal
-def build_image_from_pyramid(pyramid: List[torch.Tensor]) -> torch.Tensor:
+def build_image_from_pyramid(pyramid: list[torch.Tensor]) -> torch.Tensor:
"""
Recovers the data space latents from the Laplacian pyramid frequency space. Implementation from the paper
(Algorithm 2).
@@ -99,19 +101,19 @@ class FrequencyDecoupledGuidance(BaseGuidance):
paper. By default, we use the diffusers-native implementation that has been in the codebase for a long time.
Args:
- guidance_scales (`List[float]`, defaults to `[10.0, 5.0]`):
+ guidance_scales (`list[float]`, defaults to `[10.0, 5.0]`):
The scale parameter for frequency-decoupled guidance for each frequency component, listed from highest
frequency level to lowest. Higher values result in stronger conditioning on the text prompt, while lower
values allow for more freedom in generation. Higher values may lead to saturation and deterioration of
image quality. The FDG authors recommend using higher guidance scales for higher frequency components and
lower guidance scales for lower frequency components (so `guidance_scales` should typically be sorted in
descending order).
- guidance_rescale (`float` or `List[float]`, defaults to `0.0`):
+ guidance_rescale (`float` or `list[float]`, defaults to `0.0`):
The rescale factor applied to the noise predictions. This is used to improve image quality and fix
overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). If a list is supplied, it should be the same length as
`guidance_scales`.
- parallel_weights (`float` or `List[float]`, *optional*):
+ parallel_weights (`float` or `list[float]`, *optional*):
Optional weights for the parallel component of each frequency component of the projected CFG shift. If not
set, the weights will default to `1.0` for all components, which corresponds to using the normal CFG shift
(that is, equal weights for the parallel and orthogonal components). If set, a value in `[0, 1]` is
@@ -120,10 +122,10 @@ class FrequencyDecoupledGuidance(BaseGuidance):
Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
we use the diffusers-native implementation that has been in the codebase for a long time. See
[~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
- start (`float` or `List[float]`, defaults to `0.0`):
+ start (`float` or `list[float]`, defaults to `0.0`):
The fraction of the total number of denoising steps after which guidance starts. If a list is supplied, it
should be the same length as `guidance_scales`.
- stop (`float` or `List[float]`, defaults to `1.0`):
+ stop (`float` or `list[float]`, defaults to `1.0`):
The fraction of the total number of denoising steps after which guidance stops. If a list is supplied, it
should be the same length as `guidance_scales`.
guidance_rescale_space (`str`, defaults to `"data"`):
@@ -141,12 +143,12 @@ class FrequencyDecoupledGuidance(BaseGuidance):
@register_to_config
def __init__(
self,
- guidance_scales: Union[List[float], Tuple[float]] = [10.0, 5.0],
- guidance_rescale: Union[float, List[float], Tuple[float]] = 0.0,
- parallel_weights: Optional[Union[float, List[float], Tuple[float]]] = None,
+ guidance_scales: list[float] | tuple[float] = [10.0, 5.0],
+ guidance_rescale: float | list[float] | tuple[float] = 0.0,
+ parallel_weights: float | list[float] | tuple[float] | None = None,
use_original_formulation: bool = False,
- start: Union[float, List[float], Tuple[float]] = 0.0,
- stop: Union[float, List[float], Tuple[float]] = 1.0,
+ start: float | list[float] | tuple[float] = 0.0,
+ stop: float | list[float] | tuple[float] = 1.0,
guidance_rescale_space: str = "data",
upcast_to_double: bool = True,
enabled: bool = True,
@@ -218,7 +220,7 @@ def __init__(
f"({len(self.guidance_scales)})"
)
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -227,8 +229,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -236,7 +238,7 @@ def prepare_inputs_from_block_state(
data_batches.append(data_batch)
return data_batches
- def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput:
+ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
pred = None
if not self._is_fdg_enabled():
diff --git a/src/diffusers/guiders/guider_utils.py b/src/diffusers/guiders/guider_utils.py
index 6c328328fc3b..7be68424c345 100644
--- a/src/diffusers/guiders/guider_utils.py
+++ b/src/diffusers/guiders/guider_utils.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import os
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any
import torch
from huggingface_hub.utils import validate_hf_hub_args
@@ -51,7 +53,7 @@ def __init__(self, start: float = 0.0, stop: float = 1.0, enabled: bool = True):
self._num_inference_steps: int = None
self._timestep: torch.LongTensor = None
self._count_prepared = 0
- self._input_fields: Dict[str, Union[str, Tuple[str, str]]] = None
+ self._input_fields: dict[str, str | tuple[str, str]] = None
self._enabled = enabled
if not (0.0 <= start < 1.0):
@@ -101,11 +103,11 @@ def set_state(self, step: int, num_inference_steps: int, timestep: torch.LongTen
self._timestep = timestep
self._count_prepared = 0
- def get_state(self) -> Dict[str, Any]:
+ def get_state(self) -> dict[str, Any]:
"""
Returns the current state of the guidance technique as a dictionary. The state variables will be included in
the __repr__ method. Returns:
- `Dict[str, Any]`: A dictionary containing the current state variables including:
+ `dict[str, Any]`: A dictionary containing the current state variables including:
- step: Current inference step
- num_inference_steps: Total number of inference steps
- timestep: Current timestep tensor
@@ -163,15 +165,15 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None:
"""
pass
- def prepare_inputs(self, data: "BlockState") -> List["BlockState"]:
+ def prepare_inputs(self, data: "BlockState") -> list["BlockState"]:
raise NotImplementedError("BaseGuidance::prepare_inputs must be implemented in subclasses.")
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
raise NotImplementedError("BaseGuidance::prepare_inputs_from_block_state must be implemented in subclasses.")
- def __call__(self, data: List["BlockState"]) -> Any:
+ def __call__(self, data: list["BlockState"]) -> Any:
if not all(hasattr(d, "noise_pred") for d in data):
raise ValueError("Expected all data to have `noise_pred` attribute.")
if len(data) != self.num_conditions:
@@ -199,7 +201,7 @@ def num_conditions(self) -> int:
@classmethod
def _prepare_batch(
cls,
- data: Dict[str, Tuple[torch.Tensor, torch.Tensor]],
+ data: dict[str, tuple[torch.Tensor, torch.Tensor]],
tuple_index: int,
identifier: str,
) -> "BlockState":
@@ -208,7 +210,7 @@ def _prepare_batch(
`BaseGuidance` class. It prepares the batch based on the provided tuple index.
Args:
- input_fields (`Dict[str, Union[str, Tuple[str, str]]]`):
+ input_fields (`dict[str, str | tuple[str, str]]`):
A dictionary where the keys are the names of the fields that will be used to store the data once it is
prepared with `prepare_inputs`. The values can be either a string or a tuple of length 2, which is used
to look up the required data provided for preparation. If a string is provided, it will be used as the
@@ -242,7 +244,7 @@ def _prepare_batch(
@classmethod
def _prepare_batch_from_block_state(
cls,
- input_fields: Dict[str, Union[str, Tuple[str, str]]],
+ input_fields: dict[str, str | tuple[str, str]],
data: "BlockState",
tuple_index: int,
identifier: str,
@@ -252,7 +254,7 @@ def _prepare_batch_from_block_state(
`BaseGuidance` class. It prepares the batch based on the provided tuple index.
Args:
- input_fields (`Dict[str, Union[str, Tuple[str, str]]]`):
+ input_fields (`dict[str, str | tuple[str, str]]`):
A dictionary where the keys are the names of the fields that will be used to store the data once it is
prepared with `prepare_inputs`. The values can be either a string or a tuple of length 2, which is used
to look up the required data provided for preparation. If a string is provided, it will be used as the
@@ -288,8 +290,8 @@ def _prepare_batch_from_block_state(
@validate_hf_hub_args
def from_pretrained(
cls,
- pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
- subfolder: Optional[str] = None,
+ pretrained_model_name_or_path: str | os.PathLike | None = None,
+ subfolder: str | None = None,
return_unused_kwargs=False,
**kwargs,
) -> Self:
@@ -308,14 +310,14 @@ def from_pretrained(
The subfolder location of a model file within a larger model repository on the Hub or locally.
return_unused_kwargs (`bool`, *optional*, defaults to `False`):
Whether kwargs that are not consumed by the Python class should be returned or not.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
@@ -345,7 +347,7 @@ def from_pretrained(
)
return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs)
- def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
+ def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
"""
Save a guider configuration object to a directory so that it can be reloaded using the
[`~BaseGuidance.from_pretrained`] class method.
@@ -357,7 +359,7 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub:
Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
namespace).
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
"""
self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
@@ -365,8 +367,8 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub:
class GuiderOutput(BaseOutput):
pred: torch.Tensor
- pred_cond: Optional[torch.Tensor]
- pred_uncond: Optional[torch.Tensor]
+ pred_cond: torch.Tensor | None
+ pred_uncond: torch.Tensor | None
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
diff --git a/src/diffusers/guiders/magnitude_aware_guidance.py b/src/diffusers/guiders/magnitude_aware_guidance.py
index b81cf0d3a1f9..e83545fd889a 100644
--- a/src/diffusers/guiders/magnitude_aware_guidance.py
+++ b/src/diffusers/guiders/magnitude_aware_guidance.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
import torch
@@ -71,7 +71,7 @@ def __init__(
self.guidance_rescale = guidance_rescale
self.use_original_formulation = use_original_formulation
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -80,8 +80,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -89,7 +89,7 @@ def prepare_inputs_from_block_state(
data_batches.append(data_batch)
return data_batches
- def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput:
+ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
pred = None
if not self._is_mambo_g_enabled():
diff --git a/src/diffusers/guiders/perturbed_attention_guidance.py b/src/diffusers/guiders/perturbed_attention_guidance.py
index f233e90ca410..904d319ec3bb 100644
--- a/src/diffusers/guiders/perturbed_attention_guidance.py
+++ b/src/diffusers/guiders/perturbed_attention_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any
import torch
@@ -58,10 +60,10 @@ class PerturbedAttentionGuidance(BaseGuidance):
The fraction of the total number of denoising steps after which perturbed attention guidance starts.
perturbed_guidance_stop (`float`, defaults to `0.2`):
The fraction of the total number of denoising steps after which perturbed attention guidance stops.
- perturbed_guidance_layers (`int` or `List[int]`, *optional*):
+ perturbed_guidance_layers (`int` or `list[int]`, *optional*):
The layer indices to apply perturbed attention guidance to. Can be a single integer or a list of integers.
If not provided, `perturbed_guidance_config` must be provided.
- perturbed_guidance_config (`LayerSkipConfig` or `List[LayerSkipConfig]`, *optional*):
+ perturbed_guidance_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*):
The configuration for the perturbed attention guidance. Can be a single `LayerSkipConfig` or a list of
`LayerSkipConfig`. If not provided, `perturbed_guidance_layers` must be provided.
guidance_rescale (`float`, defaults to `0.0`):
@@ -92,8 +94,8 @@ def __init__(
perturbed_guidance_scale: float = 2.8,
perturbed_guidance_start: float = 0.01,
perturbed_guidance_stop: float = 0.2,
- perturbed_guidance_layers: Optional[Union[int, List[int]]] = None,
- perturbed_guidance_config: Union[LayerSkipConfig, List[LayerSkipConfig], Dict[str, Any]] = None,
+ perturbed_guidance_layers: int | list[int] | None = None,
+ perturbed_guidance_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None,
guidance_rescale: float = 0.0,
use_original_formulation: bool = False,
start: float = 0.0,
@@ -169,7 +171,7 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None:
registry.remove_hook(hook_name, recurse=True)
# Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.prepare_inputs
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
if self.num_conditions == 1:
tuple_indices = [0]
input_predictions = ["pred_cond"]
@@ -188,8 +190,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
if self.num_conditions == 1:
tuple_indices = [0]
input_predictions = ["pred_cond"]
@@ -211,8 +213,8 @@ def prepare_inputs_from_block_state(
def forward(
self,
pred_cond: torch.Tensor,
- pred_uncond: Optional[torch.Tensor] = None,
- pred_cond_skip: Optional[torch.Tensor] = None,
+ pred_uncond: torch.Tensor | None = None,
+ pred_cond_skip: torch.Tensor | None = None,
) -> GuiderOutput:
pred = None
diff --git a/src/diffusers/guiders/skip_layer_guidance.py b/src/diffusers/guiders/skip_layer_guidance.py
index e6109300d99c..cb7e85e179d2 100644
--- a/src/diffusers/guiders/skip_layer_guidance.py
+++ b/src/diffusers/guiders/skip_layer_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any
import torch
@@ -64,11 +66,11 @@ class SkipLayerGuidance(BaseGuidance):
The fraction of the total number of denoising steps after which skip layer guidance starts.
skip_layer_guidance_stop (`float`, defaults to `0.2`):
The fraction of the total number of denoising steps after which skip layer guidance stops.
- skip_layer_guidance_layers (`int` or `List[int]`, *optional*):
+ skip_layer_guidance_layers (`int` or `list[int]`, *optional*):
The layer indices to apply skip layer guidance to. Can be a single integer or a list of integers. If not
provided, `skip_layer_config` must be provided. The recommended values are `[7, 8, 9]` for Stable Diffusion
3.5 Medium.
- skip_layer_config (`LayerSkipConfig` or `List[LayerSkipConfig]`, *optional*):
+ skip_layer_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*):
The configuration for the skip layer guidance. Can be a single `LayerSkipConfig` or a list of
`LayerSkipConfig`. If not provided, `skip_layer_guidance_layers` must be provided.
guidance_rescale (`float`, defaults to `0.0`):
@@ -94,8 +96,8 @@ def __init__(
skip_layer_guidance_scale: float = 2.8,
skip_layer_guidance_start: float = 0.01,
skip_layer_guidance_stop: float = 0.2,
- skip_layer_guidance_layers: Optional[Union[int, List[int]]] = None,
- skip_layer_config: Union[LayerSkipConfig, List[LayerSkipConfig], Dict[str, Any]] = None,
+ skip_layer_guidance_layers: int | list[int] | None = None,
+ skip_layer_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None,
guidance_rescale: float = 0.0,
use_original_formulation: bool = False,
start: float = 0.0,
@@ -165,7 +167,7 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None:
for hook_name in self._skip_layer_hook_names:
registry.remove_hook(hook_name, recurse=True)
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
if self.num_conditions == 1:
tuple_indices = [0]
input_predictions = ["pred_cond"]
@@ -184,8 +186,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
if self.num_conditions == 1:
tuple_indices = [0]
input_predictions = ["pred_cond"]
@@ -206,8 +208,8 @@ def prepare_inputs_from_block_state(
def forward(
self,
pred_cond: torch.Tensor,
- pred_uncond: Optional[torch.Tensor] = None,
- pred_cond_skip: Optional[torch.Tensor] = None,
+ pred_uncond: torch.Tensor | None = None,
+ pred_cond_skip: torch.Tensor | None = None,
) -> GuiderOutput:
pred = None
diff --git a/src/diffusers/guiders/smoothed_energy_guidance.py b/src/diffusers/guiders/smoothed_energy_guidance.py
index 6c3906e820e0..4767607421de 100644
--- a/src/diffusers/guiders/smoothed_energy_guidance.py
+++ b/src/diffusers/guiders/smoothed_energy_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
import torch
@@ -54,11 +56,11 @@ class SmoothedEnergyGuidance(BaseGuidance):
The fraction of the total number of denoising steps after which smoothed energy guidance starts.
seg_guidance_stop (`float`, defaults to `1.0`):
The fraction of the total number of denoising steps after which smoothed energy guidance stops.
- seg_guidance_layers (`int` or `List[int]`, *optional*):
+ seg_guidance_layers (`int` or `list[int]`, *optional*):
The layer indices to apply smoothed energy guidance to. Can be a single integer or a list of integers. If
not provided, `seg_guidance_config` must be provided. The recommended values are `[7, 8, 9]` for Stable
Diffusion 3.5 Medium.
- seg_guidance_config (`SmoothedEnergyGuidanceConfig` or `List[SmoothedEnergyGuidanceConfig]`, *optional*):
+ seg_guidance_config (`SmoothedEnergyGuidanceConfig` or `list[SmoothedEnergyGuidanceConfig]`, *optional*):
The configuration for the smoothed energy layer guidance. Can be a single `SmoothedEnergyGuidanceConfig` or
a list of `SmoothedEnergyGuidanceConfig`. If not provided, `seg_guidance_layers` must be provided.
guidance_rescale (`float`, defaults to `0.0`):
@@ -86,8 +88,8 @@ def __init__(
seg_blur_threshold_inf: float = 9999.0,
seg_guidance_start: float = 0.0,
seg_guidance_stop: float = 1.0,
- seg_guidance_layers: Optional[Union[int, List[int]]] = None,
- seg_guidance_config: Union[SmoothedEnergyGuidanceConfig, List[SmoothedEnergyGuidanceConfig]] = None,
+ seg_guidance_layers: int | list[int] | None = None,
+ seg_guidance_config: SmoothedEnergyGuidanceConfig | list[SmoothedEnergyGuidanceConfig] = None,
guidance_rescale: float = 0.0,
use_original_formulation: bool = False,
start: float = 0.0,
@@ -154,7 +156,7 @@ def cleanup_models(self, denoiser: torch.nn.Module):
for hook_name in self._seg_layer_hook_names:
registry.remove_hook(hook_name, recurse=True)
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
if self.num_conditions == 1:
tuple_indices = [0]
input_predictions = ["pred_cond"]
@@ -173,8 +175,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
if self.num_conditions == 1:
tuple_indices = [0]
input_predictions = ["pred_cond"]
@@ -195,8 +197,8 @@ def prepare_inputs_from_block_state(
def forward(
self,
pred_cond: torch.Tensor,
- pred_uncond: Optional[torch.Tensor] = None,
- pred_cond_seg: Optional[torch.Tensor] = None,
+ pred_uncond: torch.Tensor | None = None,
+ pred_cond_seg: torch.Tensor | None = None,
) -> GuiderOutput:
pred = None
diff --git a/src/diffusers/guiders/tangential_classifier_free_guidance.py b/src/diffusers/guiders/tangential_classifier_free_guidance.py
index 76899c6e8494..c8911f4a69d9 100644
--- a/src/diffusers/guiders/tangential_classifier_free_guidance.py
+++ b/src/diffusers/guiders/tangential_classifier_free_guidance.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
import torch
@@ -66,7 +68,7 @@ def __init__(
self.guidance_rescale = guidance_rescale
self.use_original_formulation = use_original_formulation
- def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
+ def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -75,8 +77,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) ->
return data_batches
def prepare_inputs_from_block_state(
- self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]]
- ) -> List["BlockState"]:
+ self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
+ ) -> list["BlockState"]:
tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
data_batches = []
for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
@@ -84,7 +86,7 @@ def prepare_inputs_from_block_state(
data_batches.append(data_batch)
return data_batches
- def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput:
+ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
pred = None
if not self._is_tcfg_enabled():
diff --git a/src/diffusers/hooks/_common.py b/src/diffusers/hooks/_common.py
index f5dd1f8c7c4d..fa7ab770da6d 100644
--- a/src/diffusers/hooks/_common.py
+++ b/src/diffusers/hooks/_common.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional
-
import torch
from ..models.attention import AttentionModuleMixin, FeedForward, LuminaFeedForward
@@ -50,12 +48,13 @@
torch.nn.ConvTranspose2d,
torch.nn.ConvTranspose3d,
torch.nn.Linear,
+ torch.nn.Embedding,
# TODO(aryan): look into torch.nn.LayerNorm, torch.nn.GroupNorm later, seems to be causing some issues with CogVideoX
# because of double invocation of the same norm layer in CogVideoXLayerNorm
)
-def _get_submodule_from_fqn(module: torch.nn.Module, fqn: str) -> Optional[torch.nn.Module]:
+def _get_submodule_from_fqn(module: torch.nn.Module, fqn: str) -> torch.nn.Module | None:
for submodule_name, submodule in module.named_modules():
if submodule_name == fqn:
return submodule
diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py
index 1cbc3a35d5b9..0267dd481a89 100644
--- a/src/diffusers/hooks/_helpers.py
+++ b/src/diffusers/hooks/_helpers.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Callable, Dict, Type
+from typing import Any, Callable, Type
@dataclass
@@ -29,7 +29,7 @@ class TransformerBlockMetadata:
hidden_states_argument_name: str = "hidden_states"
_cls: Type = None
- _cached_parameter_indices: Dict[str, int] = None
+ _cached_parameter_indices: dict[str, int] = None
def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None):
kwargs = kwargs or {}
diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py
index 6f245d0befab..6130be2b8290 100644
--- a/src/diffusers/hooks/context_parallel.py
+++ b/src/diffusers/hooks/context_parallel.py
@@ -15,7 +15,7 @@
import functools
import inspect
from dataclasses import dataclass
-from typing import Dict, List, Tuple, Type, Union
+from typing import Type
import torch
import torch.distributed as dist
@@ -45,7 +45,7 @@
# TODO(aryan): consolidate with ._helpers.TransformerBlockMetadata
@dataclass
class ModuleForwardMetadata:
- cached_parameter_indices: Dict[str, int] = None
+ cached_parameter_indices: dict[str, int] = None
_cls: Type = None
def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None):
@@ -81,7 +81,7 @@ def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None)
def apply_context_parallel(
module: torch.nn.Module,
parallel_config: ContextParallelConfig,
- plan: Dict[str, ContextParallelModelPlan],
+ plan: dict[str, ContextParallelModelPlan],
) -> None:
"""Apply context parallel on a model."""
logger.debug(f"Applying context parallel with CP mesh: {parallel_config._mesh} and plan: {plan}")
@@ -110,7 +110,7 @@ def apply_context_parallel(
registry.register_hook(hook, hook_name)
-def remove_context_parallel(module: torch.nn.Module, plan: Dict[str, ContextParallelModelPlan]) -> None:
+def remove_context_parallel(module: torch.nn.Module, plan: dict[str, ContextParallelModelPlan]) -> None:
for module_id, cp_model_plan in plan.items():
submodule = _get_submodule_by_name(module, module_id)
if not isinstance(submodule, list):
@@ -328,7 +328,7 @@ def unshard_anything(
@functools.lru_cache(maxsize=64)
-def _fill_gather_shapes(shape: Tuple[int], gather_dims: Tuple[int], dim: int, world_size: int) -> List[List[int]]:
+def _fill_gather_shapes(shape: tuple[int], gather_dims: tuple[int], dim: int, world_size: int) -> list[list[int]]:
gather_shapes = []
for i in range(world_size):
rank_shape = list(copy.deepcopy(shape))
@@ -355,13 +355,13 @@ def _all_gather_anything(tensor: torch.Tensor, dim: int, group: dist.device_mesh
return gathered_tensor
-def _get_submodule_by_name(model: torch.nn.Module, name: str) -> Union[torch.nn.Module, List[torch.nn.Module]]:
+def _get_submodule_by_name(model: torch.nn.Module, name: str) -> torch.nn.Module | list[torch.nn.Module]:
if name.count("*") > 1:
raise ValueError("Wildcard '*' can only be used once in the name")
return _find_submodule_by_name(model, name)
-def _find_submodule_by_name(model: torch.nn.Module, name: str) -> Union[torch.nn.Module, List[torch.nn.Module]]:
+def _find_submodule_by_name(model: torch.nn.Module, name: str) -> torch.nn.Module | list[torch.nn.Module]:
if name == "":
return model
first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "")
diff --git a/src/diffusers/hooks/faster_cache.py b/src/diffusers/hooks/faster_cache.py
index a01afeffdb95..682cebe36c7d 100644
--- a/src/diffusers/hooks/faster_cache.py
+++ b/src/diffusers/hooks/faster_cache.py
@@ -14,7 +14,7 @@
import re
from dataclasses import dataclass
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any, Callable
import torch
@@ -60,7 +60,7 @@ class FasterCacheConfig:
Calculate the attention states every `N` iterations. If this is set to `N`, the attention computation will
be skipped `N - 1` times (i.e., cached attention states will be reused) before computing the new attention
states again.
- spatial_attention_timestep_skip_range (`Tuple[float, float]`, defaults to `(-1, 681)`):
+ spatial_attention_timestep_skip_range (`tuple[float, float]`, defaults to `(-1, 681)`):
The timestep range within which the spatial attention computation can be skipped without a significant loss
in quality. This is to be determined by the user based on the underlying model. The first value in the
tuple is the lower bound and the second value is the upper bound. Typically, diffusion timesteps for
@@ -68,17 +68,17 @@ class FasterCacheConfig:
timestep 0). For the default values, this would mean that the spatial attention computation skipping will
be applicable only after denoising timestep 681 is reached, and continue until the end of the denoising
process.
- temporal_attention_timestep_skip_range (`Tuple[float, float]`, *optional*, defaults to `None`):
+ temporal_attention_timestep_skip_range (`tuple[float, float]`, *optional*, defaults to `None`):
The timestep range within which the temporal attention computation can be skipped without a significant
loss in quality. This is to be determined by the user based on the underlying model. The first value in the
tuple is the lower bound and the second value is the upper bound. Typically, diffusion timesteps for
denoising are in the reversed range of 0 to 1000 (i.e. denoising starts at timestep 1000 and ends at
timestep 0).
- low_frequency_weight_update_timestep_range (`Tuple[int, int]`, defaults to `(99, 901)`):
+ low_frequency_weight_update_timestep_range (`tuple[int, int]`, defaults to `(99, 901)`):
The timestep range within which the low frequency weight scaling update is applied. The first value in the
tuple is the lower bound and the second value is the upper bound of the timestep range. The callback
function for the update is called only within this range.
- high_frequency_weight_update_timestep_range (`Tuple[int, int]`, defaults to `(-1, 301)`):
+ high_frequency_weight_update_timestep_range (`tuple[int, int]`, defaults to `(-1, 301)`):
The timestep range within which the high frequency weight scaling update is applied. The first value in the
tuple is the lower bound and the second value is the upper bound of the timestep range. The callback
function for the update is called only within this range.
@@ -92,15 +92,15 @@ class FasterCacheConfig:
Process the unconditional branch every `N` iterations. If this is set to `N`, the unconditional branch
computation will be skipped `N - 1` times (i.e., cached unconditional branch states will be reused) before
computing the new unconditional branch states again.
- unconditional_batch_timestep_skip_range (`Tuple[float, float]`, defaults to `(-1, 641)`):
+ unconditional_batch_timestep_skip_range (`tuple[float, float]`, defaults to `(-1, 641)`):
The timestep range within which the unconditional branch computation can be skipped without a significant
loss in quality. This is to be determined by the user based on the underlying model. The first value in the
tuple is the lower bound and the second value is the upper bound.
- spatial_attention_block_identifiers (`Tuple[str, ...]`, defaults to `("blocks.*attn1", "transformer_blocks.*attn1", "single_transformer_blocks.*attn1")`):
+ spatial_attention_block_identifiers (`tuple[str, ...]`, defaults to `("blocks.*attn1", "transformer_blocks.*attn1", "single_transformer_blocks.*attn1")`):
The identifiers to match the spatial attention blocks in the model. If the name of the block contains any
of these identifiers, FasterCache will be applied to that block. This can either be the full layer names,
partial layer names, or regex patterns. Matching will always be done using a regex match.
- temporal_attention_block_identifiers (`Tuple[str, ...]`, defaults to `("temporal_transformer_blocks.*attn1",)`):
+ temporal_attention_block_identifiers (`tuple[str, ...]`, defaults to `("temporal_transformer_blocks.*attn1",)`):
The identifiers to match the temporal attention blocks in the model. If the name of the block contains any
of these identifiers, FasterCache will be applied to that block. This can either be the full layer names,
partial layer names, or regex patterns. Matching will always be done using a regex match.
@@ -123,7 +123,7 @@ class FasterCacheConfig:
is_guidance_distilled (`bool`, defaults to `False`):
Whether the model is guidance distilled or not. If the model is guidance distilled, FasterCache will not be
applied at the denoiser-level to skip the unconditional branch computation (as there is none).
- _unconditional_conditional_input_kwargs_identifiers (`List[str]`, defaults to `("hidden_states", "encoder_hidden_states", "timestep", "attention_mask", "encoder_attention_mask")`):
+ _unconditional_conditional_input_kwargs_identifiers (`list[str]`, defaults to `("hidden_states", "encoder_hidden_states", "timestep", "attention_mask", "encoder_attention_mask")`):
The identifiers to match the input kwargs that contain the batchwise-concatenated unconditional and
conditional inputs. If the name of the input kwargs contains any of these identifiers, FasterCache will
split the inputs into unconditional and conditional branches. This must be a list of exact input kwargs
@@ -133,14 +133,14 @@ class FasterCacheConfig:
# In the paper and codebase, they hardcode these values to 2. However, it can be made configurable
# after some testing. We default to 2 if these parameters are not provided.
spatial_attention_block_skip_range: int = 2
- temporal_attention_block_skip_range: Optional[int] = None
+ temporal_attention_block_skip_range: int | None = None
- spatial_attention_timestep_skip_range: Tuple[int, int] = (-1, 681)
- temporal_attention_timestep_skip_range: Tuple[int, int] = (-1, 681)
+ spatial_attention_timestep_skip_range: tuple[int, int] = (-1, 681)
+ temporal_attention_timestep_skip_range: tuple[int, int] = (-1, 681)
# Indicator functions for low/high frequency as mentioned in Equation 11 of the paper
- low_frequency_weight_update_timestep_range: Tuple[int, int] = (99, 901)
- high_frequency_weight_update_timestep_range: Tuple[int, int] = (-1, 301)
+ low_frequency_weight_update_timestep_range: tuple[int, int] = (99, 901)
+ high_frequency_weight_update_timestep_range: tuple[int, int] = (-1, 301)
# ⍺1 and ⍺2 as mentioned in Equation 11 of the paper
alpha_low_frequency: float = 1.1
@@ -148,10 +148,10 @@ class FasterCacheConfig:
# n as described in CFG-Cache explanation in the paper - dependent on the model
unconditional_batch_skip_range: int = 5
- unconditional_batch_timestep_skip_range: Tuple[int, int] = (-1, 641)
+ unconditional_batch_timestep_skip_range: tuple[int, int] = (-1, 641)
- spatial_attention_block_identifiers: Tuple[str, ...] = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS
- temporal_attention_block_identifiers: Tuple[str, ...] = _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS
+ spatial_attention_block_identifiers: tuple[str, ...] = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS
+ temporal_attention_block_identifiers: tuple[str, ...] = _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS
attention_weight_callback: Callable[[torch.nn.Module], float] = None
low_frequency_weight_callback: Callable[[torch.nn.Module], float] = None
@@ -162,7 +162,7 @@ class FasterCacheConfig:
current_timestep_callback: Callable[[], int] = None
- _unconditional_conditional_input_kwargs_identifiers: List[str] = _UNCOND_COND_INPUT_KWARGS_IDENTIFIERS
+ _unconditional_conditional_input_kwargs_identifiers: list[str] = _UNCOND_COND_INPUT_KWARGS_IDENTIFIERS
def __repr__(self) -> str:
return (
@@ -209,7 +209,7 @@ class FasterCacheBlockState:
def __init__(self) -> None:
self.iteration: int = 0
self.batch_size: int = None
- self.cache: Tuple[torch.Tensor, torch.Tensor] = None
+ self.cache: tuple[torch.Tensor, torch.Tensor] = None
def reset(self):
self.iteration = 0
@@ -223,10 +223,10 @@ class FasterCacheDenoiserHook(ModelHook):
def __init__(
self,
unconditional_batch_skip_range: int,
- unconditional_batch_timestep_skip_range: Tuple[int, int],
+ unconditional_batch_timestep_skip_range: tuple[int, int],
tensor_format: str,
is_guidance_distilled: bool,
- uncond_cond_input_kwargs_identifiers: List[str],
+ uncond_cond_input_kwargs_identifiers: list[str],
current_timestep_callback: Callable[[], int],
low_frequency_weight_callback: Callable[[torch.nn.Module], torch.Tensor],
high_frequency_weight_callback: Callable[[torch.nn.Module], torch.Tensor],
@@ -252,7 +252,7 @@ def initialize_hook(self, module):
return module
@staticmethod
- def _get_cond_input(input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def _get_cond_input(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
# Note: this method assumes that the input tensor is batchwise-concatenated with unconditional inputs
# followed by conditional inputs.
_, cond = input.chunk(2, dim=0)
@@ -371,7 +371,7 @@ class FasterCacheBlockHook(ModelHook):
def __init__(
self,
block_skip_range: int,
- timestep_skip_range: Tuple[int, int],
+ timestep_skip_range: tuple[int, int],
is_guidance_distilled: bool,
weight_callback: Callable[[torch.nn.Module], float],
current_timestep_callback: Callable[[], int],
diff --git a/src/diffusers/hooks/first_block_cache.py b/src/diffusers/hooks/first_block_cache.py
index 862d44059301..685ccd383674 100644
--- a/src/diffusers/hooks/first_block_cache.py
+++ b/src/diffusers/hooks/first_block_cache.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Tuple, Union
import torch
@@ -53,9 +52,9 @@ class FBCSharedBlockState(BaseState):
def __init__(self) -> None:
super().__init__()
- self.head_block_output: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None
+ self.head_block_output: torch.Tensor | tuple[torch.Tensor, ...] = None
self.head_block_residual: torch.Tensor = None
- self.tail_block_residuals: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None
+ self.tail_block_residuals: torch.Tensor | tuple[torch.Tensor, ...] = None
self.should_compute: bool = True
def reset(self):
diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py
index 47f1f4199615..15c9fa44f5ff 100644
--- a/src/diffusers/hooks/group_offloading.py
+++ b/src/diffusers/hooks/group_offloading.py
@@ -17,7 +17,7 @@
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass, replace
from enum import Enum
-from typing import Dict, List, Optional, Set, Tuple, Union
+from typing import Set
import safetensors.torch
import torch
@@ -56,31 +56,31 @@ class GroupOffloadingConfig:
non_blocking: bool
record_stream: bool
low_cpu_mem_usage: bool
- num_blocks_per_group: Optional[int] = None
- offload_to_disk_path: Optional[str] = None
- stream: Optional[Union[torch.cuda.Stream, torch.Stream]] = None
- block_modules: Optional[List[str]] = None
- exclude_kwargs: Optional[List[str]] = None
- module_prefix: Optional[str] = ""
+ num_blocks_per_group: int | None = None
+ offload_to_disk_path: str | None = None
+ stream: torch.cuda.Stream | torch.Stream | None = None
+ block_modules: list[str] | None = None
+ exclude_kwargs: list[str] | None = None
+ module_prefix: str = ""
class ModuleGroup:
def __init__(
self,
- modules: List[torch.nn.Module],
+ modules: list[torch.nn.Module],
offload_device: torch.device,
onload_device: torch.device,
offload_leader: torch.nn.Module,
- onload_leader: Optional[torch.nn.Module] = None,
- parameters: Optional[List[torch.nn.Parameter]] = None,
- buffers: Optional[List[torch.Tensor]] = None,
+ onload_leader: torch.nn.Module | None = None,
+ parameters: list[torch.nn.Parameter] | None = None,
+ buffers: list[torch.Tensor] | None = None,
non_blocking: bool = False,
- stream: Union[torch.cuda.Stream, torch.Stream, None] = None,
- record_stream: Optional[bool] = False,
+ stream: torch.cuda.Stream | torch.Stream | None = None,
+ record_stream: bool | None = False,
low_cpu_mem_usage: bool = False,
onload_self: bool = True,
- offload_to_disk_path: Optional[str] = None,
- group_id: Optional[Union[int, str]] = None,
+ offload_to_disk_path: str | None = None,
+ group_id: int | str | None = None,
) -> None:
self.modules = modules
self.offload_device = offload_device
@@ -287,7 +287,7 @@ class GroupOffloadingHook(ModelHook):
def __init__(self, group: ModuleGroup, *, config: GroupOffloadingConfig) -> None:
self.group = group
- self.next_group: Optional[ModuleGroup] = None
+ self.next_group: ModuleGroup | None = None
self.config = config
def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
@@ -359,7 +359,7 @@ class LazyPrefetchGroupOffloadingHook(ModelHook):
_is_stateful = False
def __init__(self):
- self.execution_order: List[Tuple[str, torch.nn.Module]] = []
+ self.execution_order: list[tuple[str, torch.nn.Module]] = []
self._layer_execution_tracker_module_names = set()
def initialize_hook(self, module):
@@ -463,17 +463,17 @@ def pre_forward(self, module, *args, **kwargs):
def apply_group_offloading(
module: torch.nn.Module,
- onload_device: Union[str, torch.device],
- offload_device: Union[str, torch.device] = torch.device("cpu"),
- offload_type: Union[str, GroupOffloadingType] = "block_level",
- num_blocks_per_group: Optional[int] = None,
+ onload_device: str | torch.device,
+ offload_device: str | torch.device = torch.device("cpu"),
+ offload_type: str | GroupOffloadingType = "block_level",
+ num_blocks_per_group: int | None = None,
non_blocking: bool = False,
use_stream: bool = False,
record_stream: bool = False,
low_cpu_mem_usage: bool = False,
- offload_to_disk_path: Optional[str] = None,
- block_modules: Optional[List[str]] = None,
- exclude_kwargs: Optional[List[str]] = None,
+ offload_to_disk_path: str | None = None,
+ block_modules: list[str] | None = None,
+ exclude_kwargs: list[str] | None = None,
) -> None:
r"""
Applies group offloading to the internal layers of a torch.nn.Module. To understand what group offloading is, and
@@ -531,10 +531,10 @@ def apply_group_offloading(
If True, the CPU memory usage is minimized by pinning tensors on-the-fly instead of pre-pinning them. This
option only matters when using streamed CPU offloading (i.e. `use_stream=True`). This can be useful when
the CPU memory is a bottleneck but may counteract the benefits of using streams.
- block_modules (`List[str]`, *optional*):
+ block_modules (`list[str]`, *optional*):
List of module names that should be treated as blocks for offloading. If provided, only these modules will
be considered for block-level offloading. If not provided, the default block detection logic will be used.
- exclude_kwargs (`List[str]`, *optional*):
+ exclude_kwargs (`list[str]`, *optional*):
List of kwarg keys that should not be processed by send_to_device. This is useful for mutable state like
caching lists that need to maintain their object identity across forward passes. If not provided, will be
inferred from the module's `_skip_keys` attribute if it exists.
@@ -844,7 +844,7 @@ def _apply_lazy_group_offloading_hook(
def _gather_parameters_with_no_group_offloading_parent(
module: torch.nn.Module, modules_with_group_offloading: Set[str]
-) -> List[torch.nn.Parameter]:
+) -> list[torch.nn.Parameter]:
parameters = []
for name, parameter in module.named_parameters():
has_parent_with_group_offloading = False
@@ -862,7 +862,7 @@ def _gather_parameters_with_no_group_offloading_parent(
def _gather_buffers_with_no_group_offloading_parent(
module: torch.nn.Module, modules_with_group_offloading: Set[str]
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
buffers = []
for name, buffer in module.named_buffers():
has_parent_with_group_offloading = False
@@ -878,7 +878,7 @@ def _gather_buffers_with_no_group_offloading_parent(
return buffers
-def _find_parent_module_in_module_dict(name: str, module_dict: Dict[str, torch.nn.Module]) -> str:
+def _find_parent_module_in_module_dict(name: str, module_dict: dict[str, torch.nn.Module]) -> str:
atoms = name.split(".")
while len(atoms) > 0:
parent_name = ".".join(atoms)
@@ -902,7 +902,7 @@ def _raise_error_if_accelerate_model_or_sequential_hook_present(module: torch.nn
)
-def _get_top_level_group_offload_hook(module: torch.nn.Module) -> Optional[GroupOffloadingHook]:
+def _get_top_level_group_offload_hook(module: torch.nn.Module) -> GroupOffloadingHook | None:
for submodule in module.modules():
if hasattr(submodule, "_diffusers_hook"):
group_offloading_hook = submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING)
diff --git a/src/diffusers/hooks/hooks.py b/src/diffusers/hooks/hooks.py
index 6e097e5882a0..2d575b85427c 100644
--- a/src/diffusers/hooks/hooks.py
+++ b/src/diffusers/hooks/hooks.py
@@ -13,7 +13,7 @@
# limitations under the License.
import functools
-from typing import Any, Dict, Optional, Tuple
+from typing import Any
import torch
@@ -86,19 +86,19 @@ def deinitalize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
"""
return module
- def pre_forward(self, module: torch.nn.Module, *args, **kwargs) -> Tuple[Tuple[Any], Dict[str, Any]]:
+ def pre_forward(self, module: torch.nn.Module, *args, **kwargs) -> tuple[tuple[Any], dict[str, Any]]:
r"""
Hook that is executed just before the forward method of the model.
Args:
module (`torch.nn.Module`):
The module whose forward pass will be executed just after this event.
- args (`Tuple[Any]`):
+ args (`tuple[Any]`):
The positional arguments passed to the module.
- kwargs (`Dict[Str, Any]`):
+ kwargs (`dict[Str, Any]`):
The keyword arguments passed to the module.
Returns:
- `Tuple[Tuple[Any], Dict[Str, Any]]`:
+ `tuple[tuple[Any], dict[Str, Any]]`:
A tuple with the treated `args` and `kwargs`.
"""
return args, kwargs
@@ -168,7 +168,7 @@ class HookRegistry:
def __init__(self, module_ref: torch.nn.Module) -> None:
super().__init__()
- self.hooks: Dict[str, ModelHook] = {}
+ self.hooks: dict[str, ModelHook] = {}
self._module_ref = module_ref
self._hook_order = []
@@ -214,7 +214,7 @@ def new_forward(module, *args, **kwargs):
self._hook_order.append(name)
self._fn_refs.append(fn_ref)
- def get_hook(self, name: str) -> Optional[ModelHook]:
+ def get_hook(self, name: str) -> ModelHook | None:
return self.hooks.get(name, None)
def remove_hook(self, name: str, recurse: bool = True) -> None:
@@ -265,7 +265,7 @@ def check_if_exists_or_initialize(cls, module: torch.nn.Module) -> "HookRegistry
module._diffusers_hook = cls(module)
return module._diffusers_hook
- def _set_context(self, name: Optional[str] = None) -> None:
+ def _set_context(self, name: str | None = None) -> None:
for hook_name in reversed(self._hook_order):
hook = self.hooks[hook_name]
if hook._is_stateful:
diff --git a/src/diffusers/hooks/layer_skip.py b/src/diffusers/hooks/layer_skip.py
index 0ce02e987d09..112edfa2f79b 100644
--- a/src/diffusers/hooks/layer_skip.py
+++ b/src/diffusers/hooks/layer_skip.py
@@ -14,7 +14,7 @@
import math
from dataclasses import asdict, dataclass
-from typing import Callable, List, Optional
+from typing import Callable
import torch
@@ -43,7 +43,7 @@ class LayerSkipConfig:
Configuration for skipping internal transformer blocks when executing a transformer model.
Args:
- indices (`List[int]`):
+ indices (`list[int]`):
The indices of the layer to skip. This is typically the first layer in the transformer block.
fqn (`str`, defaults to `"auto"`):
The fully qualified name identifying the stack of transformer blocks. Typically, this is
@@ -63,7 +63,7 @@ class LayerSkipConfig:
skipped layers are fully retained, which is equivalent to not skipping any layers.
"""
- indices: List[int]
+ indices: list[int]
fqn: str = "auto"
skip_attention: bool = True
skip_attention_scores: bool = False
@@ -196,7 +196,7 @@ def apply_layer_skip(module: torch.nn.Module, config: LayerSkipConfig) -> None:
_apply_layer_skip_hook(module, config)
-def _apply_layer_skip_hook(module: torch.nn.Module, config: LayerSkipConfig, name: Optional[str] = None) -> None:
+def _apply_layer_skip_hook(module: torch.nn.Module, config: LayerSkipConfig, name: str | None = None) -> None:
name = name or _LAYER_SKIP_HOOK
if config.skip_attention and config.skip_attention_scores:
diff --git a/src/diffusers/hooks/layerwise_casting.py b/src/diffusers/hooks/layerwise_casting.py
index a036ad37dc2f..1edff7805fcc 100644
--- a/src/diffusers/hooks/layerwise_casting.py
+++ b/src/diffusers/hooks/layerwise_casting.py
@@ -13,7 +13,7 @@
# limitations under the License.
import re
-from typing import Optional, Tuple, Type, Union
+from typing import Type
import torch
@@ -102,8 +102,8 @@ def apply_layerwise_casting(
module: torch.nn.Module,
storage_dtype: torch.dtype,
compute_dtype: torch.dtype,
- skip_modules_pattern: Union[str, Tuple[str, ...]] = "auto",
- skip_modules_classes: Optional[Tuple[Type[torch.nn.Module], ...]] = None,
+ skip_modules_pattern: str | tuple[str, ...] = "auto",
+ skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None,
non_blocking: bool = False,
) -> None:
r"""
@@ -137,12 +137,12 @@ def apply_layerwise_casting(
The dtype to cast the module to before/after the forward pass for storage.
compute_dtype (`torch.dtype`):
The dtype to cast the module to during the forward pass for computation.
- skip_modules_pattern (`Tuple[str, ...]`, defaults to `"auto"`):
+ skip_modules_pattern (`tuple[str, ...]`, defaults to `"auto"`):
A list of patterns to match the names of the modules to skip during the layerwise casting process. If set
to `"auto"`, the default patterns are used. If set to `None`, no modules are skipped. If set to `None`
alongside `skip_modules_classes` being `None`, the layerwise casting is applied directly to the module
instead of its internal submodules.
- skip_modules_classes (`Tuple[Type[torch.nn.Module], ...]`, defaults to `None`):
+ skip_modules_classes (`tuple[Type[torch.nn.Module], ...]`, defaults to `None`):
A list of module classes to skip during the layerwise casting process.
non_blocking (`bool`, defaults to `False`):
If `True`, the weight casting operations are non-blocking.
@@ -169,8 +169,8 @@ def _apply_layerwise_casting(
module: torch.nn.Module,
storage_dtype: torch.dtype,
compute_dtype: torch.dtype,
- skip_modules_pattern: Optional[Tuple[str, ...]] = None,
- skip_modules_classes: Optional[Tuple[Type[torch.nn.Module], ...]] = None,
+ skip_modules_pattern: tuple[str, ...] | None = None,
+ skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None,
non_blocking: bool = False,
_prefix: str = "",
) -> None:
diff --git a/src/diffusers/hooks/pyramid_attention_broadcast.py b/src/diffusers/hooks/pyramid_attention_broadcast.py
index 12d6aa0616e9..ed5bd24dea01 100644
--- a/src/diffusers/hooks/pyramid_attention_broadcast.py
+++ b/src/diffusers/hooks/pyramid_attention_broadcast.py
@@ -14,7 +14,7 @@
import re
from dataclasses import dataclass
-from typing import Any, Callable, Optional, Tuple, Union
+from typing import Any, Callable
import torch
@@ -54,34 +54,34 @@ class PyramidAttentionBroadcastConfig:
The number of times a specific cross-attention broadcast is skipped before computing the attention states
to re-use. If this is set to the value `N`, the attention computation will be skipped `N - 1` times (i.e.,
old attention states will be reused) before computing the new attention states again.
- spatial_attention_timestep_skip_range (`Tuple[int, int]`, defaults to `(100, 800)`):
+ spatial_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`):
The range of timesteps to skip in the spatial attention layer. The attention computations will be
conditionally skipped if the current timestep is within the specified range.
- temporal_attention_timestep_skip_range (`Tuple[int, int]`, defaults to `(100, 800)`):
+ temporal_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`):
The range of timesteps to skip in the temporal attention layer. The attention computations will be
conditionally skipped if the current timestep is within the specified range.
- cross_attention_timestep_skip_range (`Tuple[int, int]`, defaults to `(100, 800)`):
+ cross_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`):
The range of timesteps to skip in the cross-attention layer. The attention computations will be
conditionally skipped if the current timestep is within the specified range.
- spatial_attention_block_identifiers (`Tuple[str, ...]`):
+ spatial_attention_block_identifiers (`tuple[str, ...]`):
The identifiers to match against the layer names to determine if the layer is a spatial attention layer.
- temporal_attention_block_identifiers (`Tuple[str, ...]`):
+ temporal_attention_block_identifiers (`tuple[str, ...]`):
The identifiers to match against the layer names to determine if the layer is a temporal attention layer.
- cross_attention_block_identifiers (`Tuple[str, ...]`):
+ cross_attention_block_identifiers (`tuple[str, ...]`):
The identifiers to match against the layer names to determine if the layer is a cross-attention layer.
"""
- spatial_attention_block_skip_range: Optional[int] = None
- temporal_attention_block_skip_range: Optional[int] = None
- cross_attention_block_skip_range: Optional[int] = None
+ spatial_attention_block_skip_range: int | None = None
+ temporal_attention_block_skip_range: int | None = None
+ cross_attention_block_skip_range: int | None = None
- spatial_attention_timestep_skip_range: Tuple[int, int] = (100, 800)
- temporal_attention_timestep_skip_range: Tuple[int, int] = (100, 800)
- cross_attention_timestep_skip_range: Tuple[int, int] = (100, 800)
+ spatial_attention_timestep_skip_range: tuple[int, int] = (100, 800)
+ temporal_attention_timestep_skip_range: tuple[int, int] = (100, 800)
+ cross_attention_timestep_skip_range: tuple[int, int] = (100, 800)
- spatial_attention_block_identifiers: Tuple[str, ...] = _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS
- temporal_attention_block_identifiers: Tuple[str, ...] = _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS
- cross_attention_block_identifiers: Tuple[str, ...] = _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS
+ spatial_attention_block_identifiers: tuple[str, ...] = _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS
+ temporal_attention_block_identifiers: tuple[str, ...] = _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS
+ cross_attention_block_identifiers: tuple[str, ...] = _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS
current_timestep_callback: Callable[[], int] = None
@@ -141,7 +141,7 @@ class PyramidAttentionBroadcastHook(ModelHook):
_is_stateful = True
def __init__(
- self, timestep_skip_range: Tuple[int, int], block_skip_range: int, current_timestep_callback: Callable[[], int]
+ self, timestep_skip_range: tuple[int, int], block_skip_range: int, current_timestep_callback: Callable[[], int]
) -> None:
super().__init__()
@@ -191,7 +191,7 @@ def apply_pyramid_attention_broadcast(module: torch.nn.Module, config: PyramidAt
Args:
module (`torch.nn.Module`):
The module to apply Pyramid Attention Broadcast to.
- config (`Optional[PyramidAttentionBroadcastConfig]`, `optional`, defaults to `None`):
+ config (`PyramidAttentionBroadcastConfig | None`, `optional`, defaults to `None`):
The configuration to use for Pyramid Attention Broadcast.
Example:
@@ -288,8 +288,8 @@ def _apply_pyramid_attention_broadcast_on_attention_class(
def _apply_pyramid_attention_broadcast_hook(
- module: Union[Attention, MochiAttention],
- timestep_skip_range: Tuple[int, int],
+ module: Attention | MochiAttention,
+ timestep_skip_range: tuple[int, int],
block_skip_range: int,
current_timestep_callback: Callable[[], int],
):
@@ -299,7 +299,7 @@ def _apply_pyramid_attention_broadcast_hook(
Args:
module (`torch.nn.Module`):
The module to apply Pyramid Attention Broadcast to.
- timestep_skip_range (`Tuple[int, int]`):
+ timestep_skip_range (`tuple[int, int]`):
The range of timesteps to skip in the attention layer. The attention computations will be conditionally
skipped if the current timestep is within the specified range.
block_skip_range (`int`):
diff --git a/src/diffusers/hooks/smoothed_energy_guidance_utils.py b/src/diffusers/hooks/smoothed_energy_guidance_utils.py
index 622f60764762..f413b6376d7c 100644
--- a/src/diffusers/hooks/smoothed_energy_guidance_utils.py
+++ b/src/diffusers/hooks/smoothed_energy_guidance_utils.py
@@ -14,7 +14,6 @@
import math
from dataclasses import asdict, dataclass
-from typing import List, Optional
import torch
import torch.nn.functional as F
@@ -35,21 +34,21 @@ class SmoothedEnergyGuidanceConfig:
Configuration for skipping internal transformer blocks when executing a transformer model.
Args:
- indices (`List[int]`):
+ indices (`list[int]`):
The indices of the layer to skip. This is typically the first layer in the transformer block.
fqn (`str`, defaults to `"auto"`):
The fully qualified name identifying the stack of transformer blocks. Typically, this is
`transformer_blocks`, `single_transformer_blocks`, `blocks`, `layers`, or `temporal_transformer_blocks`.
For automatic detection, set this to `"auto"`. "auto" only works on DiT models. For UNet models, you must
provide the correct fqn.
- _query_proj_identifiers (`List[str]`, defaults to `None`):
+ _query_proj_identifiers (`list[str]`, defaults to `None`):
The identifiers for the query projection layers. Typically, these are `to_q`, `query`, or `q_proj`. If
`None`, `to_q` is used by default.
"""
- indices: List[int]
+ indices: list[int]
fqn: str = "auto"
- _query_proj_identifiers: List[str] = None
+ _query_proj_identifiers: list[str] = None
def to_dict(self):
return asdict(self)
@@ -73,7 +72,7 @@ def post_forward(self, module: torch.nn.Module, output: torch.Tensor) -> torch.T
def _apply_smoothed_energy_guidance_hook(
- module: torch.nn.Module, config: SmoothedEnergyGuidanceConfig, blur_sigma: float, name: Optional[str] = None
+ module: torch.nn.Module, config: SmoothedEnergyGuidanceConfig, blur_sigma: float, name: str | None = None
) -> None:
name = name or _SMOOTHED_ENERGY_GUIDANCE_HOOK
diff --git a/src/diffusers/hooks/taylorseer_cache.py b/src/diffusers/hooks/taylorseer_cache.py
index 7cad9f4fa161..303155105e71 100644
--- a/src/diffusers/hooks/taylorseer_cache.py
+++ b/src/diffusers/hooks/taylorseer_cache.py
@@ -1,7 +1,6 @@
import math
import re
from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
import torch
import torch.nn as nn
@@ -51,12 +50,12 @@ class TaylorSeerCacheConfig:
Data type used for storing and computing Taylor series factors. Lower precision reduces memory but may
affect stability; higher precision improves accuracy at the cost of more memory.
- skip_predict_identifiers (`List[str]`, *optional*, defaults to `None`):
+ skip_predict_identifiers (`list[str]`, *optional*, defaults to `None`):
Regex patterns (using `re.fullmatch`) for module names to place as "skip" in "cache" mode. In this mode,
the module computes fully during initial or refresh steps but returns a zero tensor (matching recorded
shape) during prediction steps to skip computation cheaply.
- cache_identifiers (`List[str]`, *optional*, defaults to `None`):
+ cache_identifiers (`list[str]`, *optional*, defaults to `None`):
Regex patterns (using `re.fullmatch`) for module names to place in Taylor-series caching mode, where
outputs are approximated and cached for reuse.
@@ -82,11 +81,11 @@ def forward(x):
cache_interval: int = 5
disable_cache_before_step: int = 3
- disable_cache_after_step: Optional[int] = None
+ disable_cache_after_step: int | None = None
max_order: int = 1
- taylor_factors_dtype: Optional[torch.dtype] = torch.bfloat16
- skip_predict_identifiers: Optional[List[str]] = None
- cache_identifiers: Optional[List[str]] = None
+ taylor_factors_dtype: torch.dtype | None = torch.bfloat16
+ skip_predict_identifiers: list[str] | None = None
+ cache_identifiers: list[str] | None = None
use_lite_mode: bool = False
def __repr__(self) -> str:
@@ -106,7 +105,7 @@ def __repr__(self) -> str:
class TaylorSeerState:
def __init__(
self,
- taylor_factors_dtype: Optional[torch.dtype] = torch.bfloat16,
+ taylor_factors_dtype: torch.dtype | None = torch.bfloat16,
max_order: int = 1,
is_inactive: bool = False,
):
@@ -114,11 +113,11 @@ def __init__(
self.max_order = max_order
self.is_inactive = is_inactive
- self.module_dtypes: Tuple[torch.dtype, ...] = ()
- self.last_update_step: Optional[int] = None
- self.taylor_factors: Dict[int, Dict[int, torch.Tensor]] = {}
- self.inactive_shapes: Optional[Tuple[Tuple[int, ...], ...]] = None
- self.device: Optional[torch.device] = None
+ self.module_dtypes: tuple[torch.dtype, ...] = ()
+ self.last_update_step: int | None = None
+ self.taylor_factors: dict[int, dict[int, torch.Tensor]] = {}
+ self.inactive_shapes: tuple[tuple[int, ...], ...] | None = None
+ self.device: torch.device | None = None
self.current_step: int = -1
def reset(self) -> None:
@@ -130,7 +129,7 @@ def reset(self) -> None:
def update(
self,
- outputs: Tuple[torch.Tensor, ...],
+ outputs: tuple[torch.Tensor, ...],
) -> None:
self.module_dtypes = tuple(output.dtype for output in outputs)
self.device = outputs[0].device
@@ -139,7 +138,7 @@ def update(
self.inactive_shapes = tuple(output.shape for output in outputs)
else:
for i, features in enumerate(outputs):
- new_factors: Dict[int, torch.Tensor] = {0: features}
+ new_factors: dict[int, torch.Tensor] = {0: features}
is_first_update = self.last_update_step is None
if not is_first_update:
delta_step = self.current_step - self.last_update_step
@@ -160,7 +159,7 @@ def update(
self.last_update_step = self.current_step
@torch.compiler.disable
- def predict(self) -> List[torch.Tensor]:
+ def predict(self) -> list[torch.Tensor]:
if self.last_update_step is None:
raise ValueError("Cannot predict without prior initialization/update.")
@@ -204,7 +203,7 @@ def __init__(
disable_cache_before_step: int,
taylor_factors_dtype: torch.dtype,
state_manager: StateManager,
- disable_cache_after_step: Optional[int] = None,
+ disable_cache_after_step: int | None = None,
):
super().__init__()
self.cache_interval = cache_interval
@@ -245,7 +244,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
return outputs_list[0] if len(outputs_list) == 1 else tuple(outputs_list)
-def _resolve_patterns(config: TaylorSeerCacheConfig) -> Tuple[List[str], List[str]]:
+def _resolve_patterns(config: TaylorSeerCacheConfig) -> tuple[list[str], list[str]]:
"""
Resolve effective inactive and active pattern lists from config + templates.
"""
diff --git a/src/diffusers/hooks/utils.py b/src/diffusers/hooks/utils.py
index c5260eeebe1f..da9e398fdf29 100644
--- a/src/diffusers/hooks/utils.py
+++ b/src/diffusers/hooks/utils.py
@@ -21,8 +21,8 @@ def _get_identifiable_transformer_blocks_in_module(module: torch.nn.Module):
module_list_with_transformer_blocks = []
for name, submodule in module.named_modules():
name_endswith_identifier = any(name.endswith(identifier) for identifier in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS)
- is_modulelist = isinstance(submodule, torch.nn.ModuleList)
- if name_endswith_identifier and is_modulelist:
+ is_ModuleList = isinstance(submodule, torch.nn.ModuleList)
+ if name_endswith_identifier and is_ModuleList:
module_list_with_transformer_blocks.append((name, submodule))
return module_list_with_transformer_blocks
diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py
index abd0a25819f5..57cde8d12338 100644
--- a/src/diffusers/image_processor.py
+++ b/src/diffusers/image_processor.py
@@ -14,7 +14,6 @@
import math
import warnings
-from typing import List, Optional, Tuple, Union
import numpy as np
import PIL.Image
@@ -26,14 +25,9 @@
from .utils import CONFIG_NAME, PIL_INTERPOLATION, deprecate
-PipelineImageInput = Union[
- PIL.Image.Image,
- np.ndarray,
- torch.Tensor,
- List[PIL.Image.Image],
- List[np.ndarray],
- List[torch.Tensor],
-]
+PipelineImageInput = (
+ PIL.Image.Image | np.ndarray | torch.Tensor | list[PIL.Image.Image] | list[np.ndarray] | list[torch.Tensor]
+)
PipelineDepthInput = PipelineImageInput
@@ -47,7 +41,7 @@ def is_valid_image(image) -> bool:
- A 2D or 3D `np.ndarray` or `torch.Tensor` (grayscale or color image).
Args:
- image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`):
+ image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
The image to validate. It can be a PIL image, a NumPy array, or a torch tensor.
Returns:
@@ -68,7 +62,7 @@ def is_valid_image_imagelist(images):
- A list of valid images.
Args:
- images (`Union[np.ndarray, torch.Tensor, PIL.Image.Image, List]`):
+ images (`np.ndarray | torch.Tensor | PIL.Image.Image | list`):
The image(s) to check. Can be a batch of images (4D tensor/array), a single image, or a list of valid
images.
@@ -116,7 +110,7 @@ def __init__(
vae_scale_factor: int = 8,
vae_latent_channels: int = 4,
resample: str = "lanczos",
- reducing_gap: int = None,
+ reducing_gap: int | None = None,
do_normalize: bool = True,
do_binarize: bool = False,
do_convert_rgb: bool = False,
@@ -131,7 +125,7 @@ def __init__(
)
@staticmethod
- def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
+ def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
r"""
Convert a numpy image or a batch of images to a PIL image.
@@ -140,7 +134,7 @@ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
The image array to convert to PIL format.
Returns:
- `List[PIL.Image.Image]`:
+ `list[PIL.Image.Image]`:
A list of PIL images.
"""
if images.ndim == 3:
@@ -155,12 +149,12 @@ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
return pil_images
@staticmethod
- def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:
+ def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
r"""
Convert a PIL image or a list of PIL images to NumPy arrays.
Args:
- images (`PIL.Image.Image` or `List[PIL.Image.Image]`):
+ images (`PIL.Image.Image` or `list[PIL.Image.Image]`):
The PIL image or list of images to convert to NumPy format.
Returns:
@@ -210,7 +204,7 @@ def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
return images
@staticmethod
- def normalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
+ def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
r"""
Normalize an image array to [-1,1].
@@ -225,7 +219,7 @@ def normalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torc
return 2.0 * images - 1.0
@staticmethod
- def denormalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
+ def denormalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
r"""
Denormalize an image array to [0,1].
@@ -467,11 +461,11 @@ def _resize_and_crop(
def resize(
self,
- image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
height: int,
width: int,
resize_mode: str = "default", # "default", "fill", "crop"
- ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]:
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
"""
Resize image.
@@ -544,7 +538,7 @@ def binarize(self, image: PIL.Image.Image) -> PIL.Image.Image:
return image
def _denormalize_conditionally(
- self, images: torch.Tensor, do_denormalize: Optional[List[bool]] = None
+ self, images: torch.Tensor, do_denormalize: list[bool] | None = None
) -> torch.Tensor:
r"""
Denormalize a batch of images based on a condition list.
@@ -552,7 +546,7 @@ def _denormalize_conditionally(
Args:
images (`torch.Tensor`):
The input image tensor.
- do_denormalize (`Optional[List[bool]`, *optional*, defaults to `None`):
+ do_denormalize (`Optional[list[bool]`, *optional*, defaults to `None`):
A list of booleans indicating whether to denormalize each image in the batch. If `None`, will use the
value of `do_normalize` in the `VaeImageProcessor` config.
"""
@@ -565,25 +559,25 @@ def _denormalize_conditionally(
def get_default_height_width(
self,
- image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
- height: Optional[int] = None,
- width: Optional[int] = None,
- ) -> Tuple[int, int]:
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
+ height: int | None = None,
+ width: int | None = None,
+ ) -> tuple[int, int]:
r"""
Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`.
Args:
- image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`):
+ image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it
should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch
tensor, it should have shape `[batch, channels, height, width]`.
- height (`Optional[int]`, *optional*, defaults to `None`):
+ height (`int | None`, *optional*, defaults to `None`):
The height of the preprocessed image. If `None`, the height of the `image` input will be used.
- width (`Optional[int]`, *optional*, defaults to `None`):
+ width (`int | None`, *optional*, defaults to `None`):
The width of the preprocessed image. If `None`, the width of the `image` input will be used.
Returns:
- `Tuple[int, int]`:
+ `tuple[int, int]`:
A tuple containing the height and width, both resized to the nearest integer multiple of
`vae_scale_factor`.
"""
@@ -613,10 +607,10 @@ def get_default_height_width(
def preprocess(
self,
image: PipelineImageInput,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
resize_mode: str = "default", # "default", "fill", "crop"
- crops_coords: Optional[Tuple[int, int, int, int]] = None,
+ crops_coords: tuple[int, int, int, int] | None = None,
) -> torch.Tensor:
"""
Preprocess the image input.
@@ -638,7 +632,7 @@ def preprocess(
image to fit within the specified width and height, maintaining the aspect ratio, and then center the
image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
supported for PIL image input.
- crops_coords (`List[Tuple[int, int, int, int]]`, *optional*, defaults to `None`):
+ crops_coords (`list[tuple[int, int, int, int]]`, *optional*, defaults to `None`):
The crop coordinates for each image in the batch. If `None`, will not crop the image.
Returns:
@@ -745,8 +739,8 @@ def postprocess(
self,
image: torch.Tensor,
output_type: str = "pil",
- do_denormalize: Optional[List[bool]] = None,
- ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]:
+ do_denormalize: list[bool] | None = None,
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
"""
Postprocess the image output from tensor to `output_type`.
@@ -755,7 +749,7 @@ def postprocess(
The image input, should be a pytorch tensor with shape `B x C x H x W`.
output_type (`str`, *optional*, defaults to `pil`):
The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
- do_denormalize (`List[bool]`, *optional*, defaults to `None`):
+ do_denormalize (`list[bool]`, *optional*, defaults to `None`):
Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
`VaeImageProcessor` config.
@@ -796,7 +790,7 @@ def apply_overlay(
mask: PIL.Image.Image,
init_image: PIL.Image.Image,
image: PIL.Image.Image,
- crop_coords: Optional[Tuple[int, int, int, int]] = None,
+ crop_coords: tuple[int, int, int, int] | None = None,
) -> PIL.Image.Image:
r"""
Applies an overlay of the mask and the inpainted image on the original image.
@@ -808,7 +802,7 @@ def apply_overlay(
The original image to which the overlay is applied.
image (`PIL.Image.Image`):
The image to overlay onto the original.
- crop_coords (`Tuple[int, int, int, int]`, *optional*):
+ crop_coords (`tuple[int, int, int, int]`, *optional*):
Coordinates to crop the image. If provided, the image will be cropped accordingly.
Returns:
@@ -853,7 +847,7 @@ def __init__(
vae_scale_factor: int = 8,
vae_latent_channels: int = 4,
resample: str = "lanczos",
- reducing_gap: int = None,
+ reducing_gap: int | None = None,
do_normalize: bool = True,
do_binarize: bool = False,
do_convert_grayscale: bool = False,
@@ -887,11 +881,11 @@ def __init__(
def preprocess(
self,
image: PIL.Image.Image,
- mask: PIL.Image.Image = None,
- height: int = None,
- width: int = None,
- padding_mask_crop: Optional[int] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ mask: PIL.Image.Image | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
"""
Preprocess the image and mask.
"""
@@ -944,10 +938,10 @@ def postprocess(
self,
image: torch.Tensor,
output_type: str = "pil",
- original_image: Optional[PIL.Image.Image] = None,
- original_mask: Optional[PIL.Image.Image] = None,
- crops_coords: Optional[Tuple[int, int, int, int]] = None,
- ) -> Tuple[PIL.Image.Image, PIL.Image.Image]:
+ original_image: PIL.Image.Image | None = None,
+ original_mask: PIL.Image.Image | None = None,
+ crops_coords: tuple[int, int, int, int] | None = None,
+ ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
"""
Postprocess the image, optionally apply mask overlay
"""
@@ -998,7 +992,7 @@ def __init__(
super().__init__()
@staticmethod
- def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
+ def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
r"""
Convert a NumPy image or a batch of images to a list of PIL images.
@@ -1007,7 +1001,7 @@ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
The input NumPy array of images, which can be a single image or a batch.
Returns:
- `List[PIL.Image.Image]`:
+ `list[PIL.Image.Image]`:
A list of PIL images converted from the input NumPy array.
"""
if images.ndim == 3:
@@ -1022,12 +1016,12 @@ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
return pil_images
@staticmethod
- def depth_pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:
+ def depth_pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
r"""
Convert a PIL image or a list of PIL images to NumPy arrays.
Args:
- images (`Union[List[PIL.Image.Image], PIL.Image.Image]`):
+ images (`list[PIL.Image.Image, PIL.Image.Image]`):
The input image or list of images to be converted.
Returns:
@@ -1042,7 +1036,7 @@ def depth_pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) ->
return images
@staticmethod
- def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
+ def rgblike_to_depthmap(image: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
r"""
Convert an RGB-like depth image to a depth map.
"""
@@ -1079,7 +1073,7 @@ def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndar
else:
raise TypeError("Input image must be a torch.Tensor or np.ndarray")
- def numpy_to_depth(self, images: np.ndarray) -> List[PIL.Image.Image]:
+ def numpy_to_depth(self, images: np.ndarray) -> list[PIL.Image.Image]:
r"""
Convert a NumPy depth image or a batch of images to a list of PIL images.
@@ -1088,7 +1082,7 @@ def numpy_to_depth(self, images: np.ndarray) -> List[PIL.Image.Image]:
The input NumPy array of depth images, which can be a single image or a batch.
Returns:
- `List[PIL.Image.Image]`:
+ `list[PIL.Image.Image]`:
A list of PIL images converted from the input NumPy depth images.
"""
if images.ndim == 3:
@@ -1111,8 +1105,8 @@ def postprocess(
self,
image: torch.Tensor,
output_type: str = "pil",
- do_denormalize: Optional[List[bool]] = None,
- ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]:
+ do_denormalize: list[bool] | None = None,
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
"""
Postprocess the image output from tensor to `output_type`.
@@ -1121,7 +1115,7 @@ def postprocess(
The image input, should be a pytorch tensor with shape `B x C x H x W`.
output_type (`str`, *optional*, defaults to `pil`):
The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
- do_denormalize (`List[bool]`, *optional*, defaults to `None`):
+ do_denormalize (`list[bool]`, *optional*, defaults to `None`):
Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
`VaeImageProcessor` config.
@@ -1159,29 +1153,29 @@ def postprocess(
def preprocess(
self,
- rgb: Union[torch.Tensor, PIL.Image.Image, np.ndarray],
- depth: Union[torch.Tensor, PIL.Image.Image, np.ndarray],
- height: Optional[int] = None,
- width: Optional[int] = None,
- target_res: Optional[int] = None,
+ rgb: torch.Tensor | PIL.Image.Image | np.ndarray,
+ depth: torch.Tensor | PIL.Image.Image | np.ndarray,
+ height: int | None = None,
+ width: int | None = None,
+ target_res: int | None = None,
) -> torch.Tensor:
r"""
Preprocess the image input. Accepted formats are PIL images, NumPy arrays, or PyTorch tensors.
Args:
- rgb (`Union[torch.Tensor, PIL.Image.Image, np.ndarray]`):
+ rgb (`torch.Tensor | PIL.Image.Image | np.ndarray`):
The RGB input image, which can be a single image or a batch.
- depth (`Union[torch.Tensor, PIL.Image.Image, np.ndarray]`):
+ depth (`torch.Tensor | PIL.Image.Image | np.ndarray`):
The depth input image, which can be a single image or a batch.
- height (`Optional[int]`, *optional*, defaults to `None`):
+ height (`int | None`, *optional*, defaults to `None`):
The desired height of the processed image. If `None`, defaults to the height of the input image.
- width (`Optional[int]`, *optional*, defaults to `None`):
+ width (`int | None`, *optional*, defaults to `None`):
The desired width of the processed image. If `None`, defaults to the width of the input image.
- target_res (`Optional[int]`, *optional*, defaults to `None`):
+ target_res (`int | None`, *optional*, defaults to `None`):
Target resolution for resizing the images. If specified, overrides height and width.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing the processed RGB and depth images as PyTorch tensors.
"""
supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
@@ -1419,7 +1413,7 @@ def __init__(
)
@staticmethod
- def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[int, int]:
+ def classify_height_width_bin(height: int, width: int, ratios: dict) -> tuple[int, int]:
r"""
Returns the binned height and width based on the aspect ratio.
@@ -1429,7 +1423,7 @@ def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[in
ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width).
Returns:
- `Tuple[int, int]`: The closest binned height and width.
+ `tuple[int, int]`: The closest binned height and width.
"""
ar = float(height / width)
closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py
index dca4758ba038..13bb44e4a2a6 100644
--- a/src/diffusers/loaders/ip_adapter.py
+++ b/src/diffusers/loaders/ip_adapter.py
@@ -13,7 +13,7 @@
# limitations under the License.
from pathlib import Path
-from typing import Dict, List, Optional, Union
+from typing import List, Union
import torch
import torch.nn.functional as F
@@ -57,15 +57,15 @@ class IPAdapterMixin:
@validate_hf_hub_args
def load_ip_adapter(
self,
- pretrained_model_name_or_path_or_dict: Union[str, List[str], Dict[str, torch.Tensor]],
- subfolder: Union[str, List[str]],
- weight_name: Union[str, List[str]],
- image_encoder_folder: Optional[str] = "image_encoder",
+ pretrained_model_name_or_path_or_dict: str | list[str] | dict[str, torch.Tensor],
+ subfolder: str | list[str],
+ weight_name: str | list[str],
+ image_encoder_folder: str | None = "image_encoder",
**kwargs,
):
"""
Parameters:
- pretrained_model_name_or_path_or_dict (`str` or `List[str]` or `os.PathLike` or `List[os.PathLike]` or `dict` or `List[dict]`):
+ pretrained_model_name_or_path_or_dict (`str` or `list[str]` or `os.PathLike` or `list[os.PathLike]` or `dict` or `list[dict]`):
Can be either:
- A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
@@ -74,10 +74,10 @@ def load_ip_adapter(
with [`ModelMixin.save_pretrained`].
- A [torch state
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
- subfolder (`str` or `List[str]`):
+ subfolder (`str` or `list[str]`):
The subfolder location of a model file within a larger model repository on the Hub or locally. If a
list is passed, it should have the same length as `weight_name`.
- weight_name (`str` or `List[str]`):
+ weight_name (`str` or `list[str]`):
The name of the weight file to load. If a list is passed, it should have the same length as
`subfolder`.
image_encoder_folder (`str`, *optional*, defaults to `image_encoder`):
@@ -87,14 +87,14 @@ def load_ip_adapter(
`image_encoder_folder="image_encoder"`. If the image encoder is located in a folder other than
`subfolder`, you should pass the path to the folder that contains image encoder weights, for example,
`image_encoder_folder="different_subfolder/image_encoder"`.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
@@ -358,14 +358,14 @@ class ModularIPAdapterMixin:
@validate_hf_hub_args
def load_ip_adapter(
self,
- pretrained_model_name_or_path_or_dict: Union[str, List[str], Dict[str, torch.Tensor]],
- subfolder: Union[str, List[str]],
- weight_name: Union[str, List[str]],
+ pretrained_model_name_or_path_or_dict: str | list[str] | dict[str, torch.Tensor],
+ subfolder: str | list[str],
+ weight_name: str | list[str],
**kwargs,
):
"""
Parameters:
- pretrained_model_name_or_path_or_dict (`str` or `List[str]` or `os.PathLike` or `List[os.PathLike]` or `dict` or `List[dict]`):
+ pretrained_model_name_or_path_or_dict (`str` or `list[str]` or `os.PathLike` or `list[os.PathLike]` or `dict` or `list[dict]`):
Can be either:
- A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
@@ -374,20 +374,20 @@ def load_ip_adapter(
with [`ModelMixin.save_pretrained`].
- A [torch state
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
- subfolder (`str` or `List[str]`):
+ subfolder (`str` or `list[str]`):
The subfolder location of a model file within a larger model repository on the Hub or locally. If a
list is passed, it should have the same length as `weight_name`.
- weight_name (`str` or `List[str]`):
+ weight_name (`str` or `list[str]`):
The name of the weight file to load. If a list is passed, it should have the same length as
`subfolder`.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
@@ -608,17 +608,17 @@ class FluxIPAdapterMixin:
@validate_hf_hub_args
def load_ip_adapter(
self,
- pretrained_model_name_or_path_or_dict: Union[str, List[str], Dict[str, torch.Tensor]],
- weight_name: Union[str, List[str]],
- subfolder: Optional[Union[str, List[str]]] = "",
- image_encoder_pretrained_model_name_or_path: Optional[str] = "image_encoder",
- image_encoder_subfolder: Optional[str] = "",
+ pretrained_model_name_or_path_or_dict: str | list[str] | dict[str, torch.Tensor],
+ weight_name: str | list[str],
+ subfolder: str | list[str] | None = "",
+ image_encoder_pretrained_model_name_or_path: str | None = "image_encoder",
+ image_encoder_subfolder: str | None = "",
image_encoder_dtype: torch.dtype = torch.float16,
**kwargs,
):
"""
Parameters:
- pretrained_model_name_or_path_or_dict (`str` or `List[str]` or `os.PathLike` or `List[os.PathLike]` or `dict` or `List[dict]`):
+ pretrained_model_name_or_path_or_dict (`str` or `list[str]` or `os.PathLike` or `list[os.PathLike]` or `dict` or `list[dict]`):
Can be either:
- A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
@@ -627,10 +627,10 @@ def load_ip_adapter(
with [`ModelMixin.save_pretrained`].
- A [torch state
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
- subfolder (`str` or `List[str]`):
+ subfolder (`str` or `list[str]`):
The subfolder location of a model file within a larger model repository on the Hub or locally. If a
list is passed, it should have the same length as `weight_name`.
- weight_name (`str` or `List[str]`):
+ weight_name (`str` or `list[str]`):
The name of the weight file to load. If a list is passed, it should have the same length as
`weight_name`.
image_encoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `./image_encoder`):
@@ -640,14 +640,14 @@ def load_ip_adapter(
hosted on the Hub.
- A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
with [`ModelMixin.save_pretrained`].
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
@@ -797,13 +797,13 @@ def load_ip_adapter(
# load ip-adapter into transformer
self.transformer._load_ip_adapter_weights(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage)
- def set_ip_adapter_scale(self, scale: Union[float, List[float], List[List[float]]]):
+ def set_ip_adapter_scale(self, scale: float | list[float] | list[list[float]]):
"""
Set IP-Adapter scales per-transformer block. Input `scale` could be a single config or a list of configs for
granular control over each IP-Adapter behavior. A config can be a float or a list.
- `float` is converted to list and repeated for the number of blocks and the number of IP adapters. `List[float]`
- length match the number of blocks, it is repeated for each IP adapter. `List[List[float]]` must match the
+ `float` is converted to list and repeated for the number of blocks and the number of IP adapters. `list[float]`
+ length match the number of blocks, it is repeated for each IP adapter. `list[list[float]]` must match the
number of IP adapters and each must match the number of blocks.
Example:
@@ -918,10 +918,10 @@ def is_ip_adapter_active(self) -> bool:
@validate_hf_hub_args
def load_ip_adapter(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
weight_name: str = "ip-adapter.safetensors",
- subfolder: Optional[str] = None,
- image_encoder_folder: Optional[str] = "image_encoder",
+ subfolder: str | None = None,
+ image_encoder_folder: str | None = "image_encoder",
**kwargs,
) -> None:
"""
@@ -947,13 +947,13 @@ def load_ip_adapter(
`image_encoder_folder="image_encoder"`. If the image encoder is located in a folder other than
`subfolder`, you should pass the path to the folder that contains image encoder weights, for example,
`image_encoder_folder="different_subfolder/image_encoder"`.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py
index 3d75a7d875a4..5b5579664b55 100644
--- a/src/diffusers/loaders/lora_base.py
+++ b/src/diffusers/loaders/lora_base.py
@@ -11,13 +11,14 @@
# 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.
+from __future__ import annotations
import copy
import inspect
import json
import os
from pathlib import Path
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import safetensors
import torch
@@ -77,7 +78,7 @@ def fuse_text_encoder_lora(text_encoder, lora_scale=1.0, safe_fusing=False, adap
Controls how much to influence the outputs with the LoRA parameters.
safe_fusing (`bool`, defaults to `False`):
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
- adapter_names (`List[str]` or `str`):
+ adapter_names (`list[str]` or `str`):
The names of the adapters to use.
"""
merge_kwargs = {"safe_merge": safe_fusing}
@@ -116,20 +117,20 @@ def unfuse_text_encoder_lora(text_encoder):
def set_adapters_for_text_encoder(
- adapter_names: Union[List[str], str],
- text_encoder: Optional["PreTrainedModel"] = None, # noqa: F821
- text_encoder_weights: Optional[Union[float, List[float], List[None]]] = None,
+ adapter_names: list[str] | str,
+ text_encoder: "PreTrainedModel" | None = None, # noqa: F821
+ text_encoder_weights: float | list[float] | list[None] | None = None,
):
"""
Sets the adapter layers for the text encoder.
Args:
- adapter_names (`List[str]` or `str`):
+ adapter_names (`list[str]` or `str`):
The names of the adapters to use.
text_encoder (`torch.nn.Module`, *optional*):
The text encoder module to set the adapter layers for. If `None`, it will try to get the `text_encoder`
attribute.
- text_encoder_weights (`List[float]`, *optional*):
+ text_encoder_weights (`list[float]`, *optional*):
The weights to use for the text encoder. If `None`, the weights are set to `1.0` for all the adapters.
"""
if text_encoder is None:
@@ -159,7 +160,7 @@ def process_weights(adapter_names, weights):
set_weights_and_activate_adapters(text_encoder, adapter_names, text_encoder_weights)
-def disable_lora_for_text_encoder(text_encoder: Optional["PreTrainedModel"] = None):
+def disable_lora_for_text_encoder(text_encoder: "PreTrainedModel" | None = None):
"""
Disables the LoRA layers for the text encoder.
@@ -173,7 +174,7 @@ def disable_lora_for_text_encoder(text_encoder: Optional["PreTrainedModel"] = No
set_adapter_layers(text_encoder, enabled=False)
-def enable_lora_for_text_encoder(text_encoder: Optional["PreTrainedModel"] = None):
+def enable_lora_for_text_encoder(text_encoder: "PreTrainedModel" | None = None):
"""
Enables the LoRA layers for the text encoder.
@@ -535,10 +536,10 @@ def unload_lora_weights(self):
def fuse_lora(
self,
- components: List[str] = [],
+ components: list[str] = [],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -547,12 +548,12 @@ def fuse_lora(
> [!WARNING] > This is an experimental API.
Args:
- components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
+ components: (`list[str]`): list of LoRA-injectable components to fuse the LoRAs into.
lora_scale (`float`, defaults to 1.0):
Controls how much to influence the outputs with the LoRA parameters.
safe_fusing (`bool`, defaults to `False`):
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
- adapter_names (`List[str]`, *optional*):
+ adapter_names (`list[str]`, *optional*):
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
Example:
@@ -619,7 +620,7 @@ def fuse_lora(
self._merged_adapters = self._merged_adapters | merged_adapter_names
- def unfuse_lora(self, components: List[str] = [], **kwargs):
+ def unfuse_lora(self, components: list[str] = [], **kwargs):
r"""
Reverses the effect of
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
@@ -627,7 +628,7 @@ def unfuse_lora(self, components: List[str] = [], **kwargs):
> [!WARNING] > This is an experimental API.
Args:
- components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
+ components (`list[str]`): list of LoRA-injectable components to unfuse LoRA from.
unfuse_unet (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
unfuse_text_encoder (`bool`, defaults to `True`):
Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
@@ -674,16 +675,16 @@ def unfuse_lora(self, components: List[str] = [], **kwargs):
def set_adapters(
self,
- adapter_names: Union[List[str], str],
- adapter_weights: Optional[Union[float, Dict, List[float], List[Dict]]] = None,
+ adapter_names: list[str] | str,
+ adapter_weights: float | dict | list[float] | list[dict] | None = None,
):
"""
Set the currently active adapters for use in the pipeline.
Args:
- adapter_names (`List[str]` or `str`):
+ adapter_names (`list[str]` or `str`):
The names of the adapters to use.
- adapter_weights (`Union[List[float], float]`, *optional*):
+ adapter_weights (`list[float, float]`, *optional*):
The adapter(s) weights to use with the UNet. If `None`, the weights are set to `1.0` for all the
adapters.
@@ -835,12 +836,12 @@ def enable_lora(self):
elif issubclass(model.__class__, PreTrainedModel):
enable_lora_for_text_encoder(model)
- def delete_adapters(self, adapter_names: Union[List[str], str]):
+ def delete_adapters(self, adapter_names: list[str] | str):
"""
Delete an adapter's LoRA layers from the pipeline.
Args:
- adapter_names (`Union[List[str], str]`):
+ adapter_names (`list[str, str]`):
The names of the adapters to delete.
Example:
@@ -873,7 +874,7 @@ def delete_adapters(self, adapter_names: Union[List[str], str]):
for adapter_name in adapter_names:
delete_adapter_layers(model, adapter_name)
- def get_active_adapters(self) -> List[str]:
+ def get_active_adapters(self) -> list[str]:
"""
Gets the list of the current active adapters.
@@ -906,7 +907,7 @@ def get_active_adapters(self) -> List[str]:
return active_adapters
- def get_list_adapters(self) -> Dict[str, List[str]]:
+ def get_list_adapters(self) -> dict[str, list[str]]:
"""
Gets the current list of all available adapters in the pipeline.
"""
@@ -928,7 +929,7 @@ def get_list_adapters(self) -> Dict[str, List[str]]:
return set_adapters
- def set_lora_device(self, adapter_names: List[str], device: Union[torch.device, str, int]) -> None:
+ def set_lora_device(self, adapter_names: list[str], device: torch.device | str | int) -> None:
"""
Moves the LoRAs listed in `adapter_names` to a target device. Useful for offloading the LoRA to the CPU in case
you want to load multiple adapters and free some GPU memory.
@@ -955,9 +956,9 @@ def set_lora_device(self, adapter_names: List[str], device: Union[torch.device,
```
Args:
- adapter_names (`List[str]`):
- List of adapters to send device to.
- device (`Union[torch.device, str, int]`):
+ adapter_names (`list[str]`):
+ list of adapters to send device to.
+ device (`torch.device | str | int`):
Device to send the adapters to. Can be either a torch device, a str or an integer.
"""
if not USE_PEFT_BACKEND:
@@ -1007,13 +1008,13 @@ def pack_weights(layers, prefix):
@staticmethod
def write_lora_layers(
- state_dict: Dict[str, torch.Tensor],
+ state_dict: dict[str, torch.Tensor],
save_directory: str,
is_main_process: bool,
weight_name: str,
save_function: Callable,
safe_serialization: bool,
- lora_adapter_metadata: Optional[dict] = None,
+ lora_adapter_metadata: dict | None = None,
):
"""Writes the state dict of the LoRA layers (optionally with metadata) to disk."""
if os.path.isfile(save_directory):
@@ -1059,9 +1060,9 @@ def save_function(weights, filename):
@classmethod
def _save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- lora_layers: Dict[str, Dict[str, Union[torch.nn.Module, torch.Tensor]]],
- lora_metadata: Dict[str, Optional[dict]],
+ save_directory: str | os.PathLike,
+ lora_layers: dict[str, dict[str, torch.nn.Module | torch.Tensor]],
+ lora_metadata: dict[str, dict | None],
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py
index 9fc1a24eb994..f56632ced819 100644
--- a/src/diffusers/loaders/lora_conversion_utils.py
+++ b/src/diffusers/loaders/lora_conversion_utils.py
@@ -13,7 +13,6 @@
# limitations under the License.
import re
-from typing import List
import torch
@@ -1021,7 +1020,7 @@ def handle_qkv(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
return new_state_dict
-def _custom_replace(key: str, substrings: List[str]) -> str:
+def _custom_replace(key: str, substrings: list[str]) -> str:
# Replaces the "."s with "_"s upto the `substrings`.
# Example:
# lora_unet.foo.bar.lora_A.weight -> lora_unet_foo_bar.lora_A.weight
@@ -2321,6 +2320,14 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
prefix = "diffusion_model."
original_state_dict = {k[len(prefix) :]: v for k, v in state_dict.items()}
+ has_lora_down_up = any("lora_down" in k or "lora_up" in k for k in original_state_dict.keys())
+ if has_lora_down_up:
+ temp_state_dict = {}
+ for k, v in original_state_dict.items():
+ new_key = k.replace("lora_down", "lora_A").replace("lora_up", "lora_B")
+ temp_state_dict[new_key] = v
+ original_state_dict = temp_state_dict
+
num_double_layers = 0
num_single_layers = 0
for key in original_state_dict.keys():
@@ -2337,13 +2344,15 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
attn_prefix = f"single_transformer_blocks.{sl}.attn"
for lora_key in lora_keys:
- converted_state_dict[f"{attn_prefix}.to_qkv_mlp_proj.{lora_key}.weight"] = original_state_dict.pop(
- f"{single_block_prefix}.linear1.{lora_key}.weight"
- )
+ linear1_key = f"{single_block_prefix}.linear1.{lora_key}.weight"
+ if linear1_key in original_state_dict:
+ converted_state_dict[f"{attn_prefix}.to_qkv_mlp_proj.{lora_key}.weight"] = original_state_dict.pop(
+ linear1_key
+ )
- converted_state_dict[f"{attn_prefix}.to_out.{lora_key}.weight"] = original_state_dict.pop(
- f"{single_block_prefix}.linear2.{lora_key}.weight"
- )
+ linear2_key = f"{single_block_prefix}.linear2.{lora_key}.weight"
+ if linear2_key in original_state_dict:
+ converted_state_dict[f"{attn_prefix}.to_out.{lora_key}.weight"] = original_state_dict.pop(linear2_key)
for dl in range(num_double_layers):
transformer_block_prefix = f"transformer_blocks.{dl}"
@@ -2352,6 +2361,10 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
for attn_type in attn_types:
attn_prefix = f"{transformer_block_prefix}.attn"
qkv_key = f"double_blocks.{dl}.{attn_type}.qkv.{lora_key}.weight"
+
+ if qkv_key not in original_state_dict:
+ continue
+
fused_qkv_weight = original_state_dict.pop(qkv_key)
if lora_key == "lora_A":
@@ -2383,8 +2396,9 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
for org_proj, diff_proj in proj_mappings:
for lora_key in lora_keys:
original_key = f"double_blocks.{dl}.{org_proj}.{lora_key}.weight"
- diffusers_key = f"{transformer_block_prefix}.{diff_proj}.{lora_key}.weight"
- converted_state_dict[diffusers_key] = original_state_dict.pop(original_key)
+ if original_key in original_state_dict:
+ diffusers_key = f"{transformer_block_prefix}.{diff_proj}.{lora_key}.weight"
+ converted_state_dict[diffusers_key] = original_state_dict.pop(original_key)
mlp_mappings = [
("img_mlp.0", "ff.linear_in"),
@@ -2395,8 +2409,27 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
for org_mlp, diff_mlp in mlp_mappings:
for lora_key in lora_keys:
original_key = f"double_blocks.{dl}.{org_mlp}.{lora_key}.weight"
- diffusers_key = f"{transformer_block_prefix}.{diff_mlp}.{lora_key}.weight"
- converted_state_dict[diffusers_key] = original_state_dict.pop(original_key)
+ if original_key in original_state_dict:
+ diffusers_key = f"{transformer_block_prefix}.{diff_mlp}.{lora_key}.weight"
+ converted_state_dict[diffusers_key] = original_state_dict.pop(original_key)
+
+ extra_mappings = {
+ "img_in": "x_embedder",
+ "txt_in": "context_embedder",
+ "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
+ "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
+ "final_layer.linear": "proj_out",
+ "final_layer.adaLN_modulation.1": "norm_out.linear",
+ "single_stream_modulation.lin": "single_stream_modulation.linear",
+ "double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
+ "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
+ }
+
+ for org_key, diff_key in extra_mappings.items():
+ for lora_key in lora_keys:
+ original_key = f"{org_key}.{lora_key}.weight"
+ if original_key in original_state_dict:
+ converted_state_dict[f"{diff_key}.{lora_key}.weight"] = original_state_dict.pop(original_key)
if len(original_state_dict) > 0:
raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.")
diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py
index 24d1fd7b9308..3423a88d3368 100644
--- a/src/diffusers/loaders/lora_pipeline.py
+++ b/src/diffusers/loaders/lora_pipeline.py
@@ -13,7 +13,7 @@
# limitations under the License.
import os
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import torch
from huggingface_hub.utils import validate_hf_hub_args
@@ -141,8 +141,8 @@ class StableDiffusionLoraLoaderMixin(LoraBaseMixin):
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -244,7 +244,7 @@ def load_lora_weights(
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -264,14 +264,14 @@ def lora_state_dict(
- A [torch state
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
@@ -371,7 +371,7 @@ def load_lora_into_unet(
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
encoder lora layers.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -433,7 +433,7 @@ def load_lora_into_text_encoder(
state_dict (`dict`):
A standard state dict containing the lora layer parameters. The key should be prefixed with an
additional `text_encoder` to distinguish between unet lora layers.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -473,9 +473,9 @@ def load_lora_into_text_encoder(
@classmethod
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_lora_layers: Dict[str, torch.nn.Module] = None,
+ save_directory: str | os.PathLike,
+ unet_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
+ text_encoder_lora_layers: dict[str, torch.nn.Module] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
@@ -489,9 +489,9 @@ def save_lora_weights(
Arguments:
save_directory (`str` or `os.PathLike`):
Directory to save LoRA parameters to. Will be created if it doesn't exist.
- unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
+ unet_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`):
State dict of the LoRA layers corresponding to the `unet`.
- text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
+ text_encoder_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`):
State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
encoder LoRA state dict because it comes from 🤗 Transformers.
is_main_process (`bool`, *optional*, defaults to `True`):
@@ -535,10 +535,10 @@ def save_lora_weights(
def fuse_lora(
self,
- components: List[str] = ["unet", "text_encoder"],
+ components: list[str] = ["unet", "text_encoder"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -547,12 +547,12 @@ def fuse_lora(
> [!WARNING] > This is an experimental API.
Args:
- components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
+ components: (`list[str]`): list of LoRA-injectable components to fuse the LoRAs into.
lora_scale (`float`, defaults to 1.0):
Controls how much to influence the outputs with the LoRA parameters.
safe_fusing (`bool`, defaults to `False`):
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
- adapter_names (`List[str]`, *optional*):
+ adapter_names (`list[str]`, *optional*):
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
Example:
@@ -576,7 +576,7 @@ def fuse_lora(
**kwargs,
)
- def unfuse_lora(self, components: List[str] = ["unet", "text_encoder"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["unet", "text_encoder"], **kwargs):
r"""
Reverses the effect of
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
@@ -584,7 +584,7 @@ def unfuse_lora(self, components: List[str] = ["unet", "text_encoder"], **kwargs
> [!WARNING] > This is an experimental API.
Args:
- components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
+ components (`list[str]`): list of LoRA-injectable components to unfuse LoRA from.
unfuse_unet (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
unfuse_text_encoder (`bool`, defaults to `True`):
Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
@@ -606,8 +606,8 @@ class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin):
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -683,7 +683,7 @@ def load_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.lora_state_dict
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -703,14 +703,14 @@ def lora_state_dict(
- A [torch state
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
@@ -811,7 +811,7 @@ def load_lora_into_unet(
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
encoder lora layers.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -874,7 +874,7 @@ def load_lora_into_text_encoder(
state_dict (`dict`):
A standard state dict containing the lora layer parameters. The key should be prefixed with an
additional `text_encoder` to distinguish between unet lora layers.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -914,10 +914,10 @@ def load_lora_into_text_encoder(
@classmethod
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ unet_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
+ text_encoder_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
+ text_encoder_2_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
@@ -961,10 +961,10 @@ def save_lora_weights(
def fuse_lora(
self,
- components: List[str] = ["unet", "text_encoder", "text_encoder_2"],
+ components: list[str] = ["unet", "text_encoder", "text_encoder_2"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -978,7 +978,7 @@ def fuse_lora(
**kwargs,
)
- def unfuse_lora(self, components: List[str] = ["unet", "text_encoder", "text_encoder_2"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["unet", "text_encoder", "text_encoder_2"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -1002,7 +1002,7 @@ class SD3LoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -1054,7 +1054,7 @@ def lora_state_dict(
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
adapter_name=None,
hotswap: bool = False,
**kwargs,
@@ -1170,7 +1170,7 @@ def load_lora_into_text_encoder(
state_dict (`dict`):
A standard state dict containing the lora layer parameters. The key should be prefixed with an
additional `text_encoder` to distinguish between unet lora layers.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -1211,10 +1211,10 @@ def load_lora_into_text_encoder(
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.save_lora_weights with unet->transformer
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
+ text_encoder_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
+ text_encoder_2_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
@@ -1259,10 +1259,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.fuse_lora with unet->transformer
def fuse_lora(
self,
- components: List[str] = ["transformer", "text_encoder", "text_encoder_2"],
+ components: list[str] = ["transformer", "text_encoder", "text_encoder_2"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -1277,7 +1277,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.unfuse_lora with unet->transformer
- def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder", "text_encoder_2"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer", "text_encoder", "text_encoder_2"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -1297,7 +1297,7 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin):
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.lora_state_dict
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -1350,8 +1350,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -1425,13 +1425,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -1459,10 +1459,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -1477,7 +1477,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer", "text_encoder"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -1501,7 +1501,7 @@ class FluxLoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
return_alphas: bool = False,
**kwargs,
):
@@ -1624,8 +1624,8 @@ def lora_state_dict(
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -1763,7 +1763,7 @@ def _load_norm_into_transformer(
transformer,
prefix=None,
discard_original_layers=False,
- ) -> Dict[str, torch.Tensor]:
+ ) -> dict[str, torch.Tensor]:
# Remove prefix if present
prefix = prefix or cls.transformer_name
for key in list(state_dict.keys()):
@@ -1832,7 +1832,7 @@ def load_lora_into_text_encoder(
state_dict (`dict`):
A standard state dict containing the lora layer parameters. The key should be prefixed with an
additional `text_encoder` to distinguish between unet lora layers.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -1873,9 +1873,9 @@ def load_lora_into_text_encoder(
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.save_lora_weights with unet->transformer
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
- text_encoder_lora_layers: Dict[str, torch.nn.Module] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
+ text_encoder_lora_layers: dict[str, torch.nn.Module] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
@@ -1889,9 +1889,9 @@ def save_lora_weights(
Arguments:
save_directory (`str` or `os.PathLike`):
Directory to save LoRA parameters to. Will be created if it doesn't exist.
- transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
+ transformer_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`):
State dict of the LoRA layers corresponding to the `transformer`.
- text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
+ text_encoder_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`):
State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
encoder LoRA state dict because it comes from 🤗 Transformers.
is_main_process (`bool`, *optional*, defaults to `True`):
@@ -1935,10 +1935,10 @@ def save_lora_weights(
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -1965,7 +1965,7 @@ def fuse_lora(
**kwargs,
)
- def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer", "text_encoder"], **kwargs):
r"""
Reverses the effect of
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
@@ -1973,7 +1973,7 @@ def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], *
> [!WARNING] > This is an experimental API.
Args:
- components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
+ components (`list[str]`): list of LoRA-injectable components to unfuse LoRA from.
"""
transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
if hasattr(transformer, "_transformer_norm_layers") and transformer._transformer_norm_layers:
@@ -2322,7 +2322,7 @@ def load_lora_into_text_encoder(
state_dict (`dict`):
A standard state dict containing the lora layer parameters. The key should be prefixed with an
additional `text_encoder` to distinguish between unet lora layers.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -2362,9 +2362,9 @@ def load_lora_into_text_encoder(
@classmethod
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- text_encoder_lora_layers: Dict[str, torch.nn.Module] = None,
- transformer_lora_layers: Dict[str, torch.nn.Module] = None,
+ save_directory: str | os.PathLike,
+ text_encoder_lora_layers: dict[str, torch.nn.Module] = None,
+ transformer_lora_layers: dict[str, torch.nn.Module] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
@@ -2376,9 +2376,9 @@ def save_lora_weights(
Arguments:
save_directory (`str` or `os.PathLike`):
Directory to save LoRA parameters to. Will be created if it doesn't exist.
- unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
+ unet_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`):
State dict of the LoRA layers corresponding to the `unet`.
- text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
+ text_encoder_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`):
State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
encoder LoRA state dict because it comes from 🤗 Transformers.
is_main_process (`bool`, *optional*, defaults to `True`):
@@ -2427,7 +2427,7 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin):
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -2479,8 +2479,8 @@ def lora_state_dict(
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -2553,13 +2553,13 @@ def load_lora_into_transformer(
@classmethod
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -2586,10 +2586,10 @@ def save_lora_weights(
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -2603,7 +2603,7 @@ def fuse_lora(
**kwargs,
)
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -2623,7 +2623,7 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin):
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -2676,8 +2676,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -2751,13 +2751,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -2785,10 +2785,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -2803,7 +2803,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -2822,7 +2822,7 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -2879,8 +2879,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -2954,13 +2954,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -2988,10 +2988,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -3006,7 +3006,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -3026,7 +3026,7 @@ class LTX2LoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -3088,8 +3088,8 @@ def lora_state_dict(
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -3181,13 +3181,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -3215,10 +3215,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -3233,7 +3233,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -3253,7 +3253,7 @@ class SanaLoraLoaderMixin(LoraBaseMixin):
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -3306,8 +3306,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -3381,13 +3381,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -3415,10 +3415,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -3433,7 +3433,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -3452,7 +3452,7 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -3509,8 +3509,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -3584,13 +3584,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -3618,10 +3618,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -3636,7 +3636,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -3655,7 +3655,7 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -3713,8 +3713,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -3788,13 +3788,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -3822,10 +3822,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -3840,7 +3840,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -3860,7 +3860,7 @@ class KandinskyLoraLoaderMixin(LoraBaseMixin):
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -3913,8 +3913,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -3988,13 +3988,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -4022,10 +4022,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -4040,7 +4040,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -4059,7 +4059,7 @@ class WanLoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -4163,8 +4163,8 @@ def _maybe_expand_t2v_lora_for_i2v(
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -4262,13 +4262,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -4296,10 +4296,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -4314,7 +4314,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -4334,7 +4334,7 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin):
# Copied from diffusers.loaders.lora_pipeline.WanLoraLoaderMixin.lora_state_dict
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -4440,8 +4440,8 @@ def _maybe_expand_t2v_lora_for_i2v(
# Copied from diffusers.loaders.lora_pipeline.WanLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -4539,13 +4539,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -4573,10 +4573,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -4591,7 +4591,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -4611,7 +4611,7 @@ class CogView4LoraLoaderMixin(LoraBaseMixin):
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.lora_state_dict
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -4664,8 +4664,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -4739,13 +4739,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -4773,10 +4773,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -4791,7 +4791,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -4810,7 +4810,7 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -4867,8 +4867,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -4942,13 +4942,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -4976,10 +4976,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -4994,7 +4994,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -5013,7 +5013,7 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -5073,8 +5073,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -5148,13 +5148,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -5182,10 +5182,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -5200,7 +5200,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -5219,7 +5219,7 @@ class ZImageLoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -5279,8 +5279,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -5354,13 +5354,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -5388,10 +5388,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -5406,7 +5406,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
@@ -5425,7 +5425,7 @@ class Flux2LoraLoaderMixin(LoraBaseMixin):
@validate_hf_hub_args
def lora_state_dict(
cls,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
r"""
@@ -5472,6 +5472,10 @@ def lora_state_dict(
logger.warning(warn_msg)
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
+ is_peft_format = any(k.startswith("base_model.model.") for k in state_dict)
+ if is_peft_format:
+ state_dict = {k.replace("base_model.model.", "diffusion_model."): v for k, v in state_dict.items()}
+
is_ai_toolkit = any(k.startswith("diffusion_model.") for k in state_dict)
if is_ai_toolkit:
state_dict = _convert_non_diffusers_flux2_lora_to_diffusers(state_dict)
@@ -5482,8 +5486,8 @@ def lora_state_dict(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
- pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
- adapter_name: Optional[str] = None,
+ pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+ adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
@@ -5557,13 +5561,13 @@ def load_lora_into_transformer(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
- save_directory: Union[str, os.PathLike],
- transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+ save_directory: str | os.PathLike,
+ transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
- transformer_lora_adapter_metadata: Optional[dict] = None,
+ transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
@@ -5591,10 +5595,10 @@ def save_lora_weights(
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
- components: List[str] = ["transformer"],
+ components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
- adapter_names: Optional[List[str]] = None,
+ adapter_names: list[str] | None = None,
**kwargs,
):
r"""
@@ -5609,7 +5613,7 @@ def fuse_lora(
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
- def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
+ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
diff --git a/src/diffusers/loaders/peft.py b/src/diffusers/loaders/peft.py
index 6cdd66a0f2c4..80fb6a72869a 100644
--- a/src/diffusers/loaders/peft.py
+++ b/src/diffusers/loaders/peft.py
@@ -17,7 +17,7 @@
import os
from functools import partial
from pathlib import Path
-from typing import Dict, List, Literal, Optional, Union
+from typing import Literal
import safetensors
import torch
@@ -88,7 +88,7 @@ class PeftAdapterMixin:
_hf_peft_config_loaded = False
# kwargs for prepare_model_for_compiled_hotswap, if required
- _prepare_lora_hotswap_kwargs: Optional[dict] = None
+ _prepare_lora_hotswap_kwargs: dict | None = None
@classmethod
# Copied from diffusers.loaders.lora_base.LoraBaseMixin._optionally_disable_offloading
@@ -114,13 +114,13 @@ def load_lora_adapter(
prefix (`str`, *optional*): Prefix to filter the state dict.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
@@ -134,7 +134,7 @@ def load_lora_adapter(
allowed by Git.
subfolder (`str`, *optional*, defaults to `""`):
The subfolder location of a model file within a larger model repository on the Hub or locally.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -405,7 +405,7 @@ def save_lora_adapter(
adapter_name: str = "default",
upcast_before_saving: bool = False,
safe_serialization: bool = True,
- weight_name: Optional[str] = None,
+ weight_name: str | None = None,
):
"""
Save the LoRA parameters corresponding to the underlying model.
@@ -469,14 +469,14 @@ def save_function(weights, filename):
def set_adapters(
self,
- adapter_names: Union[List[str], str],
- weights: Optional[Union[float, Dict, List[float], List[Dict], List[None]]] = None,
+ adapter_names: list[str] | str,
+ weights: float | dict | list[float] | list[dict] | list[None] | None = None,
):
"""
Set the currently active adapters for use in the diffusion network (e.g. unet, transformer, etc.).
Args:
- adapter_names (`List[str]` or `str`):
+ adapter_names (`list[str]` or `str`):
The names of the adapters to use.
weights (`Union[List[float], float]`, *optional*):
The adapter(s) weights to use with the UNet. If `None`, the weights are set to `1.0` for all the
@@ -561,7 +561,7 @@ def add_adapter(self, adapter_config, adapter_name: str = "default") -> None:
inject_adapter_in_model(adapter_config, self, adapter_name)
self.set_adapter(adapter_name)
- def set_adapter(self, adapter_name: Union[str, List[str]]) -> None:
+ def set_adapter(self, adapter_name: str | list[str]) -> None:
"""
Sets a specific adapter by forcing the model to only use that adapter and disables the other adapters.
@@ -569,7 +569,7 @@ def set_adapter(self, adapter_name: Union[str, List[str]]) -> None:
[documentation](https://huggingface.co/docs/peft).
Args:
- adapter_name (Union[str, List[str]])):
+ adapter_name (str | list[str])):
The list of adapters to set or the adapter name in the case of a single adapter.
"""
check_peft_version(min_version=MIN_PEFT_VERSION)
@@ -655,7 +655,7 @@ def enable_adapters(self) -> None:
# support for older PEFT versions
module.disable_adapters = False
- def active_adapters(self) -> List[str]:
+ def active_adapters(self) -> list[str]:
"""
Gets the current list of active adapters of the model.
@@ -778,12 +778,12 @@ def enable_lora(self):
raise ValueError("PEFT backend is required for this method.")
set_adapter_layers(self, enabled=True)
- def delete_adapters(self, adapter_names: Union[List[str], str]):
+ def delete_adapters(self, adapter_names: list[str] | str):
"""
Delete an adapter's LoRA layers from the underlying model.
Args:
- adapter_names (`Union[List[str], str]`):
+ adapter_names (`list[str, str]`):
The names (single string or list of strings) of the adapter to delete.
Example:
diff --git a/src/diffusers/loaders/single_file.py b/src/diffusers/loaders/single_file.py
index 667f79437985..0b0d52d4a412 100644
--- a/src/diffusers/loaders/single_file.py
+++ b/src/diffusers/loaders/single_file.py
@@ -286,11 +286,11 @@ def from_single_file(cls, pretrained_model_link_or_path, **kwargs) -> Self:
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py
index ab4340fed120..c7bb2de4437a 100644
--- a/src/diffusers/loaders/single_file_model.py
+++ b/src/diffusers/loaders/single_file_model.py
@@ -15,7 +15,6 @@
import inspect
import re
from contextlib import nullcontext
-from typing import Optional
import torch
from huggingface_hub.utils import validate_hf_hub_args
@@ -235,7 +234,7 @@ class FromOriginalModelMixin:
@classmethod
@validate_hf_hub_args
- def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] = None, **kwargs) -> Self:
+ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = None, **kwargs) -> Self:
r"""
Instantiate a model from pretrained weights saved in the original `.ckpt` or `.safetensors` format. The model
is set in evaluation mode (`model.eval()`) by default.
@@ -262,11 +261,11 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
diff --git a/src/diffusers/loaders/textual_inversion.py b/src/diffusers/loaders/textual_inversion.py
index e4346700a7b8..711de5f81343 100644
--- a/src/diffusers/loaders/textual_inversion.py
+++ b/src/diffusers/loaders/textual_inversion.py
@@ -11,8 +11,9 @@
# 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.
+from __future__ import annotations
+
import json
-from typing import Dict, List, Optional, Union
import safetensors
import torch
@@ -21,7 +22,12 @@
from torch import nn
from ..models.modeling_utils import load_state_dict
-from ..utils import _get_model_file, is_accelerate_available, is_transformers_available, logging
+from ..utils import (
+ _get_model_file,
+ is_accelerate_available,
+ is_transformers_available,
+ logging,
+)
if is_transformers_available():
@@ -114,7 +120,7 @@ class TextualInversionLoaderMixin:
Load Textual Inversion tokens and embeddings to the tokenizer and text encoder.
"""
- def maybe_convert_prompt(self, prompt: Union[str, List[str]], tokenizer: "PreTrainedTokenizer"): # noqa: F821
+ def maybe_convert_prompt(self, prompt: str | list[str], tokenizer: "PreTrainedTokenizer"): # noqa: F821
r"""
Processes prompts that include a special token corresponding to a multi-vector textual inversion embedding to
be replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual
@@ -129,14 +135,14 @@ def maybe_convert_prompt(self, prompt: Union[str, List[str]], tokenizer: "PreTra
Returns:
`str` or list of `str`: The converted prompt
"""
- if not isinstance(prompt, List):
+ if not isinstance(prompt, list):
prompts = [prompt]
else:
prompts = prompt
prompts = [self._maybe_convert_prompt(p, tokenizer) for p in prompts]
- if not isinstance(prompt, List):
+ if not isinstance(prompt, list):
return prompts[0]
return prompts
@@ -265,10 +271,10 @@ def _extend_tokens_and_embeddings(tokens, embeddings, tokenizer):
@validate_hf_hub_args
def load_textual_inversion(
self,
- pretrained_model_name_or_path: Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]],
- token: Optional[Union[str, List[str]]] = None,
- tokenizer: Optional["PreTrainedTokenizer"] = None, # noqa: F821
- text_encoder: Optional["PreTrainedModel"] = None, # noqa: F821
+ pretrained_model_name_or_path: str | list[str] | dict[str, torch.Tensor] | list[dict[str, torch.Tensor]],
+ token: str | list[str] | None = None,
+ tokenizer: "PreTrainedTokenizer" | None = None, # noqa: F821
+ text_encoder: "PreTrainedModel" | None = None, # noqa: F821
**kwargs,
):
r"""
@@ -276,7 +282,7 @@ def load_textual_inversion(
Automatic1111 formats are supported).
Parameters:
- pretrained_model_name_or_path (`str` or `os.PathLike` or `List[str or os.PathLike]` or `Dict` or `List[Dict]`):
+ pretrained_model_name_or_path (`str` or `os.PathLike` or `list[str or os.PathLike]` or `Dict` or `list[Dict]`):
Can be either one of the following or a list of them:
- A string, the *model id* (for example `sd-concepts-library/low-poly-hd-logos-icons`) of a
@@ -287,7 +293,7 @@ def load_textual_inversion(
- A [torch state
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
- token (`str` or `List[str]`, *optional*):
+ token (`str` or `list[str]`, *optional*):
Override the token to use for the textual inversion weights. If `pretrained_model_name_or_path` is a
list, then `token` must also be a list of equal length.
text_encoder ([`~transformers.CLIPTextModel`], *optional*):
@@ -301,14 +307,14 @@ def load_textual_inversion(
- The saved textual inversion file is in 🤗 Diffusers format, but was saved under a specific weight
name such as `text_inv.bin`.
- The saved textual inversion file is in the Automatic1111 format.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
@@ -460,9 +466,9 @@ def load_textual_inversion(
def unload_textual_inversion(
self,
- tokens: Optional[Union[str, List[str]]] = None,
- tokenizer: Optional["PreTrainedTokenizer"] = None,
- text_encoder: Optional["PreTrainedModel"] = None,
+ tokens: str | list[str] | None = None,
+ tokenizer: "PreTrainedTokenizer" | None = None,
+ text_encoder: "PreTrainedModel" | None = None,
):
r"""
Unload Textual Inversion embeddings from the text encoder of [`StableDiffusionPipeline`]
diff --git a/src/diffusers/loaders/transformer_sd3.py b/src/diffusers/loaders/transformer_sd3.py
index e3728082efdd..d1c3fff14e60 100644
--- a/src/diffusers/loaders/transformer_sd3.py
+++ b/src/diffusers/loaders/transformer_sd3.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from contextlib import nullcontext
-from typing import Dict
from ..models.attention_processor import SD3IPAdapterJointAttnProcessor2_0
from ..models.embeddings import IPAdapterTimeImageProjection
@@ -29,8 +28,8 @@ class SD3Transformer2DLoadersMixin:
"""Load IP-Adapters and LoRA layers into a `[SD3Transformer2DModel]`."""
def _convert_ip_adapter_attn_to_diffusers(
- self, state_dict: Dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT
- ) -> Dict:
+ self, state_dict: dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT
+ ) -> dict:
if low_cpu_mem_usage:
if is_accelerate_available():
from accelerate import init_empty_weights
@@ -87,7 +86,7 @@ def _convert_ip_adapter_attn_to_diffusers(
return attn_procs
def _convert_ip_adapter_image_proj_to_diffusers(
- self, state_dict: Dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT
+ self, state_dict: dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT
) -> IPAdapterTimeImageProjection:
if low_cpu_mem_usage:
if is_accelerate_available():
@@ -155,7 +154,7 @@ def _convert_ip_adapter_image_proj_to_diffusers(
return image_proj
- def _load_ip_adapter_weights(self, state_dict: Dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT) -> None:
+ def _load_ip_adapter_weights(self, state_dict: dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT) -> None:
"""Sets IP-Adapter attention processors, image projection, and loads state_dict.
Args:
diff --git a/src/diffusers/loaders/unet.py b/src/diffusers/loaders/unet.py
index c5e56af156fc..9dab3bc667ea 100644
--- a/src/diffusers/loaders/unet.py
+++ b/src/diffusers/loaders/unet.py
@@ -15,7 +15,7 @@
from collections import defaultdict
from contextlib import nullcontext
from pathlib import Path
-from typing import Callable, Dict, Union
+from typing import Callable
import safetensors
import torch
@@ -66,7 +66,7 @@ class UNet2DConditionLoadersMixin:
unet_name = UNET_NAME
@validate_hf_hub_args
- def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):
+ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs):
r"""
Load pretrained attention processor layers into [`UNet2DConditionModel`]. Attention processor layers have to be
defined in
@@ -85,14 +85,14 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict
- A [torch state
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only (`bool`, *optional*, defaults to `False`):
@@ -106,7 +106,7 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict
allowed by Git.
subfolder (`str`, *optional*, defaults to `""`):
The subfolder location of a model file within a larger model repository on the Hub or locally.
- network_alphas (`Dict[str, float]`):
+ network_alphas (`dict[str, float]`):
The value of the network alpha used for stable learning and preventing underflow. This value has the
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
@@ -412,7 +412,7 @@ def _optionally_disable_offloading(cls, _pipeline):
def save_attn_procs(
self,
- save_directory: Union[str, os.PathLike],
+ save_directory: str | os.PathLike,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
diff --git a/src/diffusers/loaders/unet_loader_utils.py b/src/diffusers/loaders/unet_loader_utils.py
index d5b0e83cbd9e..d6aff55d6d22 100644
--- a/src/diffusers/loaders/unet_loader_utils.py
+++ b/src/diffusers/loaders/unet_loader_utils.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
-from typing import TYPE_CHECKING, Dict, List, Union
+from typing import TYPE_CHECKING
from torch import nn
@@ -40,9 +40,7 @@ def _translate_into_actual_layer_name(name):
return ".".join((updown, block, attn))
-def _maybe_expand_lora_scales(
- unet: "UNet2DConditionModel", weight_scales: List[Union[float, Dict]], default_scale=1.0
-):
+def _maybe_expand_lora_scales(unet: "UNet2DConditionModel", weight_scales: list[float | dict], default_scale=1.0):
blocks_with_transformer = {
"down": [i for i, block in enumerate(unet.down_blocks) if hasattr(block, "attentions")],
"up": [i for i, block in enumerate(unet.up_blocks) if hasattr(block, "attentions")],
@@ -64,9 +62,9 @@ def _maybe_expand_lora_scales(
def _maybe_expand_lora_scales_for_one_adapter(
- scales: Union[float, Dict],
- blocks_with_transformer: Dict[str, int],
- transformer_per_block: Dict[str, int],
+ scales: float | dict,
+ blocks_with_transformer: dict[str, int],
+ transformer_per_block: dict[str, int],
model: nn.Module,
default_scale: float = 1.0,
):
@@ -74,11 +72,11 @@ def _maybe_expand_lora_scales_for_one_adapter(
Expands the inputs into a more granular dictionary. See the example below for more details.
Parameters:
- scales (`Union[float, Dict]`):
+ scales (`float | Dict`):
Scales dict to expand.
- blocks_with_transformer (`Dict[str, int]`):
+ blocks_with_transformer (`dict[str, int]`):
Dict with keys 'up' and 'down', showing which blocks have transformer layers
- transformer_per_block (`Dict[str, int]`):
+ transformer_per_block (`dict[str, int]`):
Dict with keys 'up' and 'down', showing how many transformer layers each block has
E.g. turns
diff --git a/src/diffusers/loaders/utils.py b/src/diffusers/loaders/utils.py
index 2d39e7bfb7d2..72d6453a8bbf 100644
--- a/src/diffusers/loaders/utils.py
+++ b/src/diffusers/loaders/utils.py
@@ -12,13 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Dict
import torch
class AttnProcsLayers(torch.nn.Module):
- def __init__(self, state_dict: Dict[str, torch.Tensor]):
+ def __init__(self, state_dict: dict[str, torch.Tensor]):
super().__init__()
self.layers = torch.nn.ModuleList(state_dict.values())
self.mapping = dict(enumerate(state_dict.keys()))
diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py
index 4d1db36a7352..96953afa4f4a 100755
--- a/src/diffusers/models/__init__.py
+++ b/src/diffusers/models/__init__.py
@@ -54,6 +54,7 @@
_import_structure["autoencoders.vq_model"] = ["VQModel"]
_import_structure["cache_utils"] = ["CacheMixin"]
_import_structure["controlnets.controlnet"] = ["ControlNetModel"]
+ _import_structure["controlnets.controlnet_cosmos"] = ["CosmosControlNetModel"]
_import_structure["controlnets.controlnet_flux"] = ["FluxControlNetModel", "FluxMultiControlNetModel"]
_import_structure["controlnets.controlnet_hunyuan"] = [
"HunyuanDiT2DControlNetModel",
@@ -175,6 +176,7 @@
ControlNetModel,
ControlNetUnionModel,
ControlNetXSAdapter,
+ CosmosControlNetModel,
FluxControlNetModel,
FluxMultiControlNetModel,
HunyuanDiT2DControlNetModel,
diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py
index db45159adfc9..ed966dc8fe98 100644
--- a/src/diffusers/models/_modeling_parallel.py
+++ b/src/diffusers/models/_modeling_parallel.py
@@ -16,7 +16,7 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Literal
import torch
import torch.distributed as dist
@@ -63,8 +63,8 @@ class ContextParallelConfig:
"""
- ring_degree: Optional[int] = None
- ulysses_degree: Optional[int] = None
+ ring_degree: int | None = None
+ ulysses_degree: int | None = None
convert_to_fp32: bool = True
# TODO: support alltoall
rotate_method: Literal["allgather", "alltoall"] = "allgather"
@@ -105,11 +105,11 @@ def __post_init__(self):
raise ValueError("ulysses_anything cannot be enabled when ring_degree > 1.")
@property
- def mesh_shape(self) -> Tuple[int, int]:
+ def mesh_shape(self) -> tuple[int, int]:
return (self.ring_degree, self.ulysses_degree)
@property
- def mesh_dim_names(self) -> Tuple[str, str]:
+ def mesh_dim_names(self) -> tuple[str, str]:
"""Dimension names for the device mesh."""
return ("ring", "ulysses")
@@ -141,7 +141,7 @@ class ParallelConfig:
Configuration for context parallelism.
"""
- context_parallel_config: Optional[ContextParallelConfig] = None
+ context_parallel_config: ContextParallelConfig | None = None
_rank: int = None
_world_size: int = None
@@ -154,7 +154,7 @@ def setup(
world_size: int,
device: torch.device,
*,
- mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None,
+ mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
):
self._rank = rank
self._world_size = world_size
@@ -182,7 +182,7 @@ class ContextParallelInput:
"""
split_dim: int
- expected_dims: Optional[int] = None
+ expected_dims: int | None = None
split_output: bool = False
def __repr__(self):
@@ -203,7 +203,7 @@ class ContextParallelOutput:
"""
gather_dim: int
- expected_dims: Optional[int] = None
+ expected_dims: int | None = None
def __repr__(self):
return f"ContextParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})"
@@ -214,19 +214,17 @@ def __repr__(self):
# If the key is a string, it denotes the name of the parameter in the forward function.
# If the key is an integer, split_output must be set to True, and it denotes the index of the output
# to be split across context parallel region.
-ContextParallelInputType = Dict[
- Union[str, int], Union[ContextParallelInput, List[ContextParallelInput], Tuple[ContextParallelInput, ...]]
+ContextParallelInputType = dict[
+ str | int, ContextParallelInput | list[ContextParallelInput] | tuple[ContextParallelInput, ...]
]
# A dictionary where keys denote the output to be gathered across context parallel region, and the
# value denotes the gathering configuration.
-ContextParallelOutputType = Union[
- ContextParallelOutput, List[ContextParallelOutput], Tuple[ContextParallelOutput, ...]
-]
+ContextParallelOutputType = ContextParallelOutput | list[ContextParallelOutput] | tuple[ContextParallelOutput, ...]
# A dictionary where keys denote the module id, and the value denotes how the inputs/outputs of
# the module should be split/gathered across context parallel region.
-ContextParallelModelPlan = Dict[str, Union[ContextParallelInputType, ContextParallelOutputType]]
+ContextParallelModelPlan = dict[str, ContextParallelInputType | ContextParallelOutputType]
# Example of a ContextParallelModelPlan (QwenImageTransformer2DModel):
@@ -269,9 +267,9 @@ def __repr__(self):
# Below are utility functions for distributed communication in context parallelism.
-def gather_size_by_comm(size: int, group: dist.ProcessGroup) -> List[int]:
+def gather_size_by_comm(size: int, group: dist.ProcessGroup) -> list[int]:
r"""Gather the local size from all ranks.
- size: int, local size return: List[int], list of size from all ranks
+ size: int, local size return: list[int], list of size from all ranks
"""
# NOTE(Serving/CP Safety):
# Do NOT cache this collective result.
diff --git a/src/diffusers/models/adapter.py b/src/diffusers/models/adapter.py
index e475fe6bee88..f0652c581a3e 100644
--- a/src/diffusers/models/adapter.py
+++ b/src/diffusers/models/adapter.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
-from typing import Callable, List, Optional, Union
+from typing import Callable
import torch
import torch.nn as nn
@@ -34,11 +34,11 @@ class MultiAdapter(ModelMixin):
or saving.
Args:
- adapters (`List[T2IAdapter]`, *optional*, defaults to None):
+ adapters (`list[T2IAdapter]`, *optional*, defaults to None):
A list of `T2IAdapter` model instances.
"""
- def __init__(self, adapters: List["T2IAdapter"]):
+ def __init__(self, adapters: list["T2IAdapter"]):
super(MultiAdapter, self).__init__()
self.num_adapter = len(adapters)
@@ -73,7 +73,7 @@ def __init__(self, adapters: List["T2IAdapter"]):
self.total_downscale_factor = first_adapter_total_downscale_factor
self.downscale_factor = first_adapter_downscale_factor
- def forward(self, xs: torch.Tensor, adapter_weights: Optional[List[float]] = None) -> List[torch.Tensor]:
+ def forward(self, xs: torch.Tensor, adapter_weights: list[float] | None = None) -> list[torch.Tensor]:
r"""
Args:
xs (`torch.Tensor`):
@@ -81,7 +81,7 @@ def forward(self, xs: torch.Tensor, adapter_weights: Optional[List[float]] = Non
models, concatenated along dimension 1(channel dimension). The `channel` dimension should be equal to
`num_adapter` * number of channel per image.
- adapter_weights (`List[float]`, *optional*, defaults to None):
+ adapter_weights (`list[float]`, *optional*, defaults to None):
A list of floats representing the weights which will be multiplied by each adapter's output before
summing them together. If `None`, equal weights will be used for all adapters.
"""
@@ -104,11 +104,11 @@ def forward(self, xs: torch.Tensor, adapter_weights: Optional[List[float]] = Non
def save_pretrained(
self,
- save_directory: Union[str, os.PathLike],
+ save_directory: str | os.PathLike,
is_main_process: bool = True,
save_function: Callable = None,
safe_serialization: bool = True,
- variant: Optional[str] = None,
+ variant: str | None = None,
):
"""
Save a model and its configuration file to a specified directory, allowing it to be re-loaded with the
@@ -145,7 +145,7 @@ def save_pretrained(
model_path_to_save = model_path_to_save + f"_{idx}"
@classmethod
- def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]], **kwargs):
+ def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwargs):
r"""
Instantiate a pretrained `MultiAdapter` model from multiple pre-trained adapter models.
@@ -165,7 +165,7 @@ def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]
Override the default `torch.dtype` and load the model under this dtype.
output_loading_info(`bool`, *optional*, defaults to `False`):
Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
- device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+ device_map (`str` or `dict[str, int | str | torch.device]`, *optional*):
A map that specifies where each submodule should go. It doesn't need to be refined to each
parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
same device.
@@ -229,7 +229,7 @@ class T2IAdapter(ModelMixin, ConfigMixin):
in_channels (`int`, *optional*, defaults to `3`):
The number of channels in the adapter's input (*control image*). Set it to 1 if you're using a gray scale
image.
- channels (`List[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ channels (`list[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The number of channels in each downsample block's output hidden state. The `len(block_out_channels)`
determines the number of downsample blocks in the adapter.
num_res_blocks (`int`, *optional*, defaults to `2`):
@@ -244,7 +244,7 @@ class T2IAdapter(ModelMixin, ConfigMixin):
def __init__(
self,
in_channels: int = 3,
- channels: List[int] = [320, 640, 1280, 1280],
+ channels: list[int] = [320, 640, 1280, 1280],
num_res_blocks: int = 2,
downscale_factor: int = 8,
adapter_type: str = "full_adapter",
@@ -263,7 +263,7 @@ def __init__(
"'full_adapter_xl' or 'light_adapter'."
)
- def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+ def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
r"""
This function processes the input tensor `x` through the adapter model and returns a list of feature tensors,
each representing information extracted at a different scale from the input. The length of the list is
@@ -295,7 +295,7 @@ class FullAdapter(nn.Module):
def __init__(
self,
in_channels: int = 3,
- channels: List[int] = [320, 640, 1280, 1280],
+ channels: list[int] = [320, 640, 1280, 1280],
num_res_blocks: int = 2,
downscale_factor: int = 8,
):
@@ -318,7 +318,7 @@ def __init__(
self.total_downscale_factor = downscale_factor * 2 ** (len(channels) - 1)
- def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+ def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
r"""
This method processes the input tensor `x` through the FullAdapter model and performs operations including
pixel unshuffling, convolution, and a stack of AdapterBlocks. It returns a list of feature tensors, each
@@ -345,7 +345,7 @@ class FullAdapterXL(nn.Module):
def __init__(
self,
in_channels: int = 3,
- channels: List[int] = [320, 640, 1280, 1280],
+ channels: list[int] = [320, 640, 1280, 1280],
num_res_blocks: int = 2,
downscale_factor: int = 16,
):
@@ -370,7 +370,7 @@ def __init__(
# XL has only one downsampling AdapterBlock.
self.total_downscale_factor = downscale_factor * 2
- def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+ def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
r"""
This method takes the tensor x as input and processes it through FullAdapterXL model. It consists of operations
including unshuffling pixels, applying convolution layer and appending each block into list of feature tensors.
@@ -473,7 +473,7 @@ class LightAdapter(nn.Module):
def __init__(
self,
in_channels: int = 3,
- channels: List[int] = [320, 640, 1280],
+ channels: list[int] = [320, 640, 1280],
num_res_blocks: int = 4,
downscale_factor: int = 8,
):
@@ -496,7 +496,7 @@ def __init__(
self.total_downscale_factor = downscale_factor * (2 ** len(channels))
- def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+ def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
r"""
This method takes the input tensor x and performs downscaling and appends it in list of feature tensors. Each
feature tensor corresponds to a different level of processing within the LightAdapter.
diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py
index 8b583d1a1cce..36d0893734c7 100644
--- a/src/diffusers/models/attention.py
+++ b/src/diffusers/models/attention.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
import torch.nn as nn
@@ -38,7 +38,7 @@
class AttentionMixin:
@property
- def attn_processors(self) -> Dict[str, AttentionProcessor]:
+ def attn_processors(self) -> dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
@@ -47,7 +47,7 @@ def attn_processors(self) -> Dict[str, AttentionProcessor]:
# set recursively
processors = {}
- def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict[str, AttentionProcessor]):
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
@@ -61,7 +61,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors:
return processors
- def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+ def set_attn_processor(self, processor: AttentionProcessor | dict[str, AttentionProcessor]):
r"""
Sets the attention processor to use to compute attention.
@@ -185,7 +185,7 @@ def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None:
def set_use_xla_flash_attention(
self,
use_xla_flash_attention: bool,
- partition_spec: Optional[Tuple[Optional[str], ...]] = None,
+ partition_spec: tuple[str | None, ...] | None = None,
is_flux=False,
) -> None:
"""
@@ -194,7 +194,7 @@ def set_use_xla_flash_attention(
Args:
use_xla_flash_attention (`bool`):
Whether to use pallas flash attention kernel from `torch_xla` or not.
- partition_spec (`Tuple[]`, *optional*):
+ partition_spec (`tuple[]`, *optional*):
Specify the partition specification if using SPMD. Otherwise None.
is_flux (`bool`, *optional*, defaults to `False`):
Whether the model is a Flux model.
@@ -206,7 +206,7 @@ def set_use_xla_flash_attention(
self.set_attention_backend("_native_xla")
def set_use_memory_efficient_attention_xformers(
- self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None
+ self, use_memory_efficient_attention_xformers: bool, attention_op: Callable | None = None
) -> None:
"""
Set whether to use memory efficient attention from `xformers` or not.
@@ -402,7 +402,7 @@ def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Ten
return tensor
def get_attention_scores(
- self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+ self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor | None = None
) -> torch.Tensor:
"""
Compute the attention scores.
@@ -597,7 +597,7 @@ def __init__(
num_attention_heads: int,
attention_head_dim: int,
context_pre_only: bool = False,
- qk_norm: Optional[str] = None,
+ qk_norm: str | None = None,
use_dual_attention: bool = False,
):
super().__init__()
@@ -673,7 +673,7 @@ def __init__(
self._chunk_dim = 0
# Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward
- def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+ def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0):
# Sets chunk feed-forward
self._chunk_size = chunk_size
self._chunk_dim = dim
@@ -683,8 +683,8 @@ def forward(
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor,
temb: torch.FloatTensor,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
joint_attention_kwargs = joint_attention_kwargs or {}
if self.use_dual_attention:
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1(
@@ -790,9 +790,9 @@ def __init__(
num_attention_heads: int,
attention_head_dim: int,
dropout=0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
attention_bias: bool = False,
only_cross_attention: bool = False,
double_self_attention: bool = False,
@@ -802,11 +802,11 @@ def __init__(
norm_eps: float = 1e-5,
final_dropout: bool = False,
attention_type: str = "default",
- positional_embeddings: Optional[str] = None,
- num_positional_embeddings: Optional[int] = None,
- ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,
- ada_norm_bias: Optional[int] = None,
- ff_inner_dim: Optional[int] = None,
+ positional_embeddings: str | None = None,
+ num_positional_embeddings: int | None = None,
+ ada_norm_continous_conditioning_embedding_dim: int | None = None,
+ ada_norm_bias: int | None = None,
+ ff_inner_dim: int | None = None,
ff_bias: bool = True,
attention_out_bias: bool = True,
):
@@ -952,7 +952,7 @@ def __init__(
self._chunk_size = None
self._chunk_dim = 0
- def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+ def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0):
# Sets chunk feed-forward
self._chunk_size = chunk_size
self._chunk_dim = dim
@@ -960,13 +960,13 @@ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
def forward(
self,
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- timestep: Optional[torch.LongTensor] = None,
- cross_attention_kwargs: Dict[str, Any] = None,
- class_labels: Optional[torch.LongTensor] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ timestep: torch.LongTensor | None = None,
+ cross_attention_kwargs: dict[str, Any] = None,
+ class_labels: torch.LongTensor | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
@@ -1099,8 +1099,8 @@ def __init__(
self,
dim: int,
inner_dim: int,
- multiple_of: Optional[int] = 256,
- ffn_dim_multiplier: Optional[float] = None,
+ multiple_of: int | None = 256,
+ ffn_dim_multiplier: float | None = None,
):
super().__init__()
# custom hidden_size factor multiplier
@@ -1148,7 +1148,7 @@ def __init__(
time_mix_inner_dim: int,
num_attention_heads: int,
attention_head_dim: int,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
):
super().__init__()
self.is_res = dim == time_mix_inner_dim
@@ -1195,7 +1195,7 @@ def __init__(
self._chunk_size = None
self._chunk_dim = None
- def set_chunk_feed_forward(self, chunk_size: Optional[int], **kwargs):
+ def set_chunk_feed_forward(self, chunk_size: int | None, **kwargs):
# Sets chunk feed-forward
self._chunk_size = chunk_size
# chunk dim should be hardcoded to 1 to have better speed vs. memory trade-off
@@ -1205,7 +1205,7 @@ def forward(
self,
hidden_states: torch.Tensor,
num_frames: int,
- encoder_hidden_states: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
) -> torch.Tensor:
# Notice that normalization is always applied before the real computation in the following blocks.
# 0. Self-Attention
@@ -1268,7 +1268,7 @@ def __init__(
kv_input_dim: int,
kv_input_dim_proj_use_bias: bool,
dropout=0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
attention_out_bias: bool = True,
):
@@ -1393,9 +1393,9 @@ def __init__(
num_attention_heads: int,
attention_head_dim: int,
dropout: float = 0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
attention_bias: bool = False,
only_cross_attention: bool = False,
double_self_attention: bool = False,
@@ -1404,9 +1404,9 @@ def __init__(
norm_type: str = "layer_norm",
norm_eps: float = 1e-5,
final_dropout: bool = False,
- positional_embeddings: Optional[str] = None,
- num_positional_embeddings: Optional[int] = None,
- ff_inner_dim: Optional[int] = None,
+ positional_embeddings: str | None = None,
+ num_positional_embeddings: int | None = None,
+ ff_inner_dim: int | None = None,
ff_bias: bool = True,
attention_out_bias: bool = True,
context_length: int = 16,
@@ -1501,7 +1501,7 @@ def __init__(
self._chunk_size = None
self._chunk_dim = 0
- def _get_frame_indices(self, num_frames: int) -> List[Tuple[int, int]]:
+ def _get_frame_indices(self, num_frames: int) -> list[tuple[int, int]]:
frame_indices = []
for i in range(0, num_frames - self.context_length + 1, self.context_stride):
window_start = i
@@ -1509,7 +1509,7 @@ def _get_frame_indices(self, num_frames: int) -> List[Tuple[int, int]]:
frame_indices.append((window_start, window_end))
return frame_indices
- def _get_frame_weights(self, num_frames: int, weighting_scheme: str = "pyramid") -> List[float]:
+ def _get_frame_weights(self, num_frames: int, weighting_scheme: str = "pyramid") -> list[float]:
if weighting_scheme == "flat":
weights = [1.0] * num_frames
@@ -1548,7 +1548,7 @@ def set_free_noise_properties(
self.context_stride = context_stride
self.weighting_scheme = weighting_scheme
- def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0) -> None:
+ def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0) -> None:
# Sets chunk feed-forward
self._chunk_size = chunk_size
self._chunk_dim = dim
@@ -1556,10 +1556,10 @@ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0) -> Non
def forward(
self,
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Dict[str, Any] = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -1696,7 +1696,7 @@ class FeedForward(nn.Module):
def __init__(
self,
dim: int,
- dim_out: Optional[int] = None,
+ dim_out: int | None = None,
mult: int = 4,
dropout: float = 0.0,
activation_fn: str = "geglu",
diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py
index a94c0f5a7d43..834bce942e43 100644
--- a/src/diffusers/models/attention_dispatch.py
+++ b/src/diffusers/models/attention_dispatch.py
@@ -12,13 +12,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import contextlib
import functools
import inspect
import math
from dataclasses import dataclass
from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable
import torch
import torch.distributed as dist
@@ -60,6 +62,8 @@
_REQUIRED_XLA_VERSION = "2.2"
_REQUIRED_XFORMERS_VERSION = "0.0.29"
+logger = get_logger(__name__) # pylint: disable=invalid-name
+
_CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION)
_CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available()
_CAN_USE_AITER_ATTN = is_aiter_available() and is_aiter_version(">=", _REQUIRED_AITER_VERSION)
@@ -71,8 +75,18 @@
if _CAN_USE_FLASH_ATTN:
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward
+ try:
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
+ from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward
+ except (ImportError, OSError, RuntimeError) as e:
+ # Handle ABI mismatch or other import failures gracefully.
+ # This can happen when flash_attn was compiled against a different PyTorch version.
+ logger.warning(f"flash_attn is installed but failed to import: {e}. Falling back to native PyTorch attention.")
+ _CAN_USE_FLASH_ATTN = False
+ flash_attn_func = None
+ flash_attn_varlen_func = None
+ _wrapped_flash_attn_backward = None
+ _wrapped_flash_attn_forward = None
else:
flash_attn_func = None
flash_attn_varlen_func = None
@@ -81,26 +95,47 @@
if _CAN_USE_FLASH_ATTN_3:
- from flash_attn_interface import flash_attn_func as flash_attn_3_func
- from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func
+ try:
+ from flash_attn_interface import flash_attn_func as flash_attn_3_func
+ from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func
+ except (ImportError, OSError, RuntimeError) as e:
+ logger.warning(f"flash_attn_3 failed to import: {e}. Falling back to native attention.")
+ _CAN_USE_FLASH_ATTN_3 = False
+ flash_attn_3_func = None
+ flash_attn_3_varlen_func = None
else:
flash_attn_3_func = None
flash_attn_3_varlen_func = None
if _CAN_USE_AITER_ATTN:
- from aiter import flash_attn_func as aiter_flash_attn_func
+ try:
+ from aiter import flash_attn_func as aiter_flash_attn_func
+ except (ImportError, OSError, RuntimeError) as e:
+ logger.warning(f"aiter failed to import: {e}. Falling back to native attention.")
+ _CAN_USE_AITER_ATTN = False
+ aiter_flash_attn_func = None
else:
aiter_flash_attn_func = None
if _CAN_USE_SAGE_ATTN:
- from sageattention import (
- sageattn,
- sageattn_qk_int8_pv_fp8_cuda,
- sageattn_qk_int8_pv_fp8_cuda_sm90,
- sageattn_qk_int8_pv_fp16_cuda,
- sageattn_qk_int8_pv_fp16_triton,
- sageattn_varlen,
- )
+ try:
+ from sageattention import (
+ sageattn,
+ sageattn_qk_int8_pv_fp8_cuda,
+ sageattn_qk_int8_pv_fp8_cuda_sm90,
+ sageattn_qk_int8_pv_fp16_cuda,
+ sageattn_qk_int8_pv_fp16_triton,
+ sageattn_varlen,
+ )
+ except (ImportError, OSError, RuntimeError) as e:
+ logger.warning(f"sageattention failed to import: {e}. Falling back to native attention.")
+ _CAN_USE_SAGE_ATTN = False
+ sageattn = None
+ sageattn_qk_int8_pv_fp8_cuda = None
+ sageattn_qk_int8_pv_fp8_cuda_sm90 = None
+ sageattn_qk_int8_pv_fp16_cuda = None
+ sageattn_qk_int8_pv_fp16_triton = None
+ sageattn_varlen = None
else:
sageattn = None
sageattn_qk_int8_pv_fp16_cuda = None
@@ -111,26 +146,48 @@
if _CAN_USE_FLEX_ATTN:
- # We cannot import the flex_attention function from the package directly because it is expected (from the
- # pytorch documentation) that the user may compile it. If we import directly, we will not have access to the
- # compiled function.
- import torch.nn.attention.flex_attention as flex_attention
+ try:
+ # We cannot import the flex_attention function from the package directly because it is expected (from the
+ # pytorch documentation) that the user may compile it. If we import directly, we will not have access to the
+ # compiled function.
+ import torch.nn.attention.flex_attention as flex_attention
+ except (ImportError, OSError, RuntimeError) as e:
+ logger.warning(f"flex_attention failed to import: {e}. Falling back to native attention.")
+ _CAN_USE_FLEX_ATTN = False
+ flex_attention = None
+else:
+ flex_attention = None
if _CAN_USE_NPU_ATTN:
- from torch_npu import npu_fusion_attention
+ try:
+ from torch_npu import npu_fusion_attention
+ except (ImportError, OSError, RuntimeError) as e:
+ logger.warning(f"torch_npu failed to import: {e}. Falling back to native attention.")
+ _CAN_USE_NPU_ATTN = False
+ npu_fusion_attention = None
else:
npu_fusion_attention = None
if _CAN_USE_XLA_ATTN:
- from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention
+ try:
+ from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention
+ except (ImportError, OSError, RuntimeError) as e:
+ logger.warning(f"torch_xla failed to import: {e}. Falling back to native attention.")
+ _CAN_USE_XLA_ATTN = False
+ xla_flash_attention = None
else:
xla_flash_attention = None
if _CAN_USE_XFORMERS_ATTN:
- import xformers.ops as xops
+ try:
+ import xformers.ops as xops
+ except (ImportError, OSError, RuntimeError) as e:
+ logger.warning(f"xformers failed to import: {e}. Falling back to native attention.")
+ _CAN_USE_XFORMERS_ATTN = False
+ xops = None
else:
xops = None
@@ -156,8 +213,6 @@ def wrap(func):
_register_fake = register_fake_no_op
-logger = get_logger(__name__) # pylint: disable=invalid-name
-
# TODO(aryan): Add support for the following:
# - Sage Attention++
# - block sparse, radial and other attention methods
@@ -219,7 +274,7 @@ class _AttentionBackendRegistry:
def register(
cls,
backend: AttentionBackendName,
- constraints: Optional[List[Callable]] = None,
+ constraints: list[Callable] | None = None,
supports_context_parallel: bool = False,
):
logger.debug(f"Registering attention backend: {backend} with constraints: {constraints}")
@@ -262,12 +317,16 @@ class _HubKernelConfig:
repo_id: str
function_attr: str
- revision: Optional[str] = None
- kernel_fn: Optional[Callable] = None
+ revision: str | None = None
+ kernel_fn: Callable | None = None
+ wrapped_forward_attr: str | None = None
+ wrapped_backward_attr: str | None = None
+ wrapped_forward_fn: Callable | None = None
+ wrapped_backward_fn: Callable | None = None
# Registry for hub-based attention kernels
-_HUB_KERNELS_REGISTRY: Dict["AttentionBackendName", _HubKernelConfig] = {
+_HUB_KERNELS_REGISTRY: dict["AttentionBackendName", _HubKernelConfig] = {
# TODO: temporary revision for now. Remove when merged upstream into `main`.
AttentionBackendName._FLASH_3_HUB: _HubKernelConfig(
repo_id="kernels-community/flash-attn3", function_attr="flash_attn_func", revision="fake-ops-return-probs"
@@ -278,7 +337,11 @@ class _HubKernelConfig:
# revision="fake-ops-return-probs",
),
AttentionBackendName.FLASH_HUB: _HubKernelConfig(
- repo_id="kernels-community/flash-attn2", function_attr="flash_attn_func", revision=None
+ repo_id="kernels-community/flash-attn2",
+ function_attr="flash_attn_func",
+ revision=None,
+ wrapped_forward_attr="flash_attn_interface._wrapped_flash_attn_forward",
+ wrapped_backward_attr="flash_attn_interface._wrapped_flash_attn_backward",
),
AttentionBackendName.FLASH_VARLEN_HUB: _HubKernelConfig(
repo_id="kernels-community/flash-attn2", function_attr="flash_attn_varlen_func", revision=None
@@ -290,7 +353,7 @@ class _HubKernelConfig:
@contextlib.contextmanager
-def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE):
+def attention_backend(backend: str | AttentionBackendName = AttentionBackendName.NATIVE):
"""
Context manager to set the active attention backend.
"""
@@ -314,15 +377,15 @@ def dispatch_attention_fn(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
*,
- backend: Optional[AttentionBackendName] = None,
- parallel_config: Optional["ParallelConfig"] = None,
+ backend: AttentionBackendName | None = None,
+ parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
attention_kwargs = attention_kwargs or {}
@@ -364,7 +427,7 @@ def dispatch_attention_fn(
# A list of very simple functions to catch common errors quickly when debugging.
-def _check_attn_mask_or_causal(attn_mask: Optional[torch.Tensor], is_causal: bool, **kwargs) -> None:
+def _check_attn_mask_or_causal(attn_mask: torch.Tensor | None, is_causal: bool, **kwargs) -> None:
if attn_mask is not None and is_causal:
raise ValueError("`is_causal` cannot be True when `attn_mask` is not None.")
@@ -410,7 +473,7 @@ def _check_shape(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
**kwargs,
) -> None:
# Expected shapes:
@@ -504,7 +567,7 @@ def _prepare_for_flash_attn_or_sage_varlen_without_mask(
batch_size: int,
seq_len_q: int,
seq_len_kv: int,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
seqlens_q = torch.full((batch_size,), seq_len_q, dtype=torch.int32, device=device)
seqlens_k = torch.full((batch_size,), seq_len_kv, dtype=torch.int32, device=device)
@@ -521,7 +584,7 @@ def _prepare_for_flash_attn_or_sage_varlen_with_mask(
batch_size: int,
seq_len_q: int,
attn_mask: torch.Tensor,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
seqlens_q = torch.full((batch_size,), seq_len_q, dtype=torch.int32, device=device)
seqlens_k = attn_mask.sum(dim=1, dtype=torch.int32)
@@ -538,8 +601,8 @@ def _prepare_for_flash_attn_or_sage_varlen(
batch_size: int,
seq_len_q: int,
seq_len_kv: int,
- attn_mask: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ attn_mask: torch.Tensor | None = None,
+ device: torch.device | None = None,
) -> None:
if attn_mask is None:
return _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, device)
@@ -603,22 +666,39 @@ def _flex_attention_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):
# ===== Helpers for downloading kernels =====
+def _resolve_kernel_attr(module, attr_path: str):
+ target = module
+ for attr in attr_path.split("."):
+ if not hasattr(target, attr):
+ raise AttributeError(f"Kernel module '{module.__name__}' does not define attribute path '{attr_path}'.")
+ target = getattr(target, attr)
+ return target
+
+
def _maybe_download_kernel_for_backend(backend: AttentionBackendName) -> None:
if backend not in _HUB_KERNELS_REGISTRY:
return
config = _HUB_KERNELS_REGISTRY[backend]
- if config.kernel_fn is not None:
+ needs_kernel = config.kernel_fn is None
+ needs_wrapped_forward = config.wrapped_forward_attr is not None and config.wrapped_forward_fn is None
+ needs_wrapped_backward = config.wrapped_backward_attr is not None and config.wrapped_backward_fn is None
+
+ if not (needs_kernel or needs_wrapped_forward or needs_wrapped_backward):
return
try:
from kernels import get_kernel
kernel_module = get_kernel(config.repo_id, revision=config.revision)
- kernel_func = getattr(kernel_module, config.function_attr)
+ if needs_kernel:
+ config.kernel_fn = _resolve_kernel_attr(kernel_module, config.function_attr)
+
+ if needs_wrapped_forward:
+ config.wrapped_forward_fn = _resolve_kernel_attr(kernel_module, config.wrapped_forward_attr)
- # Cache the downloaded kernel function in the config object
- config.kernel_fn = kernel_func
+ if needs_wrapped_backward:
+ config.wrapped_backward_fn = _resolve_kernel_attr(kernel_module, config.wrapped_backward_attr)
except Exception as e:
logger.error(f"An error occurred while fetching kernel '{config.repo_id}' from the Hub: {e}")
@@ -634,19 +714,19 @@ def _wrapped_flash_attn_3(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
- softmax_scale: Optional[float] = None,
+ softmax_scale: float | None = None,
causal: bool = False,
- qv: Optional[torch.Tensor] = None,
- q_descale: Optional[torch.Tensor] = None,
- k_descale: Optional[torch.Tensor] = None,
- v_descale: Optional[torch.Tensor] = None,
+ qv: torch.Tensor | None = None,
+ q_descale: torch.Tensor | None = None,
+ k_descale: torch.Tensor | None = None,
+ v_descale: torch.Tensor | None = None,
attention_chunk: int = 0,
softcap: float = 0.0,
num_splits: int = 1,
- pack_gqa: Optional[bool] = None,
+ pack_gqa: bool | None = None,
deterministic: bool = False,
sm_margin: int = 0,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
# Hardcoded for now because pytorch does not support tuple/int type hints
window_size = (-1, -1)
out, lse, *_ = flash_attn_3_func(
@@ -676,19 +756,19 @@ def _(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
- softmax_scale: Optional[float] = None,
+ softmax_scale: float | None = None,
causal: bool = False,
- qv: Optional[torch.Tensor] = None,
- q_descale: Optional[torch.Tensor] = None,
- k_descale: Optional[torch.Tensor] = None,
- v_descale: Optional[torch.Tensor] = None,
+ qv: torch.Tensor | None = None,
+ q_descale: torch.Tensor | None = None,
+ k_descale: torch.Tensor | None = None,
+ v_descale: torch.Tensor | None = None,
attention_chunk: int = 0,
softcap: float = 0.0,
num_splits: int = 1,
- pack_gqa: Optional[bool] = None,
+ pack_gqa: bool | None = None,
deterministic: bool = False,
sm_margin: int = 0,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
window_size = (-1, -1) # noqa: F841
# A lot of the parameters here are not yet used in any way within diffusers.
# We can safely ignore for now and keep the fake op shape propagation simple.
@@ -705,14 +785,14 @@ def _native_attention_forward_op(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
_save_ctx: bool = True,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
# Native attention does not return_lse
if return_lse:
@@ -788,14 +868,14 @@ def _cudnn_attention_forward_op(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
_save_ctx: bool = True,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
if enable_gqa:
raise ValueError("`enable_gqa` is not yet supported for cuDNN attention.")
@@ -885,14 +965,14 @@ def _native_flash_attention_forward_op(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
_save_ctx: bool = True,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
if enable_gqa:
raise ValueError("`enable_gqa` is not yet supported for native flash attention.")
@@ -974,14 +1054,14 @@ def _flash_attention_forward_op(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
_save_ctx: bool = True,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
if attn_mask is not None:
raise ValueError("`attn_mask` is not yet supported for flash-attn 2.")
@@ -1069,19 +1149,250 @@ def _flash_attention_backward_op(
return grad_query, grad_key, grad_value
+def _flash_attention_hub_forward_op(
+ ctx: torch.autograd.function.FunctionCtx,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_mask: torch.Tensor | None = None,
+ dropout_p: float = 0.0,
+ is_causal: bool = False,
+ scale: float | None = None,
+ enable_gqa: bool = False,
+ return_lse: bool = False,
+ _save_ctx: bool = True,
+ _parallel_config: "ParallelConfig" | None = None,
+):
+ if attn_mask is not None:
+ raise ValueError("`attn_mask` is not yet supported for flash-attn hub kernels.")
+ if enable_gqa:
+ raise ValueError("`enable_gqa` is not yet supported for flash-attn hub kernels.")
+
+ config = _HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_HUB]
+ wrapped_forward_fn = config.wrapped_forward_fn
+ wrapped_backward_fn = config.wrapped_backward_fn
+ if wrapped_forward_fn is None or wrapped_backward_fn is None:
+ raise RuntimeError(
+ "Flash attention hub kernels must expose `_wrapped_flash_attn_forward` and `_wrapped_flash_attn_backward` "
+ "for context parallel execution."
+ )
+
+ if scale is None:
+ scale = query.shape[-1] ** (-0.5)
+
+ window_size = (-1, -1)
+ softcap = 0.0
+ alibi_slopes = None
+ deterministic = False
+ grad_enabled = any(x.requires_grad for x in (query, key, value))
+
+ if grad_enabled or (_parallel_config is not None and _parallel_config.context_parallel_config._world_size > 1):
+ dropout_p = dropout_p if dropout_p > 0 else 1e-30
+
+ with torch.set_grad_enabled(grad_enabled):
+ out, lse, S_dmask, rng_state = wrapped_forward_fn(
+ query,
+ key,
+ value,
+ dropout_p,
+ scale,
+ is_causal,
+ window_size[0],
+ window_size[1],
+ softcap,
+ alibi_slopes,
+ return_lse,
+ )
+ lse = lse.permute(0, 2, 1).contiguous()
+
+ if _save_ctx:
+ ctx.save_for_backward(query, key, value, out, lse, rng_state)
+ ctx.dropout_p = dropout_p
+ ctx.scale = scale
+ ctx.is_causal = is_causal
+ ctx.window_size = window_size
+ ctx.softcap = softcap
+ ctx.alibi_slopes = alibi_slopes
+ ctx.deterministic = deterministic
+
+ return (out, lse) if return_lse else out
+
+
+def _flash_attention_hub_backward_op(
+ ctx: torch.autograd.function.FunctionCtx,
+ grad_out: torch.Tensor,
+ *args,
+ **kwargs,
+):
+ config = _HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_HUB]
+ wrapped_backward_fn = config.wrapped_backward_fn
+ if wrapped_backward_fn is None:
+ raise RuntimeError(
+ "Flash attention hub kernels must expose `_wrapped_flash_attn_backward` for context parallel execution."
+ )
+
+ query, key, value, out, lse, rng_state = ctx.saved_tensors
+ grad_query, grad_key, grad_value = torch.empty_like(query), torch.empty_like(key), torch.empty_like(value)
+
+ _ = wrapped_backward_fn(
+ grad_out,
+ query,
+ key,
+ value,
+ out,
+ lse,
+ grad_query,
+ grad_key,
+ grad_value,
+ ctx.dropout_p,
+ ctx.scale,
+ ctx.is_causal,
+ ctx.window_size[0],
+ ctx.window_size[1],
+ ctx.softcap,
+ ctx.alibi_slopes,
+ ctx.deterministic,
+ rng_state,
+ )
+
+ grad_query = grad_query[..., : grad_out.shape[-1]]
+ grad_key = grad_key[..., : grad_out.shape[-1]]
+ grad_value = grad_value[..., : grad_out.shape[-1]]
+
+ return grad_query, grad_key, grad_value
+
+
+def _flash_attention_3_hub_forward_op(
+ ctx: torch.autograd.function.FunctionCtx,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_mask: torch.Tensor | None = None,
+ dropout_p: float = 0.0,
+ is_causal: bool = False,
+ scale: float | None = None,
+ enable_gqa: bool = False,
+ return_lse: bool = False,
+ _save_ctx: bool = True,
+ _parallel_config: "ParallelConfig" | None = None,
+ *,
+ window_size: tuple[int, int] = (-1, -1),
+ softcap: float = 0.0,
+ num_splits: int = 1,
+ pack_gqa: bool | None = None,
+ deterministic: bool = False,
+ sm_margin: int = 0,
+):
+ if attn_mask is not None:
+ raise ValueError("`attn_mask` is not yet supported for flash-attn 3 hub kernels.")
+ if dropout_p != 0.0:
+ raise ValueError("`dropout_p` is not yet supported for flash-attn 3 hub kernels.")
+ if enable_gqa:
+ raise ValueError("`enable_gqa` is not yet supported for flash-attn 3 hub kernels.")
+
+ func = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_HUB].kernel_fn
+ out = func(
+ q=query,
+ k=key,
+ v=value,
+ softmax_scale=scale,
+ causal=is_causal,
+ qv=None,
+ q_descale=None,
+ k_descale=None,
+ v_descale=None,
+ window_size=window_size,
+ softcap=softcap,
+ num_splits=num_splits,
+ pack_gqa=pack_gqa,
+ deterministic=deterministic,
+ sm_margin=sm_margin,
+ return_attn_probs=return_lse,
+ )
+
+ lse = None
+ if return_lse:
+ out, lse = out
+ lse = lse.permute(0, 2, 1).contiguous()
+
+ if _save_ctx:
+ ctx.save_for_backward(query, key, value)
+ ctx.scale = scale
+ ctx.is_causal = is_causal
+ ctx._hub_kernel = func
+
+ return (out, lse) if return_lse else out
+
+
+def _flash_attention_3_hub_backward_op(
+ ctx: torch.autograd.function.FunctionCtx,
+ grad_out: torch.Tensor,
+ *args,
+ window_size: tuple[int, int] = (-1, -1),
+ softcap: float = 0.0,
+ num_splits: int = 1,
+ pack_gqa: bool | None = None,
+ deterministic: bool = False,
+ sm_margin: int = 0,
+):
+ query, key, value = ctx.saved_tensors
+ kernel_fn = ctx._hub_kernel
+ # NOTE: Unlike the FA2 hub kernel, the FA3 hub kernel does not expose separate wrapped forward/backward
+ # primitives (no `wrapped_forward_attr`/`wrapped_backward_attr` in its `_HubKernelConfig`). We
+ # therefore rerun the forward pass under `torch.enable_grad()` and differentiate through it with
+ # `torch.autograd.grad()`. This is a second forward pass during backward; it can be avoided once
+ # the FA3 hub exposes a dedicated fused backward kernel (analogous to `_wrapped_flash_attn_backward`
+ # in the FA2 hub), at which point this can be refactored to match `_flash_attention_hub_backward_op`.
+ with torch.enable_grad():
+ query_r = query.detach().requires_grad_(True)
+ key_r = key.detach().requires_grad_(True)
+ value_r = value.detach().requires_grad_(True)
+
+ out = kernel_fn(
+ q=query_r,
+ k=key_r,
+ v=value_r,
+ softmax_scale=ctx.scale,
+ causal=ctx.is_causal,
+ qv=None,
+ q_descale=None,
+ k_descale=None,
+ v_descale=None,
+ window_size=window_size,
+ softcap=softcap,
+ num_splits=num_splits,
+ pack_gqa=pack_gqa,
+ deterministic=deterministic,
+ sm_margin=sm_margin,
+ return_attn_probs=False,
+ )
+ if isinstance(out, tuple):
+ out = out[0]
+
+ grad_query, grad_key, grad_value = torch.autograd.grad(
+ out,
+ (query_r, key_r, value_r),
+ grad_out,
+ retain_graph=False,
+ allow_unused=False,
+ )
+
+ return grad_query, grad_key, grad_value
+
+
def _sage_attention_forward_op(
ctx: torch.autograd.function.FunctionCtx,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
_save_ctx: bool = True,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
if attn_mask is not None:
raise ValueError("`attn_mask` is not yet supported for Sage attention.")
@@ -1107,6 +1418,46 @@ def _sage_attention_forward_op(
return (out, lse) if return_lse else out
+def _sage_attention_hub_forward_op(
+ ctx: torch.autograd.function.FunctionCtx,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_mask: torch.Tensor | None = None,
+ dropout_p: float = 0.0,
+ is_causal: bool = False,
+ scale: float | None = None,
+ enable_gqa: bool = False,
+ return_lse: bool = False,
+ _save_ctx: bool = True,
+ _parallel_config: "ParallelConfig" | None = None,
+):
+ if attn_mask is not None:
+ raise ValueError("`attn_mask` is not yet supported for Sage attention.")
+ if dropout_p > 0.0:
+ raise ValueError("`dropout_p` is not yet supported for Sage attention.")
+ if enable_gqa:
+ raise ValueError("`enable_gqa` is not yet supported for Sage attention.")
+
+ func = _HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB].kernel_fn
+ out = func(
+ q=query,
+ k=key,
+ v=value,
+ tensor_layout="NHD",
+ is_causal=is_causal,
+ sm_scale=scale,
+ return_lse=return_lse,
+ )
+
+ lse = None
+ if return_lse:
+ out, lse, *_ = out
+ lse = lse.permute(0, 2, 1).contiguous()
+
+ return (out, lse) if return_lse else out
+
+
def _sage_attention_backward_op(
ctx: torch.autograd.function.FunctionCtx,
grad_out: torch.Tensor,
@@ -1115,28 +1466,51 @@ def _sage_attention_backward_op(
raise NotImplementedError("Backward pass is not implemented for Sage attention.")
+def _maybe_modify_attn_mask_npu(query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None = None):
+ # Skip Attention Mask if all values are 1, `None` mask can speedup the computation
+ if attn_mask is not None and torch.all(attn_mask != 0):
+ attn_mask = None
+
+ # Reshape Attention Mask: [batch_size, seq_len_k] -> [batch_size, 1, sqe_len_q, seq_len_k]
+ # https://www.hiascend.com/document/detail/zh/Pytorch/730/apiref/torchnpuCustomsapi/docs/context/torch_npu-npu_fusion_attention.md
+ if (
+ attn_mask is not None
+ and attn_mask.ndim == 2
+ and attn_mask.shape[0] == query.shape[0]
+ and attn_mask.shape[1] == key.shape[1]
+ ):
+ B, Sq, Skv = attn_mask.shape[0], query.shape[1], key.shape[1]
+ attn_mask = ~attn_mask.to(torch.bool)
+ attn_mask = attn_mask.unsqueeze(1).expand(B, Sq, Skv).unsqueeze(1).contiguous()
+
+ return attn_mask
+
+
def _npu_attention_forward_op(
ctx: torch.autograd.function.FunctionCtx,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
_save_ctx: bool = True,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
if return_lse:
raise ValueError("NPU attention backend does not support setting `return_lse=True`.")
+ attn_mask = _maybe_modify_attn_mask_npu(query, key, attn_mask)
+
out = npu_fusion_attention(
query,
key,
value,
query.size(2), # num_heads
+ atten_mask=attn_mask,
input_layout="BSND",
pse=None,
scale=1.0 / math.sqrt(query.shape[-1]) if scale is None else scale,
@@ -1284,9 +1658,9 @@ def backward(ctx, grad_outputs):
# Below are helper functions to handle abritrary head num and abritrary sequence length for Ulysses Anything Attention.
-def _maybe_pad_qkv_head(x: torch.Tensor, H: int, group: dist.ProcessGroup) -> Tuple[torch.Tensor, int]:
+def _maybe_pad_qkv_head(x: torch.Tensor, H: int, group: dist.ProcessGroup) -> tuple[torch.Tensor, int]:
r"""Maybe pad the head dimension to be divisible by world_size.
- x: torch.Tensor, shape (B, S_LOCAL, H, D) H: int, original global head num return: Tuple[torch.Tensor, int], padded
+ x: torch.Tensor, shape (B, S_LOCAL, H, D) H: int, original global head num return: tuple[torch.Tensor, int], padded
tensor (B, S_LOCAL, H + H_PAD, D) and H_PAD
"""
world_size = dist.get_world_size(group=group)
@@ -1314,9 +1688,9 @@ def _maybe_unpad_qkv_head(x: torch.Tensor, H_PAD: int, group: dist.ProcessGroup)
return x.contiguous()
-def _maybe_pad_o_head(x: torch.Tensor, H: int, group: dist.ProcessGroup) -> Tuple[torch.Tensor, int]:
+def _maybe_pad_o_head(x: torch.Tensor, H: int, group: dist.ProcessGroup) -> tuple[torch.Tensor, int]:
r"""Maybe pad the head dimension to be divisible by world_size.
- x: torch.Tensor, shape (B, S_GLOBAL, H_LOCAL, D) H: int, original global head num return: Tuple[torch.Tensor, int],
+ x: torch.Tensor, shape (B, S_GLOBAL, H_LOCAL, D) H: int, original global head num return: tuple[torch.Tensor, int],
padded tensor (B, S_GLOBAL, H_LOCAL + H_PAD, D) and H_PAD
"""
if H is None:
@@ -1438,15 +1812,15 @@ def forward(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor],
+ attn_mask: torch.Tensor | None,
dropout_p: float,
is_causal: bool,
- scale: Optional[float],
+ scale: float | None,
enable_gqa: bool,
return_lse: bool,
forward_op,
backward_op,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
ring_mesh = _parallel_config.context_parallel_config._ring_mesh
rank = _parallel_config.context_parallel_config._ring_local_rank
@@ -1565,15 +1939,15 @@ def forward(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor],
+ attn_mask: torch.Tensor | None,
dropout_p: float,
is_causal: bool,
- scale: Optional[float],
+ scale: float | None,
enable_gqa: bool,
return_lse: bool,
forward_op,
backward_op,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
ulysses_mesh = _parallel_config.context_parallel_config._ulysses_mesh
world_size = _parallel_config.context_parallel_config.ulysses_degree
@@ -1660,15 +2034,15 @@ def forward(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor],
+ attn_mask: torch.Tensor,
dropout_p: float,
is_causal: bool,
- scale: Optional[float],
+ scale: float,
enable_gqa: bool,
return_lse: bool,
forward_op,
backward_op,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
**kwargs,
):
ulysses_mesh = _parallel_config.context_parallel_config._ulysses_mesh
@@ -1733,15 +2107,15 @@ def _templated_unified_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor],
+ attn_mask: torch.Tensor,
dropout_p: float,
is_causal: bool,
- scale: Optional[float],
+ scale: float,
enable_gqa: bool,
return_lse: bool,
forward_op,
backward_op,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
scatter_idx: int = 2,
gather_idx: int = 1,
):
@@ -1795,16 +2169,16 @@ def _templated_context_parallel_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
*,
forward_op,
backward_op,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
):
if is_causal:
raise ValueError("Causal attention is not yet supported for templated attention.")
@@ -1893,12 +2267,12 @@ def _flash_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
lse = None
if attn_mask is not None:
@@ -1940,35 +2314,53 @@ def _flash_attention(
@_AttentionBackendRegistry.register(
AttentionBackendName.FLASH_HUB,
constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
- supports_context_parallel=False,
+ supports_context_parallel=True,
)
def _flash_attention_hub(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
lse = None
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for flash-attn 2.")
func = _HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_HUB].kernel_fn
- out = func(
- q=query,
- k=key,
- v=value,
- dropout_p=dropout_p,
- softmax_scale=scale,
- causal=is_causal,
- return_attn_probs=return_lse,
- )
- if return_lse:
- out, lse, *_ = out
+ if _parallel_config is None:
+ out = func(
+ q=query,
+ k=key,
+ v=value,
+ dropout_p=dropout_p,
+ softmax_scale=scale,
+ causal=is_causal,
+ return_attn_probs=return_lse,
+ )
+ if return_lse:
+ out, lse, *_ = out
+ else:
+ out = _templated_context_parallel_attention(
+ query,
+ key,
+ value,
+ None,
+ dropout_p,
+ is_causal,
+ scale,
+ False,
+ return_lse,
+ forward_op=_flash_attention_hub_forward_op,
+ backward_op=_flash_attention_hub_backward_op,
+ _parallel_config=_parallel_config,
+ )
+ if return_lse:
+ out, lse = out
return (out, lse) if return_lse else out
@@ -1982,12 +2374,12 @@ def _flash_varlen_attention_hub(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
- scale: Optional[float] = None,
+ scale: float | None = None,
is_causal: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
batch_size, seq_len_q, _, _ = query.shape
_, seq_len_kv, _, _ = key.shape
@@ -2038,12 +2430,12 @@ def _flash_varlen_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
- scale: Optional[float] = None,
+ scale: float | None = None,
is_causal: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
batch_size, seq_len_q, _, _ = query.shape
_, seq_len_kv, _, _ = key.shape
@@ -2093,11 +2485,11 @@ def _flash_attention_3(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
- scale: Optional[float] = None,
+ attn_mask: torch.Tensor | None = None,
+ scale: float | None = None,
is_causal: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for flash-attn 3.")
@@ -2115,48 +2507,83 @@ def _flash_attention_3(
@_AttentionBackendRegistry.register(
AttentionBackendName._FLASH_3_HUB,
constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
- supports_context_parallel=False,
+ supports_context_parallel=True,
)
def _flash_attention_3_hub(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
- scale: Optional[float] = None,
+ attn_mask: torch.Tensor | None = None,
+ scale: float | None = None,
is_causal: bool = False,
- window_size: Tuple[int, int] = (-1, -1),
+ window_size: tuple[int, int] = (-1, -1),
softcap: float = 0.0,
deterministic: bool = False,
return_attn_probs: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
- if _parallel_config:
- raise NotImplementedError(f"{AttentionBackendName._FLASH_3_HUB.value} is not implemented for parallelism yet.")
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for flash-attn 3.")
func = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_HUB].kernel_fn
- out = func(
- q=query,
- k=key,
- v=value,
- softmax_scale=scale,
- causal=is_causal,
- qv=None,
- q_descale=None,
- k_descale=None,
- v_descale=None,
+ if _parallel_config is None:
+ out = func(
+ q=query,
+ k=key,
+ v=value,
+ softmax_scale=scale,
+ causal=is_causal,
+ qv=None,
+ q_descale=None,
+ k_descale=None,
+ v_descale=None,
+ window_size=window_size,
+ softcap=softcap,
+ num_splits=1,
+ pack_gqa=None,
+ deterministic=deterministic,
+ sm_margin=0,
+ return_attn_probs=return_attn_probs,
+ )
+ return (out[0], out[1]) if return_attn_probs else out
+
+ forward_op = functools.partial(
+ _flash_attention_3_hub_forward_op,
+ window_size=window_size,
+ softcap=softcap,
+ num_splits=1,
+ pack_gqa=None,
+ deterministic=deterministic,
+ sm_margin=0,
+ )
+ backward_op = functools.partial(
+ _flash_attention_3_hub_backward_op,
window_size=window_size,
softcap=softcap,
num_splits=1,
pack_gqa=None,
deterministic=deterministic,
sm_margin=0,
- return_attn_probs=return_attn_probs,
)
- # When `return_attn_probs` is True, the above returns a tuple of
- # actual outputs and lse.
- return (out[0], out[1]) if return_attn_probs else out
+ out = _templated_context_parallel_attention(
+ query,
+ key,
+ value,
+ None,
+ 0.0,
+ is_causal,
+ scale,
+ False,
+ return_attn_probs,
+ forward_op=forward_op,
+ backward_op=backward_op,
+ _parallel_config=_parallel_config,
+ )
+ if return_attn_probs:
+ out, lse = out
+ return out, lse
+
+ return out
@_AttentionBackendRegistry.register(
@@ -2168,11 +2595,11 @@ def _flash_attention_3_varlen_hub(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
- scale: Optional[float] = None,
+ attn_mask: torch.Tensor | None = None,
+ scale: float | None = None,
is_causal: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
batch_size, seq_len_q, _, _ = query.shape
_, seq_len_kv, _, _ = key.shape
@@ -2221,11 +2648,11 @@ def _flash_varlen_attention_3(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
- scale: Optional[float] = None,
+ attn_mask: torch.Tensor | None = None,
+ scale: float | None = None,
is_causal: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
batch_size, seq_len_q, _, _ = query.shape
_, seq_len_kv, _, _ = key.shape
@@ -2273,12 +2700,12 @@ def _aiter_flash_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for aiter attention")
@@ -2318,12 +2745,12 @@ def _native_flex_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[Union[torch.Tensor, "flex_attention.BlockMask"]] = None,
+ attn_mask: torch.Tensor | "flex_attention.BlockMask" | None = None,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
# TODO: should we LRU cache the block mask creation?
score_mod = None
@@ -2419,13 +2846,13 @@ def _native_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if return_lse:
raise ValueError("Native attention backend does not support setting `return_lse=True`.")
@@ -2483,13 +2910,13 @@ def _native_cudnn_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
lse = None
if _parallel_config is None and not return_lse:
@@ -2535,13 +2962,13 @@ def _native_efficient_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if return_lse:
raise ValueError("Native efficient attention backend does not support setting `return_lse=True`.")
@@ -2570,13 +2997,13 @@ def _native_flash_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for aiter attention")
@@ -2625,13 +3052,13 @@ def _native_math_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if return_lse:
raise ValueError("Native math attention backend does not support setting `return_lse=True`.")
@@ -2660,22 +3087,23 @@ def _native_npu_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
- if attn_mask is not None:
- raise ValueError("`attn_mask` is not supported for NPU attention")
if return_lse:
raise ValueError("NPU attention backend does not support setting `return_lse=True`.")
if _parallel_config is None:
+ attn_mask = _maybe_modify_attn_mask_npu(query, key, attn_mask)
+
out = npu_fusion_attention(
query,
key,
value,
query.size(2), # num_heads
+ atten_mask=attn_mask,
input_layout="BSND",
pse=None,
scale=1.0 / math.sqrt(query.shape[-1]) if scale is None else scale,
@@ -2690,7 +3118,7 @@ def _native_npu_attention(
query,
key,
value,
- None,
+ attn_mask,
dropout_p,
None,
scale,
@@ -2712,10 +3140,10 @@ def _native_xla_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
is_causal: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for XLA attention")
@@ -2742,11 +3170,11 @@ def _sage_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for sage attention")
@@ -2787,17 +3215,17 @@ def _sage_attention(
@_AttentionBackendRegistry.register(
AttentionBackendName.SAGE_HUB,
constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_shape],
- supports_context_parallel=False,
+ supports_context_parallel=True,
)
def _sage_attention_hub(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for sage attention")
@@ -2815,6 +3243,23 @@ def _sage_attention_hub(
)
if return_lse:
out, lse, *_ = out
+ else:
+ out = _templated_context_parallel_attention(
+ query,
+ key,
+ value,
+ None,
+ 0.0,
+ is_causal,
+ scale,
+ False,
+ return_lse,
+ forward_op=_sage_attention_hub_forward_op,
+ backward_op=_sage_attention_backward_op,
+ _parallel_config=_parallel_config,
+ )
+ if return_lse:
+ out, lse = out
return (out, lse) if return_lse else out
@@ -2827,11 +3272,11 @@ def _sage_varlen_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if return_lse:
raise ValueError("Sage varlen backend does not support setting `return_lse=True`.")
@@ -2882,11 +3327,11 @@ def _sage_qk_int8_pv_fp8_cuda_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for sage attention")
@@ -2909,11 +3354,11 @@ def _sage_qk_int8_pv_fp8_cuda_sm90_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for sage attention")
@@ -2936,11 +3381,11 @@ def _sage_qk_int8_pv_fp16_cuda_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for sage attention")
@@ -2963,11 +3408,11 @@ def _sage_qk_int8_pv_fp16_triton_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if attn_mask is not None:
raise ValueError("`attn_mask` is not supported for sage attention")
@@ -2990,13 +3435,13 @@ def _xformers_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
+ attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
- scale: Optional[float] = None,
+ scale: float | None = None,
enable_gqa: bool = False,
return_lse: bool = False,
- _parallel_config: Optional["ParallelConfig"] = None,
+ _parallel_config: "ParallelConfig" | None = None,
) -> torch.Tensor:
if return_lse:
raise ValueError("xformers attention backend does not support setting `return_lse=True`.")
diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py
index 66455d733aee..f83a1753dac5 100755
--- a/src/diffusers/models/attention_processor.py
+++ b/src/diffusers/models/attention_processor.py
@@ -11,9 +11,11 @@
# 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.
+from __future__ import annotations
+
import inspect
import math
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable
import torch
import torch.nn.functional as F
@@ -105,21 +107,21 @@ class Attention(nn.Module):
def __init__(
self,
query_dim: int,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
heads: int = 8,
- kv_heads: Optional[int] = None,
+ kv_heads: int | None = None,
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
upcast_attention: bool = False,
upcast_softmax: bool = False,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
cross_attention_norm_num_groups: int = 32,
- qk_norm: Optional[str] = None,
- added_kv_proj_dim: Optional[int] = None,
- added_proj_bias: Optional[bool] = True,
- norm_num_groups: Optional[int] = None,
- spatial_norm_dim: Optional[int] = None,
+ qk_norm: str | None = None,
+ added_kv_proj_dim: int | None = None,
+ added_proj_bias: bool | None = True,
+ norm_num_groups: int | None = None,
+ spatial_norm_dim: int | None = None,
out_bias: bool = True,
scale_qk: bool = True,
only_cross_attention: bool = False,
@@ -127,7 +129,7 @@ def __init__(
rescale_output_factor: float = 1.0,
residual_connection: bool = False,
_from_deprecated_attn_block: bool = False,
- processor: Optional["AttnProcessor"] = None,
+ processor: "AttnProcessor" | None = None,
out_dim: int = None,
out_context_dim: int = None,
context_pre_only=None,
@@ -309,7 +311,7 @@ def __init__(
def set_use_xla_flash_attention(
self,
use_xla_flash_attention: bool,
- partition_spec: Optional[Tuple[Optional[str], ...]] = None,
+ partition_spec: tuple[str | None, ...] | None = None,
is_flux=False,
) -> None:
r"""
@@ -318,7 +320,7 @@ def set_use_xla_flash_attention(
Args:
use_xla_flash_attention (`bool`):
Whether to use pallas flash attention kernel from `torch_xla` or not.
- partition_spec (`Tuple[]`, *optional*):
+ partition_spec (`tuple[]`, *optional*):
Specify the partition specification if using SPMD. Otherwise None.
"""
if use_xla_flash_attention:
@@ -357,7 +359,7 @@ def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None:
self.set_processor(processor)
def set_use_memory_efficient_attention_xformers(
- self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None
+ self, use_memory_efficient_attention_xformers: bool, attention_op: Callable | None = None
) -> None:
r"""
Set whether to use memory efficient attention from `xformers` or not.
@@ -567,8 +569,8 @@ def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProce
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
**cross_attention_kwargs,
) -> torch.Tensor:
r"""
@@ -655,7 +657,7 @@ def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Ten
return tensor
def get_attention_scores(
- self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+ self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor | None = None
) -> torch.Tensor:
r"""
Compute the attention scores.
@@ -868,11 +870,11 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- num_attention_heads: Optional[int] = None,
+ num_attention_heads: int | None = None,
attention_head_dim: int = 8,
mult: float = 1.0,
norm_type: str = "batch_norm",
- kernel_sizes: Tuple[int, ...] = (5,),
+ kernel_sizes: tuple[int, ...] = (5,),
eps: float = 1e-15,
residual_connection: bool = False,
):
@@ -938,8 +940,8 @@ def __init__(
dropout: float = 0.0,
bias: bool = False,
added_proj_bias: bool = True,
- out_dim: Optional[int] = None,
- out_context_dim: Optional[int] = None,
+ out_dim: int | None = None,
+ out_context_dim: int | None = None,
out_bias: bool = True,
context_pre_only: bool = False,
eps: float = 1e-5,
@@ -980,8 +982,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
**kwargs,
):
return self.processor(
@@ -1006,7 +1008,7 @@ def __call__(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
@@ -1107,9 +1109,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -1193,8 +1195,8 @@ def __init__(
self,
train_kv: bool = True,
train_q_out: bool = True,
- hidden_size: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
+ hidden_size: int | None = None,
+ cross_attention_dim: int | None = None,
out_bias: bool = True,
dropout: float = 0.0,
):
@@ -1219,8 +1221,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = hidden_states.shape
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
@@ -1284,8 +1286,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -1357,8 +1359,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -1429,7 +1431,7 @@ def __call__(
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
- attention_mask: Optional[torch.FloatTensor] = None,
+ attention_mask: torch.FloatTensor | None = None,
*args,
**kwargs,
) -> torch.FloatTensor:
@@ -1517,7 +1519,7 @@ def __call__(
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
- attention_mask: Optional[torch.FloatTensor] = None,
+ attention_mask: torch.FloatTensor | None = None,
) -> torch.FloatTensor:
residual = hidden_states
@@ -1673,7 +1675,7 @@ def __call__(
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
- attention_mask: Optional[torch.FloatTensor] = None,
+ attention_mask: torch.FloatTensor | None = None,
*args,
**kwargs,
) -> torch.FloatTensor:
@@ -1836,7 +1838,7 @@ def __call__(
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
- attention_mask: Optional[torch.FloatTensor] = None,
+ attention_mask: torch.FloatTensor | None = None,
*args,
**kwargs,
) -> torch.FloatTensor:
@@ -1915,7 +1917,7 @@ class XFormersJointAttnProcessor:
operator.
"""
- def __init__(self, attention_op: Optional[Callable] = None):
+ def __init__(self, attention_op: Callable | None = None):
self.attention_op = attention_op
def __call__(
@@ -1923,7 +1925,7 @@ def __call__(
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
- attention_mask: Optional[torch.FloatTensor] = None,
+ attention_mask: torch.FloatTensor | None = None,
*args,
**kwargs,
) -> torch.FloatTensor:
@@ -2004,10 +2006,10 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
residual = hidden_states
@@ -2287,8 +2289,8 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
text_seq_length = encoder_hidden_states.size(1)
@@ -2356,8 +2358,8 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
text_seq_length = encoder_hidden_states.size(1)
@@ -2424,15 +2426,15 @@ class XFormersAttnAddedKVProcessor:
operator.
"""
- def __init__(self, attention_op: Optional[Callable] = None):
+ def __init__(self, attention_op: Callable | None = None):
self.attention_op = attention_op
def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
residual = hidden_states
hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)
@@ -2495,16 +2497,16 @@ class XFormersAttnProcessor:
operator.
"""
- def __init__(self, attention_op: Optional[Callable] = None):
+ def __init__(self, attention_op: Callable | None = None):
self.attention_op = attention_op
def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -2593,9 +2595,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -2704,9 +2706,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -2790,7 +2792,7 @@ class XLAFlashAttnProcessor2_0:
Processor for implementing scaled dot-product attention with pallas flash attention kernel if using `torch_xla`.
"""
- def __init__(self, partition_spec: Optional[Tuple[Optional[str], ...]] = None):
+ def __init__(self, partition_spec: tuple[str | None, ...] | None = None):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError(
"XLAFlashAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
@@ -2805,9 +2807,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -2914,8 +2916,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
residual = hidden_states
is_single_frame = hidden_states.shape[1] == 1
@@ -3001,7 +3003,7 @@ def __init__(self):
def apply_partial_rotary_emb(
self,
x: torch.Tensor,
- freqs_cis: Tuple[torch.Tensor],
+ freqs_cis: tuple[torch.Tensor],
) -> torch.Tensor:
from .embeddings import apply_rotary_emb
@@ -3017,9 +3019,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
from .embeddings import apply_rotary_emb
@@ -3133,10 +3135,10 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
from .embeddings import apply_rotary_emb
@@ -3234,10 +3236,10 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
from .embeddings import apply_rotary_emb
@@ -3337,10 +3339,10 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
from .embeddings import apply_rotary_emb
@@ -3460,10 +3462,10 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
from .embeddings import apply_rotary_emb
@@ -3582,10 +3584,10 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- query_rotary_emb: Optional[torch.Tensor] = None,
- key_rotary_emb: Optional[torch.Tensor] = None,
- base_sequence_length: Optional[int] = None,
+ attention_mask: torch.Tensor | None = None,
+ query_rotary_emb: torch.Tensor | None = None,
+ key_rotary_emb: torch.Tensor | None = None,
+ base_sequence_length: int | None = None,
) -> torch.Tensor:
from .embeddings import apply_rotary_emb
@@ -3682,9 +3684,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -3792,11 +3794,11 @@ def __init__(
self,
train_kv: bool = True,
train_q_out: bool = False,
- hidden_size: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
+ hidden_size: int | None = None,
+ cross_attention_dim: int | None = None,
out_bias: bool = True,
dropout: float = 0.0,
- attention_op: Optional[Callable] = None,
+ attention_op: Callable | None = None,
):
super().__init__()
self.train_kv = train_kv
@@ -3820,8 +3822,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
@@ -3905,8 +3907,8 @@ def __init__(
self,
train_kv: bool = True,
train_q_out: bool = True,
- hidden_size: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
+ hidden_size: int | None = None,
+ cross_attention_dim: int | None = None,
out_bias: bool = True,
dropout: float = 0.0,
):
@@ -3931,8 +3933,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = hidden_states.shape
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
@@ -4012,8 +4014,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
residual = hidden_states
@@ -4099,9 +4101,9 @@ def __call__(
self,
attn: "Attention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
) -> torch.Tensor:
residual = hidden_states
@@ -4212,9 +4214,9 @@ class IPAdapterAttnProcessor(nn.Module):
The hidden size of the attention layer.
cross_attention_dim (`int`):
The number of channels in the `encoder_hidden_states`.
- num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`):
+ num_tokens (`int`, `tuple[int]` or `list[int]`, defaults to `(4,)`):
The context length of the image features.
- scale (`float` or List[`float`], defaults to 1.0):
+ scale (`float` or list[`float`], defaults to 1.0):
the weight scale of image prompt.
"""
@@ -4245,11 +4247,11 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
scale: float = 1.0,
- ip_adapter_masks: Optional[torch.Tensor] = None,
+ ip_adapter_masks: torch.Tensor | None = None,
):
residual = hidden_states
@@ -4305,7 +4307,7 @@ def __call__(
hidden_states = attn.batch_to_head_dim(hidden_states)
if ip_adapter_masks is not None:
- if not isinstance(ip_adapter_masks, List):
+ if not isinstance(ip_adapter_masks, list):
# for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width]
ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1))
if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)):
@@ -4412,9 +4414,9 @@ class IPAdapterAttnProcessor2_0(torch.nn.Module):
The hidden size of the attention layer.
cross_attention_dim (`int`):
The number of channels in the `encoder_hidden_states`.
- num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`):
+ num_tokens (`int`, `tuple[int]` or `list[int]`, defaults to `(4,)`):
The context length of the image features.
- scale (`float` or `List[float]`, defaults to 1.0):
+ scale (`float` or `list[float]`, defaults to 1.0):
the weight scale of image prompt.
"""
@@ -4450,11 +4452,11 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
scale: float = 1.0,
- ip_adapter_masks: Optional[torch.Tensor] = None,
+ ip_adapter_masks: torch.Tensor | None = None,
):
residual = hidden_states
@@ -4524,7 +4526,7 @@ def __call__(
hidden_states = hidden_states.to(query.dtype)
if ip_adapter_masks is not None:
- if not isinstance(ip_adapter_masks, List):
+ if not isinstance(ip_adapter_masks, list):
# for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width]
ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1))
if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)):
@@ -4644,9 +4646,9 @@ class IPAdapterXFormersAttnProcessor(torch.nn.Module):
The hidden size of the attention layer.
cross_attention_dim (`int`):
The number of channels in the `encoder_hidden_states`.
- num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`):
+ num_tokens (`int`, `tuple[int]` or `list[int]`, defaults to `(4,)`):
The context length of the image features.
- scale (`float` or `List[float]`, defaults to 1.0):
+ scale (`float` or `list[float]`, defaults to 1.0):
the weight scale of image prompt.
attention_op (`Callable`, *optional*, defaults to `None`):
The base
@@ -4661,7 +4663,7 @@ def __init__(
cross_attention_dim=None,
num_tokens=(4,),
scale=1.0,
- attention_op: Optional[Callable] = None,
+ attention_op: Callable | None = None,
):
super().__init__()
@@ -4690,11 +4692,11 @@ 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,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ temb: torch.FloatTensor | None = None,
scale: float = 1.0,
- ip_adapter_masks: Optional[torch.FloatTensor] = None,
+ ip_adapter_masks: torch.FloatTensor | None = None,
):
residual = hidden_states
@@ -4763,7 +4765,7 @@ def __call__(
if ip_hidden_states:
if ip_adapter_masks is not None:
- if not isinstance(ip_adapter_masks, List):
+ if not isinstance(ip_adapter_masks, list):
# for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width]
ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1))
if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)):
@@ -4911,7 +4913,7 @@ def __call__(
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
- attention_mask: Optional[torch.FloatTensor] = None,
+ attention_mask: torch.FloatTensor | None = None,
ip_hidden_states: torch.FloatTensor = None,
temb: torch.FloatTensor = None,
) -> torch.FloatTensor:
@@ -5054,9 +5056,9 @@ 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,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ temb: torch.FloatTensor | None = None,
) -> torch.Tensor:
residual = hidden_states
if attn.spatial_norm is not None:
@@ -5153,9 +5155,9 @@ 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,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ temb: torch.FloatTensor | None = None,
) -> torch.Tensor:
residual = hidden_states
if attn.spatial_norm is not None:
@@ -5345,8 +5347,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
original_dtype = hidden_states.dtype
@@ -5397,8 +5399,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
original_dtype = hidden_states.dtype
@@ -5452,8 +5454,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
original_dtype = hidden_states.dtype
@@ -5622,56 +5624,56 @@ def __new__(cls, *args, **kwargs):
FluxIPAdapterJointAttnProcessor2_0,
)
-AttentionProcessor = Union[
- AttnProcessor,
- CustomDiffusionAttnProcessor,
- AttnAddedKVProcessor,
- AttnAddedKVProcessor2_0,
- JointAttnProcessor2_0,
- PAGJointAttnProcessor2_0,
- PAGCFGJointAttnProcessor2_0,
- FusedJointAttnProcessor2_0,
- AllegroAttnProcessor2_0,
- AuraFlowAttnProcessor2_0,
- FusedAuraFlowAttnProcessor2_0,
- FluxAttnProcessor2_0,
- FluxAttnProcessor2_0_NPU,
- FusedFluxAttnProcessor2_0,
- FusedFluxAttnProcessor2_0_NPU,
- CogVideoXAttnProcessor2_0,
- FusedCogVideoXAttnProcessor2_0,
- XFormersAttnAddedKVProcessor,
- XFormersAttnProcessor,
- XLAFlashAttnProcessor2_0,
- AttnProcessorNPU,
- AttnProcessor2_0,
- MochiVaeAttnProcessor2_0,
- MochiAttnProcessor2_0,
- StableAudioAttnProcessor2_0,
- HunyuanAttnProcessor2_0,
- FusedHunyuanAttnProcessor2_0,
- PAGHunyuanAttnProcessor2_0,
- PAGCFGHunyuanAttnProcessor2_0,
- LuminaAttnProcessor2_0,
- FusedAttnProcessor2_0,
- CustomDiffusionXFormersAttnProcessor,
- CustomDiffusionAttnProcessor2_0,
- SlicedAttnProcessor,
- SlicedAttnAddedKVProcessor,
- SanaLinearAttnProcessor2_0,
- PAGCFGSanaLinearAttnProcessor2_0,
- PAGIdentitySanaLinearAttnProcessor2_0,
- SanaMultiscaleLinearAttention,
- SanaMultiscaleAttnProcessor2_0,
- SanaMultiscaleAttentionProjection,
- IPAdapterAttnProcessor,
- IPAdapterAttnProcessor2_0,
- IPAdapterXFormersAttnProcessor,
- SD3IPAdapterJointAttnProcessor2_0,
- PAGIdentitySelfAttnProcessor2_0,
- PAGCFGIdentitySelfAttnProcessor2_0,
- LoRAAttnProcessor,
- LoRAAttnProcessor2_0,
- LoRAXFormersAttnProcessor,
- LoRAAttnAddedKVProcessor,
-]
+AttentionProcessor = (
+ AttnProcessor
+ | CustomDiffusionAttnProcessor
+ | AttnAddedKVProcessor
+ | AttnAddedKVProcessor2_0
+ | JointAttnProcessor2_0
+ | PAGJointAttnProcessor2_0
+ | PAGCFGJointAttnProcessor2_0
+ | FusedJointAttnProcessor2_0
+ | AllegroAttnProcessor2_0
+ | AuraFlowAttnProcessor2_0
+ | FusedAuraFlowAttnProcessor2_0
+ | FluxAttnProcessor2_0
+ | FluxAttnProcessor2_0_NPU
+ | FusedFluxAttnProcessor2_0
+ | FusedFluxAttnProcessor2_0_NPU
+ | CogVideoXAttnProcessor2_0
+ | FusedCogVideoXAttnProcessor2_0
+ | XFormersAttnAddedKVProcessor
+ | XFormersAttnProcessor
+ | XLAFlashAttnProcessor2_0
+ | AttnProcessorNPU
+ | AttnProcessor2_0
+ | MochiVaeAttnProcessor2_0
+ | MochiAttnProcessor2_0
+ | StableAudioAttnProcessor2_0
+ | HunyuanAttnProcessor2_0
+ | FusedHunyuanAttnProcessor2_0
+ | PAGHunyuanAttnProcessor2_0
+ | PAGCFGHunyuanAttnProcessor2_0
+ | LuminaAttnProcessor2_0
+ | FusedAttnProcessor2_0
+ | CustomDiffusionXFormersAttnProcessor
+ | CustomDiffusionAttnProcessor2_0
+ | SlicedAttnProcessor
+ | SlicedAttnAddedKVProcessor
+ | SanaLinearAttnProcessor2_0
+ | PAGCFGSanaLinearAttnProcessor2_0
+ | PAGIdentitySanaLinearAttnProcessor2_0
+ | SanaMultiscaleLinearAttention
+ | SanaMultiscaleAttnProcessor2_0
+ | SanaMultiscaleAttentionProjection
+ | IPAdapterAttnProcessor
+ | IPAdapterAttnProcessor2_0
+ | IPAdapterXFormersAttnProcessor
+ | SD3IPAdapterJointAttnProcessor2_0
+ | PAGIdentitySelfAttnProcessor2_0
+ | PAGCFGIdentitySelfAttnProcessor2_0
+ | LoRAAttnProcessor
+ | LoRAAttnProcessor2_0
+ | LoRAXFormersAttnProcessor
+ | LoRAAttnAddedKVProcessor
+)
diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py
index 0a5b7fff1c66..1c001e23fe00 100644
--- a/src/diffusers/models/auto_model.py
+++ b/src/diffusers/models/auto_model.py
@@ -13,7 +13,6 @@
# limitations under the License.
import os
-from typing import Optional, Union
from huggingface_hub.utils import validate_hf_hub_args
@@ -31,13 +30,129 @@ class AutoModel(ConfigMixin):
def __init__(self, *args, **kwargs):
raise EnvironmentError(
f"{self.__class__.__name__} is designed to be instantiated "
- f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)` or "
+ f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)`, "
+ f"`{self.__class__.__name__}.from_config(config)`, or "
f"`{self.__class__.__name__}.from_pipe(pipeline)` methods."
)
+ @classmethod
+ def from_config(cls, pretrained_model_name_or_path_or_dict: str | os.PathLike | dict | None = None, **kwargs):
+ r"""
+ Instantiate a model from a config dictionary or a pretrained model configuration file with random weights (no
+ pretrained weights are loaded).
+
+ Parameters:
+ pretrained_model_name_or_path_or_dict (`str`, `os.PathLike`, or `dict`):
+ Can be either:
+
+ - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model
+ configuration hosted on the Hub.
+ - A path to a *directory* (for example `./my_model_directory`) containing a model configuration
+ file.
+ - A config dictionary.
+
+ cache_dir (`Union[str, os.PathLike]`, *optional*):
+ Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
+ is not used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force the (re-)download of the model configuration, overriding the cached version if
+ it exists.
+ proxies (`Dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint.
+ local_files_only(`bool`, *optional*, defaults to `False`):
+ Whether to only load local model configuration files or not.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files.
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether to trust remote code.
+ subfolder (`str`, *optional*, defaults to `""`):
+ The subfolder location of a model file within a larger model repository on the Hub or locally.
+
+ Returns:
+ A model object instantiated from the config with random weights.
+
+ Example:
+
+ ```py
+ from diffusers import AutoModel
+
+ model = AutoModel.from_config("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet")
+ ```
+ """
+ subfolder = kwargs.pop("subfolder", None)
+ trust_remote_code = kwargs.pop("trust_remote_code", False)
+
+ hub_kwargs_names = [
+ "cache_dir",
+ "force_download",
+ "local_files_only",
+ "proxies",
+ "revision",
+ "token",
+ ]
+ hub_kwargs = {name: kwargs.pop(name, None) for name in hub_kwargs_names}
+
+ if pretrained_model_name_or_path_or_dict is None:
+ raise ValueError(
+ "Please provide a `pretrained_model_name_or_path_or_dict` as the first positional argument."
+ )
+
+ if isinstance(pretrained_model_name_or_path_or_dict, (str, os.PathLike)):
+ pretrained_model_name_or_path = pretrained_model_name_or_path_or_dict
+ config = cls.load_config(pretrained_model_name_or_path, subfolder=subfolder, **hub_kwargs)
+ else:
+ config = pretrained_model_name_or_path_or_dict
+ pretrained_model_name_or_path = config.get("_name_or_path", None)
+
+ has_remote_code = "auto_map" in config and cls.__name__ in config["auto_map"]
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_remote_code
+ )
+
+ if has_remote_code and trust_remote_code:
+ class_ref = config["auto_map"][cls.__name__]
+ module_file, class_name = class_ref.split(".")
+ module_file = module_file + ".py"
+ model_cls = get_class_from_dynamic_module(
+ pretrained_model_name_or_path,
+ subfolder=subfolder,
+ module_file=module_file,
+ class_name=class_name,
+ **hub_kwargs,
+ )
+ else:
+ if "_class_name" in config:
+ class_name = config["_class_name"]
+ library = "diffusers"
+ elif "model_type" in config:
+ class_name = "AutoModel"
+ library = "transformers"
+ else:
+ raise ValueError(
+ f"Couldn't find a model class associated with the config: {config}. Make sure the config "
+ "contains a `_class_name` or `model_type` key."
+ )
+
+ from ..pipelines.pipeline_loading_utils import ALL_IMPORTABLE_CLASSES, get_class_obj_and_candidates
+
+ model_cls, _ = get_class_obj_and_candidates(
+ library_name=library,
+ class_name=class_name,
+ importable_classes=ALL_IMPORTABLE_CLASSES,
+ pipelines=None,
+ is_pipeline_module=False,
+ )
+
+ if model_cls is None:
+ raise ValueError(f"AutoModel can't find a model linked to {class_name}.")
+
+ return model_cls.from_config(config, **kwargs)
+
@classmethod
@validate_hf_hub_args
- def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLike]] = None, **kwargs):
+ def from_pretrained(cls, pretrained_model_or_path: str | os.PathLike | None = None, **kwargs):
r"""
Instantiate a pretrained PyTorch model from a pretrained model configuration.
@@ -53,7 +168,7 @@ def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLi
- A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
with [`~ModelMixin.save_pretrained`].
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
torch_dtype (`torch.dtype`, *optional*):
@@ -61,7 +176,7 @@ def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLi
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info (`bool`, *optional*, defaults to `False`):
@@ -83,7 +198,7 @@ def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLi
Mirror source to resolve accessibility issues if you're downloading a model in China. We do not
guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
information.
- device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+ device_map (`str` or `dict[str, int | str | torch.device]`, *optional*):
A map that specifies where each submodule should go. It doesn't need to be defined for each
parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the
same device. Defaults to `None`, meaning that the model will be loaded on CPU.
diff --git a/src/diffusers/models/autoencoders/autoencoder_asym_kl.py b/src/diffusers/models/autoencoders/autoencoder_asym_kl.py
index fa49fcfe79f8..fbd9b3e459f7 100644
--- a/src/diffusers/models/autoencoders/autoencoder_asym_kl.py
+++ b/src/diffusers/models/autoencoders/autoencoder_asym_kl.py
@@ -11,8 +11,6 @@
# 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.
-from typing import Optional, Tuple, Union
-
import torch
import torch.nn as nn
@@ -34,16 +32,16 @@ class AsymmetricAutoencoderKL(ModelMixin, AutoencoderMixin, ConfigMixin):
Parameters:
in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
out_channels (int, *optional*, defaults to 3): Number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
- Tuple of downsample block types.
- down_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
- Tuple of down block output channels.
+ down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ tuple of downsample block types.
+ down_block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
+ tuple of down block output channels.
layers_per_down_block (`int`, *optional*, defaults to `1`):
Number layers for down block.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
- Tuple of upsample block types.
- up_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
- Tuple of up block output channels.
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ tuple of upsample block types.
+ up_block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
+ tuple of up block output channels.
layers_per_up_block (`int`, *optional*, defaults to `1`):
Number layers for up block.
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
@@ -67,11 +65,11 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",),
- down_block_out_channels: Tuple[int, ...] = (64,),
+ down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",),
+ down_block_out_channels: tuple[int, ...] = (64,),
layers_per_down_block: int = 1,
- up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
- up_block_out_channels: Tuple[int, ...] = (64,),
+ up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",),
+ up_block_out_channels: tuple[int, ...] = (64,),
layers_per_up_block: int = 1,
act_fn: str = "silu",
latent_channels: int = 4,
@@ -111,7 +109,7 @@ def __init__(
self.register_to_config(force_upcast=False)
@apply_forward_hook
- def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[AutoencoderKLOutput, Tuple[torch.Tensor]]:
+ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput | tuple[torch.Tensor]:
h = self.encoder(x)
moments = self.quant_conv(h)
posterior = DiagonalGaussianDistribution(moments)
@@ -124,10 +122,10 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[Autoencoder
def _decode(
self,
z: torch.Tensor,
- image: Optional[torch.Tensor] = None,
- mask: Optional[torch.Tensor] = None,
+ image: torch.Tensor | None = None,
+ mask: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ ) -> DecoderOutput | tuple[torch.Tensor]:
z = self.post_quant_conv(z)
dec = self.decoder(z, image, mask)
@@ -140,11 +138,11 @@ def _decode(
def decode(
self,
z: torch.Tensor,
- generator: Optional[torch.Generator] = None,
- image: Optional[torch.Tensor] = None,
- mask: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ image: torch.Tensor | None = None,
+ mask: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ ) -> DecoderOutput | tuple[torch.Tensor]:
decoded = self._decode(z, image, mask).sample
if not return_dict:
@@ -155,11 +153,11 @@ def decode(
def forward(
self,
sample: torch.Tensor,
- mask: Optional[torch.Tensor] = None,
+ mask: torch.Tensor | None = None,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | tuple[torch.Tensor]:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_dc.py b/src/diffusers/models/autoencoders/autoencoder_dc.py
index 19b666fdc4a8..02a83d79aba5 100644
--- a/src/diffusers/models/autoencoders/autoencoder_dc.py
+++ b/src/diffusers/models/autoencoders/autoencoder_dc.py
@@ -13,7 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -68,7 +67,7 @@ def __init__(
in_channels: int,
mult: float = 1.0,
attention_head_dim: int = 32,
- qkv_multiscales: Tuple[int, ...] = (5,),
+ qkv_multiscales: tuple[int, ...] = (5,),
norm_type: str = "batch_norm",
) -> None:
super().__init__()
@@ -102,7 +101,7 @@ def get_block(
attention_head_dim: int,
norm_type: str,
act_fn: str,
- qkv_multiscales: Tuple[int, ...] = (),
+ qkv_multiscales: tuple[int, ...] = (),
):
if block_type == "ResBlock":
block = ResBlock(in_channels, out_channels, norm_type, act_fn)
@@ -205,10 +204,10 @@ def __init__(
in_channels: int,
latent_channels: int,
attention_head_dim: int = 32,
- block_type: Union[str, Tuple[str]] = "ResBlock",
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512, 1024, 1024),
- layers_per_block: Tuple[int, ...] = (2, 2, 2, 2, 2, 2),
- qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)),
+ block_type: str | tuple[str] = "ResBlock",
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024),
+ layers_per_block: tuple[int, ...] = (2, 2, 2, 2, 2, 2),
+ qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)),
downsample_block_type: str = "pixel_unshuffle",
out_shortcut: bool = True,
):
@@ -291,12 +290,12 @@ def __init__(
in_channels: int,
latent_channels: int,
attention_head_dim: int = 32,
- block_type: Union[str, Tuple[str]] = "ResBlock",
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512, 1024, 1024),
- layers_per_block: Tuple[int, ...] = (2, 2, 2, 2, 2, 2),
- qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)),
- norm_type: Union[str, Tuple[str]] = "rms_norm",
- act_fn: Union[str, Tuple[str]] = "silu",
+ block_type: str | tuple[str] = "ResBlock",
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024),
+ layers_per_block: tuple[int, ...] = (2, 2, 2, 2, 2, 2),
+ qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)),
+ norm_type: str | tuple[str] = "rms_norm",
+ act_fn: str | tuple[str] = "silu",
upsample_block_type: str = "pixel_shuffle",
in_shortcut: bool = True,
conv_act_fn: str = "relu",
@@ -391,29 +390,29 @@ class AutoencoderDC(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModel
The number of input channels in samples.
latent_channels (`int`, defaults to `32`):
The number of channels in the latent space representation.
- encoder_block_types (`Union[str, Tuple[str]]`, defaults to `"ResBlock"`):
+ encoder_block_types (`str | tuple[str]`, defaults to `"ResBlock"`):
The type(s) of block to use in the encoder.
- decoder_block_types (`Union[str, Tuple[str]]`, defaults to `"ResBlock"`):
+ decoder_block_types (`str | tuple[str]`, defaults to `"ResBlock"`):
The type(s) of block to use in the decoder.
- encoder_block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512, 1024, 1024)`):
+ encoder_block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512, 1024, 1024)`):
The number of output channels for each block in the encoder.
- decoder_block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512, 1024, 1024)`):
+ decoder_block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512, 1024, 1024)`):
The number of output channels for each block in the decoder.
- encoder_layers_per_block (`Tuple[int]`, defaults to `(2, 2, 2, 3, 3, 3)`):
+ encoder_layers_per_block (`tuple[int]`, defaults to `(2, 2, 2, 3, 3, 3)`):
The number of layers per block in the encoder.
- decoder_layers_per_block (`Tuple[int]`, defaults to `(3, 3, 3, 3, 3, 3)`):
+ decoder_layers_per_block (`tuple[int]`, defaults to `(3, 3, 3, 3, 3, 3)`):
The number of layers per block in the decoder.
- encoder_qkv_multiscales (`Tuple[Tuple[int, ...], ...]`, defaults to `((), (), (), (5,), (5,), (5,))`):
+ encoder_qkv_multiscales (`tuple[tuple[int, ...], ...]`, defaults to `((), (), (), (5,), (5,), (5,))`):
Multi-scale configurations for the encoder's QKV (query-key-value) transformations.
- decoder_qkv_multiscales (`Tuple[Tuple[int, ...], ...]`, defaults to `((), (), (), (5,), (5,), (5,))`):
+ decoder_qkv_multiscales (`tuple[tuple[int, ...], ...]`, defaults to `((), (), (), (5,), (5,), (5,))`):
Multi-scale configurations for the decoder's QKV (query-key-value) transformations.
upsample_block_type (`str`, defaults to `"pixel_shuffle"`):
The type of block to use for upsampling in the decoder.
downsample_block_type (`str`, defaults to `"pixel_unshuffle"`):
The type of block to use for downsampling in the encoder.
- decoder_norm_types (`Union[str, Tuple[str]]`, defaults to `"rms_norm"`):
+ decoder_norm_types (`str | tuple[str]`, defaults to `"rms_norm"`):
The normalization type(s) to use in the decoder.
- decoder_act_fns (`Union[str, Tuple[str]]`, defaults to `"silu"`):
+ decoder_act_fns (`str | tuple[str]`, defaults to `"silu"`):
The activation function(s) to use in the decoder.
encoder_out_shortcut (`bool`, defaults to `True`):
Whether to use shortcut at the end of the encoder.
@@ -436,18 +435,18 @@ def __init__(
in_channels: int = 3,
latent_channels: int = 32,
attention_head_dim: int = 32,
- encoder_block_types: Union[str, Tuple[str]] = "ResBlock",
- decoder_block_types: Union[str, Tuple[str]] = "ResBlock",
- encoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512, 1024, 1024),
- decoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512, 1024, 1024),
- encoder_layers_per_block: Tuple[int, ...] = (2, 2, 2, 3, 3, 3),
- decoder_layers_per_block: Tuple[int, ...] = (3, 3, 3, 3, 3, 3),
- encoder_qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)),
- decoder_qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)),
+ encoder_block_types: str | tuple[str] = "ResBlock",
+ decoder_block_types: str | tuple[str] = "ResBlock",
+ encoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024),
+ decoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024),
+ encoder_layers_per_block: tuple[int, ...] = (2, 2, 2, 3, 3, 3),
+ decoder_layers_per_block: tuple[int, ...] = (3, 3, 3, 3, 3, 3),
+ encoder_qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)),
+ decoder_qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)),
upsample_block_type: str = "pixel_shuffle",
downsample_block_type: str = "pixel_unshuffle",
- decoder_norm_types: Union[str, Tuple[str]] = "rms_norm",
- decoder_act_fns: Union[str, Tuple[str]] = "silu",
+ decoder_norm_types: str | tuple[str] = "rms_norm",
+ decoder_act_fns: str | tuple[str] = "silu",
encoder_out_shortcut: bool = True,
decoder_in_shortcut: bool = True,
decoder_conv_act_fn: str = "relu",
@@ -506,10 +505,10 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
) -> None:
r"""
Enable tiled AE decoding. When this option is enabled, the AE will split the input tensor into tiles to compute
@@ -547,7 +546,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
return encoded
@apply_forward_hook
- def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[EncoderOutput, Tuple[torch.Tensor]]:
+ def encode(self, x: torch.Tensor, return_dict: bool = True) -> EncoderOutput | tuple[torch.Tensor]:
r"""
Encode a batch of images into latents.
@@ -581,7 +580,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
return decoded
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
r"""
Decode a batch of images.
@@ -665,7 +664,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> torch.Tenso
return (encoded,)
return EncoderOutput(latent=encoded)
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, height, width = z.shape
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl.py b/src/diffusers/models/autoencoders/autoencoder_kl.py
index 95991dca3304..d2e7318f5679 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl.py
@@ -11,8 +11,6 @@
# 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.
-from typing import Optional, Tuple, Union
-
import torch
import torch.nn as nn
@@ -47,12 +45,12 @@ class AutoencoderKL(
Parameters:
in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
out_channels (int, *optional*, defaults to 3): Number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
- Tuple of downsample block types.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
- Tuple of upsample block types.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
- Tuple of block output channels.
+ down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ tuple of downsample block types.
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ tuple of upsample block types.
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
+ tuple of block output channels.
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
sample_size (`int`, *optional*, defaults to `32`): Sample input size.
@@ -81,18 +79,18 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",),
- up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
- block_out_channels: Tuple[int, ...] = (64,),
+ down_block_types: tuple[str] = ("DownEncoderBlock2D",),
+ up_block_types: tuple[str] = ("UpDecoderBlock2D",),
+ block_out_channels: tuple[int] = (64,),
layers_per_block: int = 1,
act_fn: str = "silu",
latent_channels: int = 4,
norm_num_groups: int = 32,
sample_size: int = 32,
scaling_factor: float = 0.18215,
- shift_factor: Optional[float] = None,
- latents_mean: Optional[Tuple[float]] = None,
- latents_std: Optional[Tuple[float]] = None,
+ shift_factor: float | None = None,
+ latents_mean: tuple[float] | None = None,
+ latents_std: tuple[float] | None = None,
force_upcast: bool = True,
use_quant_conv: bool = True,
use_post_quant_conv: bool = True,
@@ -172,7 +170,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -198,7 +196,7 @@ def encode(
return AutoencoderKLOutput(latent_dist=posterior)
- def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
return self.tiled_decode(z, return_dict=return_dict)
@@ -215,7 +213,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
@apply_forward_hook
def decode(
self, z: torch.FloatTensor, return_dict: bool = True, generator=None
- ) -> Union[DecoderOutput, torch.FloatTensor]:
+ ) -> DecoderOutput | torch.FloatTensor:
"""
Decode a batch of images.
@@ -363,7 +361,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder
return AutoencoderKLOutput(latent_dist=posterior)
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -417,8 +415,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py b/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py
index 6756586460d3..463f8f41bc10 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py
@@ -14,7 +14,6 @@
# limitations under the License.
import math
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -40,7 +39,7 @@ class AllegroTemporalConvLayer(nn.Module):
def __init__(
self,
in_dim: int,
- out_dim: Optional[int] = None,
+ out_dim: int | None = None,
dropout: float = 0.0,
norm_num_groups: int = 32,
up_sample: bool = False,
@@ -234,7 +233,7 @@ def __init__(
output_scale_factor: float = 1.0,
spatial_upsample: bool = True,
temporal_upsample: bool = False,
- temb_channels: Optional[int] = None,
+ temb_channels: int | None = None,
):
super().__init__()
@@ -417,14 +416,14 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"AllegroDownBlock3D",
"AllegroDownBlock3D",
"AllegroDownBlock3D",
"AllegroDownBlock3D",
),
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
- temporal_downsample_blocks: Tuple[bool, ...] = [True, True, False, False],
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+ temporal_downsample_blocks: tuple[bool, ...] = [True, True, False, False],
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -544,14 +543,14 @@ def __init__(
self,
in_channels: int = 4,
out_channels: int = 3,
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"AllegroUpBlock3D",
"AllegroUpBlock3D",
"AllegroUpBlock3D",
"AllegroUpBlock3D",
),
- temporal_upsample_blocks: Tuple[bool, ...] = [False, True, True, False],
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+ temporal_upsample_blocks: tuple[bool, ...] = [False, True, True, False],
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -687,14 +686,14 @@ class AutoencoderKLAllegro(ModelMixin, AutoencoderMixin, ConfigMixin):
Number of channels in the input image.
out_channels (int, defaults to `3`):
Number of channels in the output.
- down_block_types (`Tuple[str, ...]`, defaults to `("AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D")`):
- Tuple of strings denoting which types of down blocks to use.
- up_block_types (`Tuple[str, ...]`, defaults to `("AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D")`):
- Tuple of strings denoting which types of up blocks to use.
- block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
- Tuple of integers denoting number of output channels in each block.
- temporal_downsample_blocks (`Tuple[bool, ...]`, defaults to `(True, True, False, False)`):
- Tuple of booleans denoting which blocks to enable temporal downsampling in.
+ down_block_types (`tuple[str, ...]`, defaults to `("AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D")`):
+ tuple of strings denoting which types of down blocks to use.
+ up_block_types (`tuple[str, ...]`, defaults to `("AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D")`):
+ tuple of strings denoting which types of up blocks to use.
+ block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+ tuple of integers denoting number of output channels in each block.
+ temporal_downsample_blocks (`tuple[bool, ...]`, defaults to `(True, True, False, False)`):
+ tuple of booleans denoting which blocks to enable temporal downsampling in.
latent_channels (`int`, defaults to `4`):
Number of channels in latents.
layers_per_block (`int`, defaults to `2`):
@@ -727,21 +726,21 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"AllegroDownBlock3D",
"AllegroDownBlock3D",
"AllegroDownBlock3D",
"AllegroDownBlock3D",
),
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"AllegroUpBlock3D",
"AllegroUpBlock3D",
"AllegroUpBlock3D",
"AllegroUpBlock3D",
),
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
- temporal_downsample_blocks: Tuple[bool, ...] = (True, True, False, False),
- temporal_upsample_blocks: Tuple[bool, ...] = (False, True, True, False),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+ temporal_downsample_blocks: tuple[bool, ...] = (True, True, False, False),
+ temporal_upsample_blocks: tuple[bool, ...] = (False, True, True, False),
latent_channels: int = 4,
layers_per_block: int = 2,
act_fn: str = "silu",
@@ -807,7 +806,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
r"""
Encode a batch of videos into latents.
@@ -842,7 +841,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
raise NotImplementedError("Decoding without tiling has not been implemented yet.")
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
"""
Decode a batch of videos.
@@ -1044,8 +1043,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py
index 79433f7b9232..9921e3932465 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py
@@ -13,8 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Dict, Optional, Tuple, Union
-
import numpy as np
import torch
import torch.nn as nn
@@ -72,7 +70,7 @@ class CogVideoXCausalConv3d(nn.Module):
Args:
in_channels (`int`): Number of channels in the input tensor.
out_channels (`int`): Number of output channels produced by the convolution.
- kernel_size (`int` or `Tuple[int, int, int]`): Kernel size of the convolutional kernel.
+ kernel_size (`int` or `tuple[int, int, int]`): Kernel size of the convolutional kernel.
stride (`int`, defaults to `1`): Stride of the convolution.
dilation (`int`, defaults to `1`): Dilation rate of the convolution.
pad_mode (`str`, defaults to `"constant"`): Padding mode.
@@ -82,7 +80,7 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int, int]],
+ kernel_size: int | tuple[int, int, int],
stride: int = 1,
dilation: int = 1,
pad_mode: str = "constant",
@@ -123,7 +121,7 @@ def __init__(
)
def fake_context_parallel_forward(
- self, inputs: torch.Tensor, conv_cache: Optional[torch.Tensor] = None
+ self, inputs: torch.Tensor, conv_cache: torch.Tensor | None = None
) -> torch.Tensor:
if self.pad_mode == "replicate":
inputs = F.pad(inputs, self.time_causal_padding, mode="replicate")
@@ -134,7 +132,7 @@ def fake_context_parallel_forward(
inputs = torch.cat(cached_inputs + [inputs], dim=2)
return inputs
- def forward(self, inputs: torch.Tensor, conv_cache: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, inputs: torch.Tensor, conv_cache: torch.Tensor | None = None) -> torch.Tensor:
inputs = self.fake_context_parallel_forward(inputs, conv_cache)
if self.pad_mode == "replicate":
@@ -174,7 +172,7 @@ def __init__(
self.conv_b = CogVideoXCausalConv3d(zq_channels, f_channels, kernel_size=1, stride=1)
def forward(
- self, f: torch.Tensor, zq: torch.Tensor, conv_cache: Optional[Dict[str, torch.Tensor]] = None
+ self, f: torch.Tensor, zq: torch.Tensor, conv_cache: dict[str, torch.Tensor] | None = None
) -> torch.Tensor:
new_conv_cache = {}
conv_cache = conv_cache or {}
@@ -227,14 +225,14 @@ class CogVideoXResnetBlock3D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
dropout: float = 0.0,
temb_channels: int = 512,
groups: int = 32,
eps: float = 1e-6,
non_linearity: str = "swish",
conv_shortcut: bool = False,
- spatial_norm_dim: Optional[int] = None,
+ spatial_norm_dim: int | None = None,
pad_mode: str = "first",
):
super().__init__()
@@ -287,9 +285,9 @@ def __init__(
def forward(
self,
inputs: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- zq: Optional[torch.Tensor] = None,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ temb: torch.Tensor | None = None,
+ zq: torch.Tensor | None = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
new_conv_cache = {}
conv_cache = conv_cache or {}
@@ -409,9 +407,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- zq: Optional[torch.Tensor] = None,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ temb: torch.Tensor | None = None,
+ zq: torch.Tensor | None = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
r"""Forward method of the `CogVideoXDownBlock3D` class."""
@@ -477,7 +475,7 @@ def __init__(
resnet_eps: float = 1e-6,
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
- spatial_norm_dim: Optional[int] = None,
+ spatial_norm_dim: int | None = None,
pad_mode: str = "first",
):
super().__init__()
@@ -504,9 +502,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- zq: Optional[torch.Tensor] = None,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ temb: torch.Tensor | None = None,
+ zq: torch.Tensor | None = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
r"""Forward method of the `CogVideoXMidBlock3D` class."""
@@ -611,9 +609,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- zq: Optional[torch.Tensor] = None,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ temb: torch.Tensor | None = None,
+ zq: torch.Tensor | None = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
r"""Forward method of the `CogVideoXUpBlock3D` class."""
@@ -652,10 +650,10 @@ class CogVideoXEncoder3D(nn.Module):
The number of input channels.
out_channels (`int`, *optional*, defaults to 3):
The number of output channels.
- down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ down_block_types (`tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available
options.
- block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`):
The number of output channels for each block.
act_fn (`str`, *optional*, defaults to `"silu"`):
The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
@@ -671,13 +669,13 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 16,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
),
- block_out_channels: Tuple[int, ...] = (128, 256, 256, 512),
+ block_out_channels: tuple[int, ...] = (128, 256, 256, 512),
layers_per_block: int = 3,
act_fn: str = "silu",
norm_eps: float = 1e-6,
@@ -743,8 +741,8 @@ def __init__(
def forward(
self,
sample: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ temb: torch.Tensor | None = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
r"""The forward method of the `CogVideoXEncoder3D` class."""
@@ -805,9 +803,9 @@ class CogVideoXDecoder3D(nn.Module):
The number of input channels.
out_channels (`int`, *optional*, defaults to 3):
The number of output channels.
- up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ up_block_types (`tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options.
- block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`):
The number of output channels for each block.
act_fn (`str`, *optional*, defaults to `"silu"`):
The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
@@ -823,13 +821,13 @@ def __init__(
self,
in_channels: int = 16,
out_channels: int = 3,
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
),
- block_out_channels: Tuple[int, ...] = (128, 256, 256, 512),
+ block_out_channels: tuple[int, ...] = (128, 256, 256, 512),
layers_per_block: int = 3,
act_fn: str = "silu",
norm_eps: float = 1e-6,
@@ -902,8 +900,8 @@ def __init__(
def forward(
self,
sample: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ temb: torch.Tensor | None = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
r"""The forward method of the `CogVideoXDecoder3D` class."""
@@ -966,12 +964,12 @@ class AutoencoderKLCogVideoX(ModelMixin, AutoencoderMixin, ConfigMixin, FromOrig
Parameters:
in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
out_channels (int, *optional*, defaults to 3): Number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
- Tuple of downsample block types.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
- Tuple of upsample block types.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
- Tuple of block output channels.
+ down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ tuple of downsample block types.
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ tuple of upsample block types.
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
+ tuple of block output channels.
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
sample_size (`int`, *optional*, defaults to `32`): Sample input size.
scaling_factor (`float`, *optional*, defaults to `1.15258426`):
@@ -995,19 +993,19 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str] = (
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
),
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str] = (
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
),
- block_out_channels: Tuple[int, ...] = (128, 256, 256, 512),
+ block_out_channels: tuple[int] = (128, 256, 256, 512),
latent_channels: int = 16,
layers_per_block: int = 3,
act_fn: str = "silu",
@@ -1017,9 +1015,9 @@ def __init__(
sample_height: int = 480,
sample_width: int = 720,
scaling_factor: float = 1.15258426,
- shift_factor: Optional[float] = None,
- latents_mean: Optional[Tuple[float]] = None,
- latents_std: Optional[Tuple[float]] = None,
+ shift_factor: float | None = None,
+ latents_mean: tuple[float] | None = None,
+ latents_std: tuple[float] | None = None,
force_upcast: float = True,
use_quant_conv: bool = False,
use_post_quant_conv: bool = False,
@@ -1090,10 +1088,10 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_overlap_factor_height: Optional[float] = None,
- tile_overlap_factor_width: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_overlap_factor_height: float | None = None,
+ tile_overlap_factor_width: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -1153,7 +1151,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -1178,7 +1176,7 @@ def encode(
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)
- def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
if self.use_tiling and (width > self.tile_latent_min_width or height > self.tile_latent_min_height):
@@ -1207,7 +1205,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
return DecoderOutput(sample=dec)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
"""
Decode a batch of images.
@@ -1321,7 +1319,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
enc = torch.cat(result_rows, dim=3)
return enc
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -1409,8 +1407,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[torch.Tensor, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> torch.Tensor | torch.Tensor:
x = sample
posterior = self.encode(x).latent_dist
if sample_posterior:
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py b/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py
index b17522d1c424..4fe1f62890be 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py
@@ -12,8 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import math
-from typing import List, Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -47,9 +48,9 @@ def __init__(
self,
in_channels: int = 1,
out_channels: int = 1,
- kernel_size: Union[int, Tuple[int, int, int]] = (3, 3, 3),
- dilation: Union[int, Tuple[int, int, int]] = (1, 1, 1),
- stride: Union[int, Tuple[int, int, int]] = (1, 1, 1),
+ kernel_size: int | tuple[int, int, int] = (3, 3, 3),
+ dilation: int | tuple[int, int, int] = (1, 1, 1),
+ stride: int | tuple[int, int, int] = (1, 1, 1),
padding: int = 1,
pad_mode: str = "constant",
) -> None:
@@ -419,7 +420,7 @@ def __init__(
attention_head_dim: int,
num_groups: int = 1,
dropout: float = 0.0,
- processor: Union["CosmosSpatialAttentionProcessor2_0", "CosmosTemporalAttentionProcessor2_0"] = None,
+ processor: "CosmosSpatialAttentionProcessor2_0" | "CosmosTemporalAttentionProcessor2_0" = None,
) -> None:
super().__init__()
self.num_attention_heads = num_attention_heads
@@ -438,7 +439,7 @@ def __init__(
if self.processor is None:
raise ValueError("CosmosCausalAttention requires a processor.")
- def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
return self.processor(self, hidden_states=hidden_states, attention_mask=attention_mask)
@@ -450,7 +451,7 @@ def __init__(self):
)
def __call__(
- self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+ self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None
) -> torch.Tensor:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
residual = hidden_states
@@ -489,7 +490,7 @@ def __init__(self):
)
def __call__(
- self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+ self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None
) -> torch.Tensor:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
residual = hidden_states
@@ -711,9 +712,9 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 16,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
num_resnet_blocks: int = 2,
- attention_resolutions: Tuple[int, ...] = (32,),
+ attention_resolutions: tuple[int, ...] = (32,),
resolution: int = 1024,
patch_size: int = 4,
patch_type: str = "haar",
@@ -795,9 +796,9 @@ def __init__(
self,
in_channels: int = 16,
out_channels: int = 3,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
num_resnet_blocks: int = 2,
- attention_resolutions: Tuple[int, ...] = (32,),
+ attention_resolutions: tuple[int, ...] = (32,),
resolution: int = 1024,
patch_size: int = 4,
patch_type: str = "haar",
@@ -886,12 +887,12 @@ class AutoencoderKLCosmos(ModelMixin, AutoencoderMixin, ConfigMixin):
Number of output channels.
latent_channels (`int`, defaults to `16`):
Number of latent channels.
- encoder_block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+ encoder_block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
Number of output channels for each encoder down block.
- decode_block_out_channels (`Tuple[int, ...]`, defaults to `(256, 512, 512, 512)`):
+ decode_block_out_channels (`tuple[int, ...]`, defaults to `(256, 512, 512, 512)`):
Number of output channels for each decoder up block.
- attention_resolutions (`Tuple[int, ...]`, defaults to `(32,)`):
- List of image/video resolutions at which to apply attention.
+ attention_resolutions (`tuple[int, ...]`, defaults to `(32,)`):
+ list of image/video resolutions at which to apply attention.
resolution (`int`, defaults to `1024`):
Base image/video resolution used for computing whether a block should have attention layers.
num_layers (`int`, defaults to `2`):
@@ -924,9 +925,9 @@ def __init__(
in_channels: int = 3,
out_channels: int = 3,
latent_channels: int = 16,
- encoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
- decode_block_out_channels: Tuple[int, ...] = (256, 512, 512, 512),
- attention_resolutions: Tuple[int, ...] = (32,),
+ encoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+ decode_block_out_channels: tuple[int, ...] = (256, 512, 512, 512),
+ attention_resolutions: tuple[int, ...] = (32,),
resolution: int = 1024,
num_layers: int = 2,
patch_size: int = 4,
@@ -934,8 +935,8 @@ def __init__(
scaling_factor: float = 1.0,
spatial_compression_ratio: int = 8,
temporal_compression_ratio: int = 8,
- latents_mean: Optional[List[float]] = LATENTS_MEAN,
- latents_std: Optional[List[float]] = LATENTS_STD,
+ latents_mean: list[float] | None = LATENTS_MEAN,
+ latents_std: list[float] | None = LATENTS_STD,
) -> None:
super().__init__()
@@ -999,12 +1000,12 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_min_num_frames: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
- tile_sample_stride_num_frames: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_min_num_frames: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
+ tile_sample_stride_num_frames: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -1050,7 +1051,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> torch.Tensor:
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)
- def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
z = self.post_quant_conv(z)
dec = self.decoder(z)
@@ -1059,7 +1060,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
return DecoderOutput(sample=dec)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
if self.use_slicing and z.shape[0] > 1:
decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
decoded = torch.cat(decoded_slices)
@@ -1075,8 +1076,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[Tuple[torch.Tensor], DecoderOutput]:
+ generator: torch.Generator | None = None,
+ ) -> tuple[torch.Tensor] | DecoderOutput:
x = sample
posterior = self.encode(x).latent_dist
if sample_posterior:
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py b/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py
index 3325d33c06bf..c1345d5de73f 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import math
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -48,11 +47,11 @@ class AutoencoderKLFlux2(
Parameters:
in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
out_channels (int, *optional*, defaults to 3): Number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
Tuple of downsample block types.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
Tuple of upsample block types.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
Tuple of block output channels.
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
@@ -74,19 +73,19 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
),
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
),
- block_out_channels: Tuple[int, ...] = (
+ block_out_channels: tuple[int, ...] = (
128,
256,
512,
@@ -103,7 +102,7 @@ def __init__(
mid_block_add_attention: bool = True,
batch_norm_eps: float = 1e-4,
batch_norm_momentum: float = 0.1,
- patch_size: Tuple[int, int] = (2, 2),
+ patch_size: tuple[int, int] = (2, 2),
):
super().__init__()
@@ -187,7 +186,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -213,7 +212,7 @@ def encode(
return AutoencoderKLOutput(latent_dist=posterior)
- def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
return self.tiled_decode(z, return_dict=return_dict)
@@ -230,7 +229,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
@apply_forward_hook
def decode(
self, z: torch.FloatTensor, return_dict: bool = True, generator=None
- ) -> Union[DecoderOutput, torch.FloatTensor]:
+ ) -> DecoderOutput | torch.FloatTensor:
"""
Decode a batch of images.
@@ -378,7 +377,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder
return AutoencoderKLOutput(latent_dist=posterior)
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -432,8 +431,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py
index ddc0aed6b0ff..a19c267b6d36 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
-
import numpy as np
import torch
import torch.nn as nn
@@ -50,10 +48,10 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int, int]] = 3,
- stride: Union[int, Tuple[int, int, int]] = 1,
- padding: Union[int, Tuple[int, int, int]] = 0,
- dilation: Union[int, Tuple[int, int, int]] = 1,
+ kernel_size: int | tuple[int, int, int] = 3,
+ stride: int | tuple[int, int, int] = 1,
+ padding: int | tuple[int, int, int] = 0,
+ dilation: int | tuple[int, int, int] = 1,
bias: bool = True,
pad_mode: str = "replicate",
) -> None:
@@ -82,11 +80,11 @@ class HunyuanVideoUpsampleCausal3D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
kernel_size: int = 3,
stride: int = 1,
bias: bool = True,
- upsample_factor: Tuple[float, float, float] = (2, 2, 2),
+ upsample_factor: tuple[float, float, float] = (2, 2, 2),
) -> None:
super().__init__()
@@ -124,7 +122,7 @@ class HunyuanVideoDownsampleCausal3D(nn.Module):
def __init__(
self,
channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
padding: int = 1,
kernel_size: int = 3,
bias: bool = True,
@@ -144,7 +142,7 @@ class HunyuanVideoResnetBlockCausal3D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
dropout: float = 0.0,
groups: int = 32,
eps: float = 1e-6,
@@ -357,7 +355,7 @@ def __init__(
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
add_upsample: bool = True,
- upsample_scale_factor: Tuple[int, int, int] = (2, 2, 2),
+ upsample_scale_factor: tuple[int, int, int] = (2, 2, 2),
) -> None:
super().__init__()
resnets = []
@@ -418,13 +416,13 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
),
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -526,13 +524,13 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
),
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -641,19 +639,19 @@ def __init__(
in_channels: int = 3,
out_channels: int = 3,
latent_channels: int = 16,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
),
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
),
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+ block_out_channels: tuple[int] = (128, 256, 512, 512),
layers_per_block: int = 2,
act_fn: str = "silu",
norm_num_groups: int = 32,
@@ -725,12 +723,12 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_min_num_frames: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
- tile_sample_stride_num_frames: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_min_num_frames: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
+ tile_sample_stride_num_frames: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -779,7 +777,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
r"""
Encode a batch of images into latents.
@@ -804,7 +802,7 @@ def encode(
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)
- def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -825,7 +823,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
return DecoderOutput(sample=dec)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images.
@@ -924,7 +922,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
return enc
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -1013,7 +1011,7 @@ def _temporal_tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames]
return enc
- def _temporal_tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def _temporal_tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
@@ -1054,8 +1052,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py
index c02f11bef40a..6922ac853554 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -238,7 +237,7 @@ def __init__(
self,
in_channels: int,
z_channels: int,
- block_out_channels: Tuple[int, ...],
+ block_out_channels: tuple[int, ...],
num_res_blocks: int,
spatial_compression_ratio: int,
non_linearity: str = "silu",
@@ -329,7 +328,7 @@ class HunyuanImageDecoder2D(nn.Module):
Number of latent channels.
out_channels : int
Number of output channels.
- block_out_channels : Tuple[int, ...]
+ block_out_channels : tuple[int, ...]
Output channels for each block.
num_res_blocks : int
Number of residual blocks per block.
@@ -344,7 +343,7 @@ def __init__(
self,
z_channels: int,
out_channels: int,
- block_out_channels: Tuple[int, ...],
+ block_out_channels: tuple[int, ...],
num_res_blocks: int,
spatial_compression_ratio: int,
upsample_match_channel: bool = True,
@@ -427,7 +426,7 @@ def __init__(
in_channels: int,
out_channels: int,
latent_channels: int,
- block_out_channels: Tuple[int, ...],
+ block_out_channels: tuple[int, ...],
layers_per_block: int,
spatial_compression_ratio: int,
sample_size: int,
@@ -467,8 +466,8 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_size: Optional[int] = None,
- tile_overlap_factor: Optional[float] = None,
+ tile_sample_min_size: int | None = None,
+ tile_overlap_factor: float | None = None,
) -> None:
r"""
Enable spatial tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles
@@ -500,7 +499,7 @@ def _encode(self, x: torch.Tensor):
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
r"""
Encode a batch of images into latents.
@@ -539,7 +538,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True):
return DecoderOutput(sample=dec)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images.
@@ -620,7 +619,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
return moments
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
"""
Decode latent using spatial tiling strategy.
@@ -670,8 +669,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py
index 973574e616bf..81957e2feed4 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -37,10 +36,10 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int, int]] = 3,
- stride: Union[int, Tuple[int, int, int]] = 1,
- padding: Union[int, Tuple[int, int, int]] = 0,
- dilation: Union[int, Tuple[int, int, int]] = 1,
+ kernel_size: int | tuple[int, int, int] = 3,
+ stride: int | tuple[int, int, int] = 1,
+ padding: int | tuple[int, int, int] = 0,
+ dilation: int | tuple[int, int, int] = 1,
bias: bool = True,
pad_mode: str = "replicate",
) -> None:
@@ -226,7 +225,7 @@ class HunyuanImageRefinerResnetBlock(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
non_linearity: str = "swish",
) -> None:
super().__init__()
@@ -315,7 +314,7 @@ def __init__(
in_channels: int,
out_channels: int,
num_layers: int = 1,
- downsample_out_channels: Optional[int] = None,
+ downsample_out_channels: int | None = None,
add_temporal_downsample: int = True,
) -> None:
super().__init__()
@@ -364,7 +363,7 @@ def __init__(
in_channels: int,
out_channels: int,
num_layers: int = 1,
- upsample_out_channels: Optional[int] = None,
+ upsample_out_channels: int | None = None,
add_temporal_upsample: bool = True,
) -> None:
super().__init__()
@@ -422,7 +421,7 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 64,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 1024, 1024),
layers_per_block: int = 2,
temporal_compression_ratio: int = 4,
spatial_compression_ratio: int = 16,
@@ -509,7 +508,7 @@ def __init__(
self,
in_channels: int = 32,
out_channels: int = 3,
- block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128),
+ block_out_channels: tuple[int, ...] = (1024, 1024, 512, 256, 128),
layers_per_block: int = 2,
spatial_compression_ratio: int = 16,
temporal_compression_ratio: int = 4,
@@ -601,7 +600,7 @@ def __init__(
in_channels: int = 3,
out_channels: int = 3,
latent_channels: int = 32,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 1024, 1024),
layers_per_block: int = 2,
spatial_compression_ratio: int = 16,
temporal_compression_ratio: int = 4,
@@ -655,11 +654,11 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
- tile_overlap_factor: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
+ tile_overlap_factor: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -697,7 +696,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
r"""
Encode a batch of images into latents.
@@ -735,7 +734,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
return dec
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images.
@@ -893,8 +892,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py
index c662c1657513..2c38b174a100 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -37,10 +36,10 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int, int]] = 3,
- stride: Union[int, Tuple[int, int, int]] = 1,
- padding: Union[int, Tuple[int, int, int]] = 0,
- dilation: Union[int, Tuple[int, int, int]] = 1,
+ kernel_size: int | tuple[int, int, int] = 3,
+ stride: int | tuple[int, int, int] = 1,
+ padding: int | tuple[int, int, int] = 0,
+ dilation: int | tuple[int, int, int] = 1,
bias: bool = True,
pad_mode: str = "replicate",
) -> None:
@@ -268,7 +267,7 @@ class HunyuanVideo15ResnetBlock(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
non_linearity: str = "swish",
) -> None:
super().__init__()
@@ -357,7 +356,7 @@ def __init__(
in_channels: int,
out_channels: int,
num_layers: int = 1,
- downsample_out_channels: Optional[int] = None,
+ downsample_out_channels: int | None = None,
add_temporal_downsample: int = True,
) -> None:
super().__init__()
@@ -406,7 +405,7 @@ def __init__(
in_channels: int,
out_channels: int,
num_layers: int = 1,
- upsample_out_channels: Optional[int] = None,
+ upsample_out_channels: int | None = None,
add_temporal_upsample: bool = True,
) -> None:
super().__init__()
@@ -464,7 +463,7 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 64,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 1024, 1024),
layers_per_block: int = 2,
temporal_compression_ratio: int = 4,
spatial_compression_ratio: int = 16,
@@ -550,7 +549,7 @@ def __init__(
self,
in_channels: int = 32,
out_channels: int = 3,
- block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128),
+ block_out_channels: tuple[int, ...] = (1024, 1024, 512, 256, 128),
layers_per_block: int = 2,
spatial_compression_ratio: int = 16,
temporal_compression_ratio: int = 4,
@@ -642,7 +641,7 @@ def __init__(
in_channels: int = 3,
out_channels: int = 3,
latent_channels: int = 32,
- block_out_channels: Tuple[int] = (128, 256, 512, 1024, 1024),
+ block_out_channels: tuple[int] = (128, 256, 512, 1024, 1024),
layers_per_block: int = 2,
spatial_compression_ratio: int = 16,
temporal_compression_ratio: int = 4,
@@ -695,11 +694,11 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_latent_min_height: Optional[int] = None,
- tile_latent_min_width: Optional[int] = None,
- tile_overlap_factor: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_latent_min_height: int | None = None,
+ tile_latent_min_width: int | None = None,
+ tile_overlap_factor: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -735,7 +734,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
r"""
Encode a batch of images into latents.
@@ -771,7 +770,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
return dec
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images.
@@ -926,8 +925,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py
index 47f2081b7e45..a7acc105e9ec 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py
@@ -13,8 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
-
import torch
import torch.nn as nn
@@ -34,9 +32,9 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int, int]] = 3,
- stride: Union[int, Tuple[int, int, int]] = 1,
- dilation: Union[int, Tuple[int, int, int]] = 1,
+ kernel_size: int | tuple[int, int, int] = 3,
+ stride: int | tuple[int, int, int] = 1,
+ dilation: int | tuple[int, int, int] = 1,
groups: int = 1,
padding_mode: str = "zeros",
is_causal: bool = True,
@@ -104,7 +102,7 @@ class LTXVideoResnetBlock3d(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
dropout: float = 0.0,
eps: float = 1e-6,
elementwise_affine: bool = False,
@@ -149,7 +147,7 @@ def __init__(
self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5)
def forward(
- self, inputs: torch.Tensor, temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None
+ self, inputs: torch.Tensor, temb: torch.Tensor | None = None, generator: torch.Generator | None = None
) -> torch.Tensor:
hidden_states = inputs
@@ -201,7 +199,7 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- stride: Union[int, Tuple[int, int, int]] = 1,
+ stride: int | tuple[int, int, int] = 1,
is_causal: bool = True,
padding_mode: str = "zeros",
) -> None:
@@ -249,7 +247,7 @@ class LTXVideoUpsampler3d(nn.Module):
def __init__(
self,
in_channels: int,
- stride: Union[int, Tuple[int, int, int]] = 1,
+ stride: int | tuple[int, int, int] = 1,
is_causal: bool = True,
residual: bool = False,
upscale_factor: int = 1,
@@ -326,7 +324,7 @@ class LTXVideoDownBlock3D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
resnet_eps: float = 1e-6,
@@ -382,8 +380,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ temb: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
) -> torch.Tensor:
r"""Forward method of the `LTXDownBlock3D` class."""
@@ -432,7 +430,7 @@ class LTXVideo095DownBlock3D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
resnet_eps: float = 1e-6,
@@ -497,8 +495,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ temb: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
) -> torch.Tensor:
r"""Forward method of the `LTXDownBlock3D` class."""
@@ -575,8 +573,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ temb: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
) -> torch.Tensor:
r"""Forward method of the `LTXMidBlock3D` class."""
@@ -628,7 +626,7 @@ class LTXVideoUpBlock3d(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
resnet_eps: float = 1e-6,
@@ -696,8 +694,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ temb: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
) -> torch.Tensor:
if self.conv_in is not None:
hidden_states = self.conv_in(hidden_states, temb, generator)
@@ -735,11 +733,11 @@ class LTXVideoEncoder3d(nn.Module):
Number of input channels.
out_channels (`int`, defaults to 128):
Number of latent channels.
- block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+ block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
The number of output channels for each block.
- spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`:
+ spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`:
Whether a block should contain spatio-temporal downscaling layers or not.
- layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
+ layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
The number of layers per block.
patch_size (`int`, defaults to `4`):
The size of spatial patches.
@@ -755,16 +753,16 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 128,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
- down_block_types: Tuple[str, ...] = (
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+ down_block_types: tuple[str, ...] = (
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
),
- spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False),
- layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4),
- downsample_type: Tuple[str, ...] = ("conv", "conv", "conv", "conv"),
+ spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False),
+ layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4),
+ downsample_type: tuple[str, ...] = ("conv", "conv", "conv", "conv"),
patch_size: int = 4,
patch_size_t: int = 1,
resnet_norm_eps: float = 1e-6,
@@ -888,11 +886,11 @@ class LTXVideoDecoder3d(nn.Module):
Number of latent channels.
out_channels (`int`, defaults to 3):
Number of output channels.
- block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+ block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
The number of output channels for each block.
- spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`:
+ spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`:
Whether a block should contain spatio-temporal upscaling layers or not.
- layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
+ layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
The number of layers per block.
patch_size (`int`, defaults to `4`):
The size of spatial patches.
@@ -910,17 +908,17 @@ def __init__(
self,
in_channels: int = 128,
out_channels: int = 3,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
- spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False),
- layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+ spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False),
+ layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4),
patch_size: int = 4,
patch_size_t: int = 1,
resnet_norm_eps: float = 1e-6,
is_causal: bool = False,
- inject_noise: Tuple[bool, ...] = (False, False, False, False),
+ inject_noise: tuple[bool, ...] = (False, False, False, False),
timestep_conditioning: bool = False,
- upsample_residual: Tuple[bool, ...] = (False, False, False, False),
- upsample_factor: Tuple[bool, ...] = (1, 1, 1, 1),
+ upsample_residual: tuple[bool, ...] = (False, False, False, False),
+ upsample_factor: tuple[bool, ...] = (1, 1, 1, 1),
) -> None:
super().__init__()
@@ -989,7 +987,7 @@ def __init__(
self.gradient_checkpointing = False
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = self.conv_in(hidden_states)
if self.timestep_scale_multiplier is not None:
@@ -1049,11 +1047,11 @@ class AutoencoderKLLTXVideo(ModelMixin, AutoencoderMixin, ConfigMixin, FromOrigi
Number of output channels.
latent_channels (`int`, defaults to `128`):
Number of latent channels.
- block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+ block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
The number of output channels for each block.
- spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`:
+ spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`:
Whether a block should contain spatio-temporal downscaling or not.
- layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
+ layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
The number of layers per block.
patch_size (`int`, defaults to `4`):
The size of spatial patches.
@@ -1082,22 +1080,22 @@ def __init__(
in_channels: int = 3,
out_channels: int = 3,
latent_channels: int = 128,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
- down_block_types: Tuple[str, ...] = (
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+ down_block_types: tuple[str, ...] = (
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
),
- decoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
- layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4),
- decoder_layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4),
- spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False),
- decoder_spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False),
- decoder_inject_noise: Tuple[bool, ...] = (False, False, False, False, False),
- downsample_type: Tuple[str, ...] = ("conv", "conv", "conv", "conv"),
- upsample_residual: Tuple[bool, ...] = (False, False, False, False),
- upsample_factor: Tuple[int, ...] = (1, 1, 1, 1),
+ decoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+ layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4),
+ decoder_layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4),
+ spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False),
+ decoder_spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False),
+ decoder_inject_noise: tuple[bool, ...] = (False, False, False, False, False),
+ downsample_type: tuple[str, ...] = ("conv", "conv", "conv", "conv"),
+ upsample_residual: tuple[bool, ...] = (False, False, False, False),
+ upsample_factor: tuple[int, ...] = (1, 1, 1, 1),
timestep_conditioning: bool = False,
patch_size: int = 4,
patch_size_t: int = 1,
@@ -1187,12 +1185,12 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_min_num_frames: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
- tile_sample_stride_num_frames: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_min_num_frames: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
+ tile_sample_stride_num_frames: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -1235,7 +1233,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -1260,8 +1258,8 @@ def encode(
return AutoencoderKLOutput(latent_dist=posterior)
def _decode(
- self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, return_dict: bool = True
- ) -> Union[DecoderOutput, torch.Tensor]:
+ self, z: torch.Tensor, temb: torch.Tensor | None = None, return_dict: bool = True
+ ) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -1282,8 +1280,8 @@ def _decode(
@apply_forward_hook
def decode(
- self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, return_dict: bool = True
- ) -> Union[DecoderOutput, torch.Tensor]:
+ self, z: torch.Tensor, temb: torch.Tensor | None = None, return_dict: bool = True
+ ) -> DecoderOutput | torch.Tensor:
"""
Decode a batch of images.
@@ -1389,8 +1387,8 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
return enc
def tiled_decode(
- self, z: torch.Tensor, temb: Optional[torch.Tensor], return_dict: bool = True
- ) -> Union[DecoderOutput, torch.Tensor]:
+ self, z: torch.Tensor, temb: torch.Tensor | None, return_dict: bool = True
+ ) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -1479,8 +1477,8 @@ def _temporal_tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
return enc
def _temporal_tiled_decode(
- self, z: torch.Tensor, temb: Optional[torch.Tensor], return_dict: bool = True
- ) -> Union[DecoderOutput, torch.Tensor]:
+ self, z: torch.Tensor, temb: torch.Tensor | None, return_dict: bool = True
+ ) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
@@ -1519,11 +1517,11 @@ def _temporal_tiled_decode(
def forward(
self,
sample: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[torch.Tensor, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> torch.Tensor | torch.Tensor:
x = sample
posterior = self.encode(x).latent_dist
if sample_posterior:
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
index 01dd55a938b6..7c04bd715c25 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
@@ -13,7 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -48,7 +47,7 @@ def __init__(self, channel_dim: int = 1, eps: float = 1e-8) -> None:
self.channel_dim = channel_dim
self.eps = eps
- def forward(self, x: torch.Tensor, channel_dim: Optional[int] = None) -> torch.Tensor:
+ def forward(self, x: torch.Tensor, channel_dim: int | None = None) -> torch.Tensor:
"""
Apply RMS normalization along the configured dimension.
"""
@@ -66,9 +65,9 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int, int]] = 3,
- stride: Union[int, Tuple[int, int, int]] = 1,
- dilation: Union[int, Tuple[int, int, int]] = 1,
+ kernel_size: int | tuple[int, int, int] = 3,
+ stride: int | tuple[int, int, int] = 1,
+ dilation: int | tuple[int, int, int] = 1,
groups: int = 1,
spatial_padding_mode: str = "zeros",
):
@@ -136,7 +135,7 @@ class LTX2VideoResnetBlock3d(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
dropout: float = 0.0,
eps: float = 1e-6,
elementwise_affine: bool = False,
@@ -188,8 +187,8 @@ def __init__(
def forward(
self,
inputs: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ temb: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
causal: bool = True,
) -> torch.Tensor:
hidden_states = inputs
@@ -243,7 +242,7 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- stride: Union[int, Tuple[int, int, int]] = 1,
+ stride: int | tuple[int, int, int] = 1,
spatial_padding_mode: str = "zeros",
) -> None:
super().__init__()
@@ -290,7 +289,7 @@ class LTXVideoUpsampler3d(nn.Module):
def __init__(
self,
in_channels: int,
- stride: Union[int, Tuple[int, int, int]] = 1,
+ stride: int | tuple[int, int, int] = 1,
residual: bool = False,
upscale_factor: int = 1,
spatial_padding_mode: str = "zeros",
@@ -366,7 +365,7 @@ class LTX2VideoDownBlock3D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
resnet_eps: float = 1e-6,
@@ -440,8 +439,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ temb: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
causal: bool = True,
) -> torch.Tensor:
r"""Forward method of the `LTXDownBlock3D` class."""
@@ -520,8 +519,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ temb: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
causal: bool = True,
) -> torch.Tensor:
r"""Forward method of the `LTXMidBlock3D` class."""
@@ -575,7 +574,7 @@ class LTX2VideoUpBlock3d(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
resnet_eps: float = 1e-6,
@@ -643,8 +642,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ temb: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
causal: bool = True,
) -> torch.Tensor:
if self.conv_in is not None:
@@ -685,13 +684,13 @@ class LTX2VideoEncoder3d(nn.Module):
Number of input channels.
out_channels (`int`, defaults to 128):
Number of latent channels.
- block_out_channels (`Tuple[int, ...]`, defaults to `(256, 512, 1024, 2048)`):
+ block_out_channels (`tuple[int, ...]`, defaults to `(256, 512, 1024, 2048)`):
The number of output channels for each block.
- spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, True)`:
+ spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, True)`:
Whether a block should contain spatio-temporal downscaling layers or not.
- layers_per_block (`Tuple[int, ...]`, defaults to `(4, 6, 6, 2, 2)`):
+ layers_per_block (`tuple[int, ...]`, defaults to `(4, 6, 6, 2, 2)`):
The number of layers per block.
- downsample_type (`Tuple[str, ...]`, defaults to `("spatial", "temporal", "spatiotemporal", "spatiotemporal")`):
+ downsample_type (`tuple[str, ...]`, defaults to `("spatial", "temporal", "spatiotemporal", "spatiotemporal")`):
The spatiotemporal downsampling pattern per block. Per-layer values can be
- `"spatial"` (downsample spatial dims by 2x)
- `"temporal"` (downsample temporal dim by 2x)
@@ -710,16 +709,16 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 128,
- block_out_channels: Tuple[int, ...] = (256, 512, 1024, 2048),
- down_block_types: Tuple[str, ...] = (
+ block_out_channels: tuple[int, ...] = (256, 512, 1024, 2048),
+ down_block_types: tuple[str, ...] = (
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
),
- spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, True),
- layers_per_block: Tuple[int, ...] = (4, 6, 6, 2, 2),
- downsample_type: Tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
+ spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, True),
+ layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2),
+ downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
patch_size: int = 4,
patch_size_t: int = 1,
resnet_norm_eps: float = 1e-6,
@@ -786,7 +785,7 @@ def __init__(
self.gradient_checkpointing = False
- def forward(self, hidden_states: torch.Tensor, causal: Optional[bool] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, causal: bool | None = None) -> torch.Tensor:
r"""The forward method of the `LTXVideoEncoder3d` class."""
p = self.patch_size
@@ -838,11 +837,11 @@ class LTX2VideoDecoder3d(nn.Module):
Number of latent channels.
out_channels (`int`, defaults to 3):
Number of output channels.
- block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+ block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
The number of output channels for each block.
- spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`:
+ spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`:
Whether a block should contain spatio-temporal upscaling layers or not.
- layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
+ layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
The number of layers per block.
patch_size (`int`, defaults to `4`):
The size of spatial patches.
@@ -860,17 +859,17 @@ def __init__(
self,
in_channels: int = 128,
out_channels: int = 3,
- block_out_channels: Tuple[int, ...] = (256, 512, 1024),
- spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True),
- layers_per_block: Tuple[int, ...] = (5, 5, 5, 5),
+ block_out_channels: tuple[int, ...] = (256, 512, 1024),
+ spatio_temporal_scaling: tuple[bool, ...] = (True, True, True),
+ layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
patch_size: int = 4,
patch_size_t: int = 1,
resnet_norm_eps: float = 1e-6,
is_causal: bool = False,
- inject_noise: Tuple[bool, ...] = (False, False, False),
+ inject_noise: tuple[bool, ...] = (False, False, False),
timestep_conditioning: bool = False,
- upsample_residual: Tuple[bool, ...] = (True, True, True),
- upsample_factor: Tuple[bool, ...] = (2, 2, 2),
+ upsample_residual: tuple[bool, ...] = (True, True, True),
+ upsample_factor: tuple[bool, ...] = (2, 2, 2),
spatial_padding_mode: str = "reflect",
) -> None:
super().__init__()
@@ -952,8 +951,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- causal: Optional[bool] = None,
+ temb: torch.Tensor | None = None,
+ causal: bool | None = None,
) -> torch.Tensor:
causal = causal or self.is_causal
@@ -1016,11 +1015,11 @@ class AutoencoderKLLTX2Video(ModelMixin, AutoencoderMixin, ConfigMixin, FromOrig
Number of output channels.
latent_channels (`int`, defaults to `128`):
Number of latent channels.
- block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+ block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
The number of output channels for each block.
- spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`:
+ spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`:
Whether a block should contain spatio-temporal downscaling or not.
- layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
+ layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
The number of layers per block.
patch_size (`int`, defaults to `4`):
The size of spatial patches.
@@ -1049,22 +1048,22 @@ def __init__(
in_channels: int = 3,
out_channels: int = 3,
latent_channels: int = 128,
- block_out_channels: Tuple[int, ...] = (256, 512, 1024, 2048),
- down_block_types: Tuple[str, ...] = (
+ block_out_channels: tuple[int, ...] = (256, 512, 1024, 2048),
+ down_block_types: tuple[str, ...] = (
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
),
- decoder_block_out_channels: Tuple[int, ...] = (256, 512, 1024),
- layers_per_block: Tuple[int, ...] = (4, 6, 6, 2, 2),
- decoder_layers_per_block: Tuple[int, ...] = (5, 5, 5, 5),
- spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, True),
- decoder_spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True),
- decoder_inject_noise: Tuple[bool, ...] = (False, False, False, False),
- downsample_type: Tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
- upsample_residual: Tuple[bool, ...] = (True, True, True),
- upsample_factor: Tuple[int, ...] = (2, 2, 2),
+ decoder_block_out_channels: tuple[int, ...] = (256, 512, 1024),
+ layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2),
+ decoder_layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
+ spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, True),
+ decoder_spatio_temporal_scaling: tuple[bool, ...] = (True, True, True),
+ decoder_inject_noise: tuple[bool, ...] = (False, False, False, False),
+ downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
+ upsample_residual: tuple[bool, ...] = (True, True, True),
+ upsample_factor: tuple[int, ...] = (2, 2, 2),
timestep_conditioning: bool = False,
patch_size: int = 4,
patch_size_t: int = 1,
@@ -1158,12 +1157,12 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_min_num_frames: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
- tile_sample_stride_num_frames: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_min_num_frames: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
+ tile_sample_stride_num_frames: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -1190,7 +1189,7 @@ def enable_tiling(
self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
- def _encode(self, x: torch.Tensor, causal: Optional[bool] = None) -> torch.Tensor:
+ def _encode(self, x: torch.Tensor, causal: bool | None = None) -> torch.Tensor:
batch_size, num_channels, num_frames, height, width = x.shape
if self.use_framewise_decoding and num_frames > self.tile_sample_min_num_frames:
@@ -1205,8 +1204,8 @@ def _encode(self, x: torch.Tensor, causal: Optional[bool] = None) -> torch.Tenso
@apply_forward_hook
def encode(
- self, x: torch.Tensor, causal: Optional[bool] = None, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ self, x: torch.Tensor, causal: bool | None = None, return_dict: bool = True
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -1233,10 +1232,10 @@ def encode(
def _decode(
self,
z: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- causal: Optional[bool] = None,
+ temb: torch.Tensor | None = None,
+ causal: bool | None = None,
return_dict: bool = True,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ ) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -1259,10 +1258,10 @@ def _decode(
def decode(
self,
z: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- causal: Optional[bool] = None,
+ temb: torch.Tensor | None = None,
+ causal: bool | None = None,
return_dict: bool = True,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ ) -> DecoderOutput | torch.Tensor:
"""
Decode a batch of images.
@@ -1317,7 +1316,7 @@ def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.
)
return b
- def tiled_encode(self, x: torch.Tensor, causal: Optional[bool] = None) -> torch.Tensor:
+ def tiled_encode(self, x: torch.Tensor, causal: bool | None = None) -> torch.Tensor:
r"""Encode a batch of images using a tiled encoder.
Args:
@@ -1370,8 +1369,8 @@ def tiled_encode(self, x: torch.Tensor, causal: Optional[bool] = None) -> torch.
return enc
def tiled_decode(
- self, z: torch.Tensor, temb: Optional[torch.Tensor], causal: Optional[bool] = None, return_dict: bool = True
- ) -> Union[DecoderOutput, torch.Tensor]:
+ self, z: torch.Tensor, temb: torch.Tensor | None, causal: bool | None = None, return_dict: bool = True
+ ) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -1431,7 +1430,7 @@ def tiled_decode(
return DecoderOutput(sample=dec)
- def _temporal_tiled_encode(self, x: torch.Tensor, causal: Optional[bool] = None) -> AutoencoderKLOutput:
+ def _temporal_tiled_encode(self, x: torch.Tensor, causal: bool | None = None) -> AutoencoderKLOutput:
batch_size, num_channels, num_frames, height, width = x.shape
latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1
@@ -1462,8 +1461,8 @@ def _temporal_tiled_encode(self, x: torch.Tensor, causal: Optional[bool] = None)
return enc
def _temporal_tiled_decode(
- self, z: torch.Tensor, temb: Optional[torch.Tensor], causal: Optional[bool] = None, return_dict: bool = True
- ) -> Union[DecoderOutput, torch.Tensor]:
+ self, z: torch.Tensor, temb: torch.Tensor | None, causal: bool | None = None, return_dict: bool = True
+ ) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
@@ -1502,13 +1501,13 @@ def _temporal_tiled_decode(
def forward(
self,
sample: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
sample_posterior: bool = False,
- encoder_causal: Optional[bool] = None,
- decoder_causal: Optional[bool] = None,
+ encoder_causal: bool | None = None,
+ decoder_causal: bool | None = None,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[torch.Tensor, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> torch.Tensor | torch.Tensor:
x = sample
posterior = self.encode(x, causal=encoder_causal).latent_dist
if sample_posterior:
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py
index b29629a29f80..f9390dab5b74 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py
@@ -13,7 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -38,9 +37,9 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int]],
+ kernel_size: int | tuple[int, int],
stride: int = 1,
- dilation: Union[int, Tuple[int, int]] = 1,
+ dilation: int | tuple[int, int] = 1,
groups: int = 1,
bias: bool = True,
causality_axis: str = "height",
@@ -140,7 +139,7 @@ class LTX2AudioResnetBlock(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
conv_shortcut: bool = False,
dropout: float = 0.0,
temb_channels: int = 512,
@@ -201,7 +200,7 @@ def __init__(
else:
self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
- def forward(self, x: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, x: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
h = self.norm1(x)
h = self.non_linearity(h)
h = self.conv1(h)
@@ -221,7 +220,7 @@ def forward(self, x: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch
class LTX2AudioDownsample(nn.Module):
- def __init__(self, in_channels: int, with_conv: bool, causality_axis: Optional[str] = "height") -> None:
+ def __init__(self, in_channels: int, with_conv: bool, causality_axis: str | None = "height") -> None:
super().__init__()
self.with_conv = with_conv
self.causality_axis = causality_axis
@@ -255,7 +254,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
class LTX2AudioUpsample(nn.Module):
- def __init__(self, in_channels: int, with_conv: bool, causality_axis: Optional[str] = "height") -> None:
+ def __init__(self, in_channels: int, with_conv: bool, causality_axis: str | None = "height") -> None:
super().__init__()
self.with_conv = with_conv
self.causality_axis = causality_axis
@@ -313,7 +312,7 @@ def unpatchify(self, audio_latents: torch.Tensor, channels: int, mel_bins: int)
return audio_latents.view(batch, time, channels, mel_bins).permute(0, 2, 1, 3)
@property
- def patch_size(self) -> Tuple[int, int, int]:
+ def patch_size(self) -> tuple[int, int, int]:
return self._patch_size
@@ -323,19 +322,19 @@ def __init__(
base_channels: int = 128,
output_channels: int = 1,
num_res_blocks: int = 2,
- attn_resolutions: Optional[Tuple[int, ...]] = None,
+ attn_resolutions: tuple[int, ...] | None = None,
in_channels: int = 2,
resolution: int = 256,
latent_channels: int = 8,
- ch_mult: Tuple[int, ...] = (1, 2, 4),
+ ch_mult: tuple[int, ...] = (1, 2, 4),
norm_type: str = "group",
- causality_axis: Optional[str] = "width",
+ causality_axis: str | None = "width",
dropout: float = 0.0,
mid_block_add_attention: bool = False,
sample_rate: int = 16000,
mel_hop_length: int = 160,
is_causal: bool = True,
- mel_bins: Optional[int] = 64,
+ mel_bins: int | None = 64,
double_z: bool = True,
):
super().__init__()
@@ -480,19 +479,19 @@ def __init__(
base_channels: int = 128,
output_channels: int = 1,
num_res_blocks: int = 2,
- attn_resolutions: Optional[Tuple[int, ...]] = None,
+ attn_resolutions: tuple[int, ...] | None = None,
in_channels: int = 2,
resolution: int = 256,
latent_channels: int = 8,
- ch_mult: Tuple[int, ...] = (1, 2, 4),
+ ch_mult: tuple[int, ...] = (1, 2, 4),
norm_type: str = "group",
- causality_axis: Optional[str] = "width",
+ causality_axis: str | None = "width",
dropout: float = 0.0,
mid_block_add_attention: bool = False,
sample_rate: int = 16000,
mel_hop_length: int = 160,
is_causal: bool = True,
- mel_bins: Optional[int] = 64,
+ mel_bins: int | None = 64,
) -> None:
super().__init__()
@@ -678,20 +677,20 @@ def __init__(
self,
base_channels: int = 128,
output_channels: int = 2,
- ch_mult: Tuple[int, ...] = (1, 2, 4),
+ ch_mult: tuple[int, ...] = (1, 2, 4),
num_res_blocks: int = 2,
- attn_resolutions: Optional[Tuple[int, ...]] = None,
+ attn_resolutions: tuple[int, ...] | None = None,
in_channels: int = 2,
resolution: int = 256,
latent_channels: int = 8,
norm_type: str = "pixel",
- causality_axis: Optional[str] = "height",
+ causality_axis: str | None = "height",
dropout: float = 0.0,
mid_block_add_attention: bool = False,
sample_rate: int = 16000,
mel_hop_length: int = 160,
is_causal: bool = True,
- mel_bins: Optional[int] = 64,
+ mel_bins: int | None = 64,
double_z: bool = True,
) -> None:
super().__init__()
@@ -774,7 +773,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
return self.decoder(z)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
if self.use_slicing and z.shape[0] > 1:
decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)]
decoded = torch.cat(decoded_slices)
@@ -791,8 +790,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
posterior = self.encode(sample).latent_dist
if sample_posterior:
z = posterior.sample(generator=generator)
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py b/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py
index 97ca9d669264..ea0e2cd00d52 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py
@@ -14,7 +14,6 @@
# limitations under the License.
import math
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -37,10 +36,10 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, ...]] = 3,
- stride: Union[int, Tuple[int, ...]] = 1,
- padding: Union[int, Tuple[int, ...]] = 1,
- dilation: Union[int, Tuple[int, ...]] = 1,
+ kernel_size: int | tuple[int, ...] = 3,
+ stride: int | tuple[int, ...] = 1,
+ padding: int | tuple[int, ...] = 1,
+ dilation: int | tuple[int, ...] = 1,
groups: int = 1,
bias: bool = True,
padding_mode: str = "zeros",
@@ -437,13 +436,13 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 8,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"SpatialDownBlock3D",
"SpatialTemporalDownBlock3D",
"SpatialTemporalDownBlock3D",
"SpatialTemporalDownBlock3D",
),
- block_out_channels: Tuple[int, ...] = [128, 256, 512, 512],
+ block_out_channels: tuple[int, ...] = [128, 256, 512, 512],
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -553,13 +552,13 @@ def __init__(
self,
in_channels: int = 8,
out_channels: int = 3,
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"SpatialUpBlock3D",
"SpatialTemporalUpBlock3D",
"SpatialTemporalUpBlock3D",
"SpatialTemporalUpBlock3D",
),
- block_out_channels: Tuple[int, ...] = [128, 256, 512, 512],
+ block_out_channels: tuple[int, ...] = [128, 256, 512, 512],
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -680,14 +679,14 @@ def __init__(
in_channels: int = 3,
latent_channels: int = 16,
out_channels: int = 3,
- block_out_channels: Tuple[int, ...] = [128, 256, 512, 512],
- down_block_types: Tuple[str, ...] = [
+ block_out_channels: tuple[int, ...] = [128, 256, 512, 512],
+ down_block_types: tuple[str, ...] = [
"SpatialDownBlock3D",
"SpatialTemporalDownBlock3D",
"SpatialTemporalDownBlock3D",
"SpatialTemporalDownBlock3D",
],
- up_block_types: Tuple[str, ...] = [
+ up_block_types: tuple[str, ...] = [
"SpatialUpBlock3D",
"SpatialTemporalUpBlock3D",
"SpatialTemporalUpBlock3D",
@@ -771,12 +770,12 @@ def _clear_conv_cache(self):
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_min_num_frames: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
- tile_sample_stride_num_frames: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_min_num_frames: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
+ tile_sample_stride_num_frames: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -808,7 +807,7 @@ def enable_tiling(
@apply_forward_hook
def _encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -838,7 +837,7 @@ def _encode(
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -863,7 +862,7 @@ def encode(
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)
- def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -890,7 +889,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
return DecoderOutput(sample=dec)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
"""
Decode a batch of images.
@@ -983,7 +982,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder
moments = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
return moments
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
sample_height = height * self.spatial_compression_ratio
sample_width = width * self.spatial_compression_ratio
@@ -1049,8 +1048,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py
index 7a64ac7de172..a0f831c867b0 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py
@@ -14,7 +14,6 @@
# limitations under the License.
import functools
-from typing import Dict, Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -83,7 +82,7 @@ class MochiResnetBlock3D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
act_fn: str = "swish",
):
super().__init__()
@@ -106,7 +105,7 @@ def __init__(
def forward(
self,
inputs: torch.Tensor,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
new_conv_cache = {}
conv_cache = conv_cache or {}
@@ -193,7 +192,7 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
chunk_size: int = 2**15,
) -> torch.Tensor:
r"""Forward method of the `MochiUpBlock3D` class."""
@@ -294,7 +293,7 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
r"""Forward method of the `MochiMidBlock3D` class."""
@@ -368,7 +367,7 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- conv_cache: Optional[Dict[str, torch.Tensor]] = None,
+ conv_cache: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
r"""Forward method of the `MochiUpBlock3D` class."""
@@ -445,13 +444,13 @@ class MochiEncoder3D(nn.Module):
The number of input channels.
out_channels (`int`, *optional*):
The number of output channels.
- block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`):
+ block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`):
The number of output channels for each block.
- layers_per_block (`Tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`):
+ layers_per_block (`tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`):
The number of resnet blocks for each block.
- temporal_expansions (`Tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`):
+ temporal_expansions (`tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`):
The temporal expansion factor for each of the up blocks.
- spatial_expansions (`Tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`):
+ spatial_expansions (`tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`):
The spatial expansion factor for each of the up blocks.
non_linearity (`str`, *optional*, defaults to `"swish"`):
The non-linearity to use in the decoder.
@@ -461,11 +460,11 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 768),
- layers_per_block: Tuple[int, ...] = (3, 3, 4, 6, 3),
- temporal_expansions: Tuple[int, ...] = (1, 2, 3),
- spatial_expansions: Tuple[int, ...] = (2, 2, 2),
- add_attention_block: Tuple[bool, ...] = (False, True, True, True, True),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 768),
+ layers_per_block: tuple[int, ...] = (3, 3, 4, 6, 3),
+ temporal_expansions: tuple[int, ...] = (1, 2, 3),
+ spatial_expansions: tuple[int, ...] = (2, 2, 2),
+ add_attention_block: tuple[bool, ...] = (False, True, True, True, True),
act_fn: str = "swish",
):
super().__init__()
@@ -499,9 +498,7 @@ def __init__(
self.gradient_checkpointing = False
- def forward(
- self, hidden_states: torch.Tensor, conv_cache: Optional[Dict[str, torch.Tensor]] = None
- ) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, conv_cache: dict[str, torch.Tensor] | None = None) -> torch.Tensor:
r"""Forward method of the `MochiEncoder3D` class."""
new_conv_cache = {}
@@ -558,13 +555,13 @@ class MochiDecoder3D(nn.Module):
The number of input channels.
out_channels (`int`, *optional*):
The number of output channels.
- block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`):
+ block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`):
The number of output channels for each block.
- layers_per_block (`Tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`):
+ layers_per_block (`tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`):
The number of resnet blocks for each block.
- temporal_expansions (`Tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`):
+ temporal_expansions (`tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`):
The temporal expansion factor for each of the up blocks.
- spatial_expansions (`Tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`):
+ spatial_expansions (`tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`):
The spatial expansion factor for each of the up blocks.
non_linearity (`str`, *optional*, defaults to `"swish"`):
The non-linearity to use in the decoder.
@@ -574,10 +571,10 @@ def __init__(
self,
in_channels: int, # 12
out_channels: int, # 3
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 768),
- layers_per_block: Tuple[int, ...] = (3, 3, 4, 6, 3),
- temporal_expansions: Tuple[int, ...] = (1, 2, 3),
- spatial_expansions: Tuple[int, ...] = (2, 2, 2),
+ block_out_channels: tuple[int, ...] = (128, 256, 512, 768),
+ layers_per_block: tuple[int, ...] = (3, 3, 4, 6, 3),
+ temporal_expansions: tuple[int, ...] = (1, 2, 3),
+ spatial_expansions: tuple[int, ...] = (2, 2, 2),
act_fn: str = "swish",
):
super().__init__()
@@ -612,9 +609,7 @@ def __init__(
self.gradient_checkpointing = False
- def forward(
- self, hidden_states: torch.Tensor, conv_cache: Optional[Dict[str, torch.Tensor]] = None
- ) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, conv_cache: dict[str, torch.Tensor] | None = None) -> torch.Tensor:
r"""Forward method of the `MochiDecoder3D` class."""
new_conv_cache = {}
@@ -668,8 +663,8 @@ class AutoencoderKLMochi(ModelMixin, AutoencoderMixin, ConfigMixin):
Parameters:
in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
out_channels (int, *optional*, defaults to 3): Number of channels in the output.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
- Tuple of block output channels.
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
+ tuple of block output channels.
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
scaling_factor (`float`, *optional*, defaults to `1.15258426`):
The component-wise standard deviation of the trained latent space computed using the first batch of the
@@ -688,15 +683,15 @@ def __init__(
self,
in_channels: int = 15,
out_channels: int = 3,
- encoder_block_out_channels: Tuple[int, ...] = (64, 128, 256, 384),
- decoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 768),
+ encoder_block_out_channels: tuple[int] = (64, 128, 256, 384),
+ decoder_block_out_channels: tuple[int] = (128, 256, 512, 768),
latent_channels: int = 12,
- layers_per_block: Tuple[int, ...] = (3, 3, 4, 6, 3),
+ layers_per_block: tuple[int, ...] = (3, 3, 4, 6, 3),
act_fn: str = "silu",
- temporal_expansions: Tuple[int, ...] = (1, 2, 3),
- spatial_expansions: Tuple[int, ...] = (2, 2, 2),
- add_attention_block: Tuple[bool, ...] = (False, True, True, True, True),
- latents_mean: Tuple[float, ...] = (
+ temporal_expansions: tuple[int, ...] = (1, 2, 3),
+ spatial_expansions: tuple[int, ...] = (2, 2, 2),
+ add_attention_block: tuple[bool, ...] = (False, True, True, True, True),
+ latents_mean: tuple[float, ...] = (
-0.06730895953510081,
-0.038011381506090416,
-0.07477820912866141,
@@ -710,7 +705,7 @@ def __init__(
-0.011931556316503654,
-0.0321993391887285,
),
- latents_std: Tuple[float, ...] = (
+ latents_std: tuple[float, ...] = (
0.9263795028493863,
0.9248894543193766,
0.9393059390890617,
@@ -790,10 +785,10 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -860,7 +855,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -885,7 +880,7 @@ def encode(
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)
- def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
batch_size, num_channels, num_frames, height, width = z.shape
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -915,7 +910,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
return DecoderOutput(sample=dec)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
"""
Decode a batch of images.
@@ -1013,7 +1008,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
return enc
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -1096,8 +1091,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[torch.Tensor, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> torch.Tensor | torch.Tensor:
x = sample
posterior = self.encode(x).latent_dist
if sample_posterior:
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py
index 7f7266146e6b..f2ca0f42a272 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py
@@ -18,8 +18,6 @@
# - GitHub: https://github.com/Wan-Video/Wan2.1
# - Paper: https://huggingface.co/papers/2503.20314
-from typing import List, Optional, Tuple, Union
-
import torch
import torch.nn as nn
import torch.nn.functional as F
@@ -58,9 +56,9 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int, int]],
- stride: Union[int, Tuple[int, int, int]] = 1,
- padding: Union[int, Tuple[int, int, int]] = 0,
+ kernel_size: int | tuple[int, int, int],
+ stride: int | tuple[int, int, int] = 1,
+ padding: int | tuple[int, int, int] = 0,
) -> None:
super().__init__(
in_channels=in_channels,
@@ -497,7 +495,7 @@ def __init__(
out_dim: int,
num_res_blocks: int,
dropout: float = 0.0,
- upsample_mode: Optional[str] = None,
+ upsample_mode: str | None = None,
non_linearity: str = "silu",
):
super().__init__()
@@ -681,14 +679,14 @@ def __init__(
self,
base_dim: int = 96,
z_dim: int = 16,
- dim_mult: Tuple[int, ...] = (1, 2, 4, 4),
+ dim_mult: list[int] = [1, 2, 4, 4],
num_res_blocks: int = 2,
- attn_scales: List[float] = [],
- temperal_downsample: List[bool] = [False, True, True],
+ attn_scales: list[float] = [],
+ temperal_downsample: list[bool] = [False, True, True],
dropout: float = 0.0,
input_channels: int = 3,
- latents_mean: List[float] = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921],
- latents_std: List[float] = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160],
+ latents_mean: list[float] = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921],
+ latents_std: list[float] = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160],
) -> None:
# fmt: on
super().__init__()
@@ -738,10 +736,10 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -809,7 +807,7 @@ def _encode(self, x: torch.Tensor):
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
r"""
Encode a batch of images into latents.
@@ -859,7 +857,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True):
return DecoderOutput(sample=out)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images.
@@ -965,7 +963,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
return enc
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -1033,8 +1031,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py
index 7a307b1eacd8..95d4b0b7b535 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import itertools
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -32,7 +31,7 @@ def __init__(
self,
in_channels: int = 4,
out_channels: int = 3,
- block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+ block_out_channels: tuple[int] = (128, 256, 512, 512),
layers_per_block: int = 2,
):
super().__init__()
@@ -146,10 +145,10 @@ class AutoencoderKLTemporalDecoder(ModelMixin, AttentionMixin, AutoencoderMixin,
Parameters:
in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
out_channels (int, *optional*, defaults to 3): Number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
- Tuple of downsample block types.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
- Tuple of block output channels.
+ down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ tuple of downsample block types.
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
+ tuple of block output channels.
layers_per_block: (`int`, *optional*, defaults to 1): Number of layers per block.
latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
sample_size (`int`, *optional*, defaults to `32`): Sample input size.
@@ -173,8 +172,8 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",),
- block_out_channels: Tuple[int, ...] = (64,),
+ down_block_types: tuple[str] = ("DownEncoderBlock2D",),
+ block_out_channels: tuple[int] = (64,),
layers_per_block: int = 1,
latent_channels: int = 4,
sample_size: int = 32,
@@ -219,7 +218,7 @@ def set_default_attn_processor(self):
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -249,7 +248,7 @@ def decode(
z: torch.Tensor,
num_frames: int,
return_dict: bool = True,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ ) -> DecoderOutput | torch.Tensor:
"""
Decode a batch of images.
@@ -278,9 +277,9 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
num_frames: int = 1,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ ) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
index 761dff2dc61a..ea5d2efe642f 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import torch
import torch.nn as nn
import torch.nn.functional as F
@@ -149,9 +147,9 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int, int]],
- stride: Union[int, Tuple[int, int, int]] = 1,
- padding: Union[int, Tuple[int, int, int]] = 0,
+ kernel_size: int | tuple[int, int, int],
+ stride: int | tuple[int, int, int] = 1,
+ padding: int | tuple[int, int, int] = 0,
) -> None:
super().__init__(
in_channels=in_channels,
@@ -730,7 +728,7 @@ def __init__(
out_dim: int,
num_res_blocks: int,
dropout: float = 0.0,
- upsample_mode: Optional[str] = None,
+ upsample_mode: str | None = None,
non_linearity: str = "silu",
):
super().__init__()
@@ -971,14 +969,14 @@ class AutoencoderKLWan(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalMo
def __init__(
self,
base_dim: int = 96,
- decoder_base_dim: Optional[int] = None,
+ decoder_base_dim: int | None = None,
z_dim: int = 16,
- dim_mult: List[int] = [1, 2, 4, 4],
+ dim_mult: list[int] = [1, 2, 4, 4],
num_res_blocks: int = 2,
- attn_scales: List[float] = [],
- temperal_downsample: List[bool] = [False, True, True],
+ attn_scales: list[float] = [],
+ temperal_downsample: list[bool] = [False, True, True],
dropout: float = 0.0,
- latents_mean: List[float] = [
+ latents_mean: list[float] = [
-0.7571,
-0.7089,
-0.9113,
@@ -996,7 +994,7 @@ def __init__(
0.2503,
-0.2921,
],
- latents_std: List[float] = [
+ latents_std: list[float] = [
2.8184,
1.4541,
2.3275,
@@ -1017,9 +1015,9 @@ def __init__(
is_residual: bool = False,
in_channels: int = 3,
out_channels: int = 3,
- patch_size: Optional[int] = None,
- scale_factor_temporal: Optional[int] = 4,
- scale_factor_spatial: Optional[int] = 8,
+ patch_size: int | None = None,
+ scale_factor_temporal: int | None = 4,
+ scale_factor_spatial: int | None = 8,
) -> None:
super().__init__()
@@ -1087,10 +1085,10 @@ def __init__(
def enable_tiling(
self,
- tile_sample_min_height: Optional[int] = None,
- tile_sample_min_width: Optional[int] = None,
- tile_sample_stride_height: Optional[float] = None,
- tile_sample_stride_width: Optional[float] = None,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_stride_height: float | None = None,
+ tile_sample_stride_width: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -1155,7 +1153,7 @@ def _encode(self, x: torch.Tensor):
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
r"""
Encode a batch of images into latents.
@@ -1211,7 +1209,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True):
return DecoderOutput(sample=out)
@apply_forward_hook
- def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images.
@@ -1323,7 +1321,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
return enc
- def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Decode a batch of images using a tiled decoder.
@@ -1406,8 +1404,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | torch.Tensor:
"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_oobleck.py b/src/diffusers/models/autoencoders/autoencoder_oobleck.py
index d83264559209..239317cffd71 100644
--- a/src/diffusers/models/autoencoders/autoencoder_oobleck.py
+++ b/src/diffusers/models/autoencoders/autoencoder_oobleck.py
@@ -13,7 +13,6 @@
# limitations under the License.
import math
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -153,7 +152,7 @@ def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
self.logvar = torch.log(self.var)
self.deterministic = deterministic
- def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+ def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
# make sure sample is on the same device as the parameters and has same dtype
sample = randn_tensor(
self.mean.shape,
@@ -303,9 +302,9 @@ class AutoencoderOobleck(ModelMixin, AutoencoderMixin, ConfigMixin):
Parameters:
encoder_hidden_size (`int`, *optional*, defaults to 128):
Intermediate representation dimension for the encoder.
- downsampling_ratios (`List[int]`, *optional*, defaults to `[2, 4, 4, 8, 8]`):
+ downsampling_ratios (`list[int]`, *optional*, defaults to `[2, 4, 4, 8, 8]`):
Ratios for downsampling in the encoder. These are used in reverse order for upsampling in the decoder.
- channel_multiples (`List[int]`, *optional*, defaults to `[1, 2, 4, 8, 16]`):
+ channel_multiples (`list[int]`, *optional*, defaults to `[1, 2, 4, 8, 16]`):
Multiples used to determine the hidden sizes of the hidden layers.
decoder_channels (`int`, *optional*, defaults to 128):
Intermediate representation dimension for the decoder.
@@ -360,7 +359,7 @@ def __init__(
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[AutoencoderOobleckOutput, Tuple[OobleckDiagonalGaussianDistribution]]:
+ ) -> AutoencoderOobleckOutput | tuple[OobleckDiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -386,7 +385,7 @@ def encode(
return AutoencoderOobleckOutput(latent_dist=posterior)
- def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[OobleckDecoderOutput, torch.Tensor]:
+ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> OobleckDecoderOutput | torch.Tensor:
dec = self.decoder(z)
if not return_dict:
@@ -397,7 +396,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[OobleckDec
@apply_forward_hook
def decode(
self, z: torch.FloatTensor, return_dict: bool = True, generator=None
- ) -> Union[OobleckDecoderOutput, torch.FloatTensor]:
+ ) -> OobleckDecoderOutput | torch.FloatTensor:
"""
Decode a batch of images.
@@ -428,8 +427,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[OobleckDecoderOutput, torch.Tensor]:
+ generator: torch.Generator | None = None,
+ ) -> OobleckDecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_tiny.py b/src/diffusers/models/autoencoders/autoencoder_tiny.py
index b9ac713d7392..14bc57ca9022 100644
--- a/src/diffusers/models/autoencoders/autoencoder_tiny.py
+++ b/src/diffusers/models/autoencoders/autoencoder_tiny.py
@@ -14,7 +14,6 @@
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import torch
@@ -50,11 +49,11 @@ class AutoencoderTiny(ModelMixin, AutoencoderMixin, ConfigMixin):
Parameters:
in_channels (`int`, *optional*, defaults to 3): Number of channels in the input image.
out_channels (`int`, *optional*, defaults to 3): Number of channels in the output.
- encoder_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`):
- Tuple of integers representing the number of output channels for each encoder block. The length of the
+ encoder_block_out_channels (`tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`):
+ tuple of integers representing the number of output channels for each encoder block. The length of the
tuple should be equal to the number of encoder blocks.
- decoder_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`):
- Tuple of integers representing the number of output channels for each decoder block. The length of the
+ decoder_block_out_channels (`tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`):
+ tuple of integers representing the number of output channels for each decoder block. The length of the
tuple should be equal to the number of decoder blocks.
act_fn (`str`, *optional*, defaults to `"relu"`):
Activation function to be used throughout the model.
@@ -64,12 +63,12 @@ class AutoencoderTiny(ModelMixin, AutoencoderMixin, ConfigMixin):
upsampling_scaling_factor (`int`, *optional*, defaults to 2):
Scaling factor for upsampling in the decoder. It determines the size of the output image during the
upsampling process.
- num_encoder_blocks (`Tuple[int]`, *optional*, defaults to `(1, 3, 3, 3)`):
- Tuple of integers representing the number of encoder blocks at each stage of the encoding process. The
+ num_encoder_blocks (`tuple[int]`, *optional*, defaults to `(1, 3, 3, 3)`):
+ tuple of integers representing the number of encoder blocks at each stage of the encoding process. The
length of the tuple should be equal to the number of stages in the encoder. Each stage has a different
number of encoder blocks.
- num_decoder_blocks (`Tuple[int]`, *optional*, defaults to `(3, 3, 3, 1)`):
- Tuple of integers representing the number of decoder blocks at each stage of the decoding process. The
+ num_decoder_blocks (`tuple[int]`, *optional*, defaults to `(3, 3, 3, 1)`):
+ tuple of integers representing the number of decoder blocks at each stage of the decoding process. The
length of the tuple should be equal to the number of stages in the decoder. Each stage has a different
number of decoder blocks.
latent_magnitude (`float`, *optional*, defaults to 3.0):
@@ -99,14 +98,14 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- encoder_block_out_channels: Tuple[int, ...] = (64, 64, 64, 64),
- decoder_block_out_channels: Tuple[int, ...] = (64, 64, 64, 64),
+ encoder_block_out_channels: tuple[int, ...] = (64, 64, 64, 64),
+ decoder_block_out_channels: tuple[int, ...] = (64, 64, 64, 64),
act_fn: str = "relu",
upsample_fn: str = "nearest",
latent_channels: int = 4,
upsampling_scaling_factor: int = 2,
- num_encoder_blocks: Tuple[int, ...] = (1, 3, 3, 3),
- num_decoder_blocks: Tuple[int, ...] = (3, 3, 3, 1),
+ num_encoder_blocks: tuple[int, ...] = (1, 3, 3, 3),
+ num_decoder_blocks: tuple[int, ...] = (3, 3, 3, 1),
latent_magnitude: int = 3,
latent_shift: float = 0.5,
force_upcast: bool = False,
@@ -258,7 +257,7 @@ def _tiled_decode(self, x: torch.Tensor) -> torch.Tensor:
return out
@apply_forward_hook
- def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[AutoencoderTinyOutput, Tuple[torch.Tensor]]:
+ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderTinyOutput | tuple[torch.Tensor]:
if self.use_slicing and x.shape[0] > 1:
output = [
self._tiled_encode(x_slice) if self.use_tiling else self.encoder(x_slice) for x_slice in x.split(1)
@@ -274,8 +273,8 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[Autoencoder
@apply_forward_hook
def decode(
- self, x: torch.Tensor, generator: Optional[torch.Generator] = None, return_dict: bool = True
- ) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ self, x: torch.Tensor, generator: torch.Generator | None = None, return_dict: bool = True
+ ) -> DecoderOutput | tuple[torch.Tensor]:
if self.use_slicing and x.shape[0] > 1:
output = [
self._tiled_decode(x_slice) if self.use_tiling else self.decoder(x_slice) for x_slice in x.split(1)
@@ -293,7 +292,7 @@ def forward(
self,
sample: torch.Tensor,
return_dict: bool = True,
- ) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ ) -> DecoderOutput | tuple[torch.Tensor]:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/consistency_decoder_vae.py b/src/diffusers/models/autoencoders/consistency_decoder_vae.py
index db9404f4ac70..2d53b745e0fb 100644
--- a/src/diffusers/models/autoencoders/consistency_decoder_vae.py
+++ b/src/diffusers/models/autoencoders/consistency_decoder_vae.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import torch
import torch.nn.functional as F
@@ -77,9 +76,9 @@ def __init__(
latent_channels: int = 4,
sample_size: int = 32,
encoder_act_fn: str = "silu",
- encoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+ encoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
encoder_double_z: bool = True,
- encoder_down_block_types: Tuple[str, ...] = (
+ encoder_down_block_types: tuple[str, ...] = (
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
@@ -90,8 +89,8 @@ def __init__(
encoder_norm_num_groups: int = 32,
encoder_out_channels: int = 4,
decoder_add_attention: bool = False,
- decoder_block_out_channels: Tuple[int, ...] = (320, 640, 1024, 1024),
- decoder_down_block_types: Tuple[str, ...] = (
+ decoder_block_out_channels: tuple[int, ...] = (320, 640, 1024, 1024),
+ decoder_down_block_types: tuple[str, ...] = (
"ResnetDownsampleBlock2D",
"ResnetDownsampleBlock2D",
"ResnetDownsampleBlock2D",
@@ -106,7 +105,7 @@ def __init__(
decoder_out_channels: int = 6,
decoder_resnet_time_scale_shift: str = "scale_shift",
decoder_time_embedding_type: str = "learned",
- decoder_up_block_types: Tuple[str, ...] = (
+ decoder_up_block_types: tuple[str, ...] = (
"ResnetUpsampleBlock2D",
"ResnetUpsampleBlock2D",
"ResnetUpsampleBlock2D",
@@ -186,7 +185,7 @@ def set_default_attn_processor(self):
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
- ) -> Union[ConsistencyDecoderVAEOutput, Tuple[DiagonalGaussianDistribution]]:
+ ) -> ConsistencyDecoderVAEOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
@@ -222,21 +221,21 @@ def encode(
def decode(
self,
z: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
num_inference_steps: int = 2,
- ) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ ) -> DecoderOutput | tuple[torch.Tensor]:
"""
Decodes the input latent vector `z` using the consistency decoder VAE model.
Args:
z (torch.Tensor): The input latent vector.
- generator (Optional[torch.Generator]): The random number generator. Default is None.
+ generator (torch.Generator | None): The random number generator. Default is None.
return_dict (bool): Whether to return the output as a dictionary. Default is True.
num_inference_steps (int): The number of inference steps. Default is 2.
Returns:
- Union[DecoderOutput, Tuple[torch.Tensor]]: The decoded output.
+ DecoderOutput | tuple[torch.Tensor]: The decoded output.
"""
z = (z * self.config.scaling_factor - self.means) / self.stds
@@ -279,7 +278,7 @@ def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.
b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
return b
- def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[ConsistencyDecoderVAEOutput, Tuple]:
+ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> ConsistencyDecoderVAEOutput | tuple:
r"""Encode a batch of images using a tiled encoder.
When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
@@ -339,8 +338,8 @@ def forward(
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
- generator: Optional[torch.Generator] = None,
- ) -> Union[DecoderOutput, Tuple[torch.Tensor]]:
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | tuple[torch.Tensor]:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/vae.py b/src/diffusers/models/autoencoders/vae.py
index 9c6031a988f9..042cb5c10021 100644
--- a/src/diffusers/models/autoencoders/vae.py
+++ b/src/diffusers/models/autoencoders/vae.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple
import numpy as np
import torch
@@ -54,7 +53,7 @@ class DecoderOutput(BaseOutput):
"""
sample: torch.Tensor
- commit_loss: Optional[torch.FloatTensor] = None
+ commit_loss: torch.FloatTensor | None = None
class Encoder(nn.Module):
@@ -66,10 +65,10 @@ class Encoder(nn.Module):
The number of input channels.
out_channels (`int`, *optional*, defaults to 3):
The number of output channels.
- down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ down_block_types (`tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available
options.
- block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`):
The number of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2):
The number of layers per block.
@@ -85,8 +84,8 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",),
- block_out_channels: Tuple[int, ...] = (64,),
+ down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",),
+ block_out_channels: tuple[int, ...] = (64,),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -187,9 +186,9 @@ class Decoder(nn.Module):
The number of input channels.
out_channels (`int`, *optional*, defaults to 3):
The number of output channels.
- up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ up_block_types (`tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options.
- block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`):
The number of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2):
The number of layers per block.
@@ -205,8 +204,8 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
- block_out_channels: Tuple[int, ...] = (64,),
+ up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",),
+ block_out_channels: tuple[int, ...] = (64,),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -280,7 +279,7 @@ def __init__(
def forward(
self,
sample: torch.Tensor,
- latent_embeds: Optional[torch.Tensor] = None,
+ latent_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
r"""The forward method of the `Decoder` class."""
@@ -402,9 +401,9 @@ class MaskConditionDecoder(nn.Module):
The number of input channels.
out_channels (`int`, *optional*, defaults to 3):
The number of output channels.
- up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ up_block_types (`tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options.
- block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`):
The number of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2):
The number of layers per block.
@@ -420,8 +419,8 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
- block_out_channels: Tuple[int, ...] = (64,),
+ up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",),
+ block_out_channels: tuple[int, ...] = (64,),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
@@ -500,9 +499,9 @@ def __init__(
def forward(
self,
z: torch.Tensor,
- image: Optional[torch.Tensor] = None,
- mask: Optional[torch.Tensor] = None,
- latent_embeds: Optional[torch.Tensor] = None,
+ image: torch.Tensor | None = None,
+ mask: torch.Tensor | None = None,
+ latent_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
r"""The forward method of the `MaskConditionDecoder` class."""
sample = z
@@ -633,7 +632,7 @@ def unmap_to_all(self, inds: torch.LongTensor) -> torch.LongTensor:
back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds)
return back.reshape(ishape)
- def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Tuple]:
+ def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, tuple]:
# reshape z -> (batch, height, width, channel) and flatten
z = z.permute(0, 2, 3, 1).contiguous()
z_flattened = z.view(-1, self.vq_embed_dim)
@@ -667,7 +666,7 @@ def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Tuple]:
return z_q, loss, (perplexity, min_encodings, min_encoding_indices)
- def get_codebook_entry(self, indices: torch.LongTensor, shape: Tuple[int, ...]) -> torch.Tensor:
+ def get_codebook_entry(self, indices: torch.LongTensor, shape: tuple[int, ...]) -> torch.Tensor:
# shape specifying (batch, height, width, channel)
if self.remap is not None:
indices = indices.reshape(shape[0], -1) # add batch axis
@@ -698,7 +697,7 @@ def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
self.mean, device=self.parameters.device, dtype=self.parameters.dtype
)
- def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+ def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
# make sure sample is on the same device as the parameters and has same dtype
sample = randn_tensor(
self.mean.shape,
@@ -728,7 +727,7 @@ def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
dim=[1, 2, 3],
)
- def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor:
+ def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = [1, 2, 3]) -> torch.Tensor:
if self.deterministic:
return torch.Tensor([0.0])
logtwopi = np.log(2.0 * np.pi)
@@ -745,7 +744,7 @@ class IdentityDistribution(object):
def __init__(self, parameters: torch.Tensor):
self.parameters = parameters
- def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+ def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
return self.parameters
def mode(self) -> torch.Tensor:
@@ -761,10 +760,10 @@ class EncoderTiny(nn.Module):
The number of input channels.
out_channels (`int`):
The number of output channels.
- num_blocks (`Tuple[int, ...]`):
+ num_blocks (`tuple[int, ...]`):
Each value of the tuple represents a Conv2d layer followed by `value` number of `AutoencoderTinyBlock`'s to
use.
- block_out_channels (`Tuple[int, ...]`):
+ block_out_channels (`tuple[int, ...]`):
The number of output channels for each block.
act_fn (`str`):
The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
@@ -774,8 +773,8 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- num_blocks: Tuple[int, ...],
- block_out_channels: Tuple[int, ...],
+ num_blocks: tuple[int, ...],
+ block_out_channels: tuple[int, ...],
act_fn: str,
):
super().__init__()
@@ -827,10 +826,10 @@ class DecoderTiny(nn.Module):
The number of input channels.
out_channels (`int`):
The number of output channels.
- num_blocks (`Tuple[int, ...]`):
+ num_blocks (`tuple[int, ...]`):
Each value of the tuple represents a Conv2d layer followed by `value` number of `AutoencoderTinyBlock`'s to
use.
- block_out_channels (`Tuple[int, ...]`):
+ block_out_channels (`tuple[int, ...]`):
The number of output channels for each block.
upsampling_scaling_factor (`int`):
The scaling factor to use for upsampling.
@@ -842,8 +841,8 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- num_blocks: Tuple[int, ...],
- block_out_channels: Tuple[int, ...],
+ num_blocks: tuple[int, ...],
+ block_out_channels: tuple[int, ...],
upsampling_scaling_factor: int,
act_fn: str,
upsample_fn: str,
diff --git a/src/diffusers/models/autoencoders/vq_model.py b/src/diffusers/models/autoencoders/vq_model.py
index 82436473dfc6..9214fb7faad6 100644
--- a/src/diffusers/models/autoencoders/vq_model.py
+++ b/src/diffusers/models/autoencoders/vq_model.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -48,12 +47,12 @@ class VQModel(ModelMixin, AutoencoderMixin, ConfigMixin):
Parameters:
in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
out_channels (int, *optional*, defaults to 3): Number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
- Tuple of downsample block types.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
- Tuple of upsample block types.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
- Tuple of block output channels.
+ down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ tuple of downsample block types.
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ tuple of upsample block types.
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
+ tuple of block output channels.
layers_per_block (`int`, *optional*, defaults to `1`): Number of layers per block.
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
latent_channels (`int`, *optional*, defaults to `3`): Number of channels in the latent space.
@@ -80,16 +79,16 @@ def __init__(
self,
in_channels: int = 3,
out_channels: int = 3,
- down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",),
- up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
- block_out_channels: Tuple[int, ...] = (64,),
+ down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",),
+ up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",),
+ block_out_channels: tuple[int, ...] = (64,),
layers_per_block: int = 1,
act_fn: str = "silu",
latent_channels: int = 3,
sample_size: int = 32,
num_vq_embeddings: int = 256,
norm_num_groups: int = 32,
- vq_embed_dim: Optional[int] = None,
+ vq_embed_dim: int | None = None,
scaling_factor: float = 0.18215,
norm_type: str = "group", # group, spatial
mid_block_add_attention=True,
@@ -143,7 +142,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput:
@apply_forward_hook
def decode(
self, h: torch.Tensor, force_not_quantize: bool = False, return_dict: bool = True, shape=None
- ) -> Union[DecoderOutput, torch.Tensor]:
+ ) -> DecoderOutput | torch.Tensor:
# also go through quantization layer
if not force_not_quantize:
quant, commit_loss, _ = self.quantize(h)
@@ -161,9 +160,7 @@ def decode(
return DecoderOutput(sample=dec, commit_loss=commit_loss)
- def forward(
- self, sample: torch.Tensor, return_dict: bool = True
- ) -> Union[DecoderOutput, Tuple[torch.Tensor, ...]]:
+ def forward(self, sample: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor, ...]:
r"""
The [`VQModel`] forward method.
diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py
index 04c90668a1db..5f9587a1b4de 100644
--- a/src/diffusers/models/cache_utils.py
+++ b/src/diffusers/models/cache_utils.py
@@ -41,7 +41,7 @@ def enable_cache(self, config) -> None:
Enable caching techniques on the model.
Args:
- config (`Union[PyramidAttentionBroadcastConfig, FasterCacheConfig, FirstBlockCacheConfig]`):
+ config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig`):
The configuration for applying the caching technique. Currently supported caching techniques are:
- [`~hooks.PyramidAttentionBroadcastConfig`]
- [`~hooks.FasterCacheConfig`]
diff --git a/src/diffusers/models/controlnets/__init__.py b/src/diffusers/models/controlnets/__init__.py
index fee7f231e899..853a2207f903 100644
--- a/src/diffusers/models/controlnets/__init__.py
+++ b/src/diffusers/models/controlnets/__init__.py
@@ -3,6 +3,7 @@
if is_torch_available():
from .controlnet import ControlNetModel, ControlNetOutput
+ from .controlnet_cosmos import CosmosControlNetModel
from .controlnet_flux import FluxControlNetModel, FluxControlNetOutput, FluxMultiControlNetModel
from .controlnet_hunyuan import (
HunyuanControlNetOutput,
diff --git a/src/diffusers/models/controlnets/controlnet.py b/src/diffusers/models/controlnets/controlnet.py
index 0b5b9fa3efba..8c2ff2fdd123 100644
--- a/src/diffusers/models/controlnets/controlnet.py
+++ b/src/diffusers/models/controlnets/controlnet.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
from torch import nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
from ...loaders.single_file_model import FromOriginalModelMixin
-from ...utils import BaseOutput, logging
+from ...utils import BaseOutput, apply_lora_scale, logging
from ..attention import AttentionMixin
from ..attention_processor import (
ADDED_KV_ATTENTION_PROCESSORS,
@@ -58,7 +58,7 @@ class ControlNetOutput(BaseOutput):
Output can be used to condition the original UNet's middle block activation.
"""
- down_block_res_samples: Tuple[torch.Tensor]
+ down_block_res_samples: tuple[torch.Tensor]
mid_block_res_sample: torch.Tensor
@@ -76,7 +76,7 @@ def __init__(
self,
conditioning_embedding_channels: int,
conditioning_channels: int = 3,
- block_out_channels: Tuple[int, ...] = (16, 32, 96, 256),
+ block_out_channels: tuple[int, ...] = (16, 32, 96, 256),
):
super().__init__()
@@ -120,7 +120,7 @@ class ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModel
The frequency shift to apply to the time embedding.
down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
The tuple of downsample blocks to use.
- only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`):
+ only_cross_attention (`bool | tuple[bool]`, defaults to `False`):
block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, defaults to 2):
@@ -138,7 +138,7 @@ class ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModel
The epsilon to use for the normalization.
cross_attention_dim (`int`, defaults to 1280):
The dimension of the cross attention features.
- transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
+ transformer_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 1):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
[`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
[`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
@@ -148,7 +148,7 @@ class ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModel
encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
- attention_head_dim (`Union[int, Tuple[int]]`, defaults to 8):
+ attention_head_dim (`int | tuple[int]`, defaults to 8):
The dimension of the attention heads.
use_linear_projection (`bool`, defaults to `False`):
class_embed_type (`str`, *optional*, defaults to `None`):
@@ -185,37 +185,37 @@ def __init__(
conditioning_channels: int = 3,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
- only_cross_attention: Union[bool, Tuple[bool]] = False,
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+ mid_block_type: str | None = "UNetMidBlock2DCrossAttn",
+ only_cross_attention: bool | tuple[bool] = False,
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
layers_per_block: int = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
act_fn: str = "silu",
- norm_num_groups: Optional[int] = 32,
+ norm_num_groups: int | None = 32,
norm_eps: float = 1e-5,
cross_attention_dim: int = 1280,
- transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
- encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
- attention_head_dim: Union[int, Tuple[int, ...]] = 8,
- num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
+ transformer_layers_per_block: int | tuple[int, ...] = 1,
+ encoder_hid_dim: int | None = None,
+ encoder_hid_dim_type: str | None = None,
+ attention_head_dim: int | tuple[int, ...] = 8,
+ num_attention_heads: int | tuple[int, ...] | None = None,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
- addition_time_embed_dim: Optional[int] = None,
- num_class_embeds: Optional[int] = None,
+ class_embed_type: str | None = None,
+ addition_embed_type: str | None = None,
+ addition_time_embed_dim: int | None = None,
+ num_class_embeds: int | None = None,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
- projection_class_embeddings_input_dim: Optional[int] = None,
+ projection_class_embeddings_input_dim: int | None = None,
controlnet_conditioning_channel_order: str = "rgb",
- conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+ conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256),
global_pool_conditions: bool = False,
addition_embed_type_num_heads: int = 64,
):
@@ -445,7 +445,7 @@ def from_unet(
cls,
unet: UNet2DConditionModel,
controlnet_conditioning_channel_order: str = "rgb",
- conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+ conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256),
load_weights_from_unet: bool = True,
conditioning_channels: int = 3,
):
@@ -533,7 +533,7 @@ def set_default_attn_processor(self):
self.set_attn_processor(processor)
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice
- def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None:
+ def set_attention_slice(self, slice_size: str | int | list[int]) -> None:
r"""
Enable sliced attention computation.
@@ -587,7 +587,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
# Recursively walk through all the children.
# Any children which exposes the set_attention_slice method
# gets the message
- def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]):
if hasattr(module, "set_attention_slice"):
module.set_attention_slice(slice_size.pop())
@@ -598,28 +598,29 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i
for module in self.children():
fn_recursive_set_attention_slice(module, reversed_slice_size)
+ @apply_lora_scale("cross_attention_kwargs")
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
controlnet_cond: torch.Tensor,
conditioning_scale: float = 1.0,
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guess_mode: bool = False,
return_dict: bool = True,
- ) -> Union[ControlNetOutput, Tuple[Tuple[torch.Tensor, ...], torch.Tensor]]:
+ ) -> ControlNetOutput | tuple[tuple[torch.Tensor, ...], torch.Tensor]:
"""
The [`ControlNetModel`] forward method.
Args:
sample (`torch.Tensor`):
The noisy input tensor.
- timestep (`Union[torch.Tensor, float, int]`):
+ timestep (`torch.Tensor | float | int`):
The number of timesteps to denoise an input.
encoder_hidden_states (`torch.Tensor`):
The encoder hidden states.
diff --git a/src/diffusers/models/controlnets/controlnet_cosmos.py b/src/diffusers/models/controlnets/controlnet_cosmos.py
new file mode 100644
index 000000000000..e39f8dfb568a
--- /dev/null
+++ b/src/diffusers/models/controlnets/controlnet_cosmos.py
@@ -0,0 +1,317 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from ...configuration_utils import ConfigMixin, register_to_config
+from ...loaders import FromOriginalModelMixin
+from ...utils import BaseOutput, is_torchvision_available, logging
+from ..modeling_utils import ModelMixin
+from ..transformers.transformer_cosmos import (
+ CosmosEmbedding,
+ CosmosLearnablePositionalEmbed,
+ CosmosPatchEmbed,
+ CosmosRotaryPosEmbed,
+ CosmosTransformerBlock,
+)
+
+
+if is_torchvision_available():
+ from torchvision import transforms
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+@dataclass
+class CosmosControlNetOutput(BaseOutput):
+ """
+ Output of [`CosmosControlNetModel`].
+
+ Args:
+ control_block_samples (`list[torch.Tensor]`):
+ List of control block activations to be injected into transformer blocks.
+ """
+
+ control_block_samples: List[torch.Tensor]
+
+
+class CosmosControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+ r"""
+ ControlNet for Cosmos Transfer2.5.
+
+ This model duplicates the shared embedding modules from the transformer (patch_embed, time_embed,
+ learnable_pos_embed, img_context_proj) to enable proper CPU offloading. The forward() method computes everything
+ internally from raw inputs.
+ """
+
+ _supports_gradient_checkpointing = True
+ _skip_layerwise_casting_patterns = ["patch_embed", "patch_embed_base", "time_embed"]
+ _no_split_modules = ["CosmosTransformerBlock"]
+ _keep_in_fp32_modules = ["learnable_pos_embed"]
+
+ @register_to_config
+ def __init__(
+ self,
+ n_controlnet_blocks: int = 4,
+ in_channels: int = 130,
+ latent_channels: int = 18, # base latent channels (latents + condition_mask) + padding_mask
+ model_channels: int = 2048,
+ num_attention_heads: int = 32,
+ attention_head_dim: int = 128,
+ mlp_ratio: float = 4.0,
+ text_embed_dim: int = 1024,
+ adaln_lora_dim: int = 256,
+ patch_size: Tuple[int, int, int] = (1, 2, 2),
+ max_size: Tuple[int, int, int] = (128, 240, 240),
+ rope_scale: Tuple[float, float, float] = (2.0, 1.0, 1.0),
+ extra_pos_embed_type: str | None = None,
+ img_context_dim_in: int | None = None,
+ img_context_dim_out: int = 2048,
+ use_crossattn_projection: bool = False,
+ crossattn_proj_in_channels: int = 1024,
+ encoder_hidden_states_channels: int = 1024,
+ ):
+ super().__init__()
+
+ self.patch_embed = CosmosPatchEmbed(in_channels, model_channels, patch_size, bias=False)
+
+ self.patch_embed_base = CosmosPatchEmbed(latent_channels, model_channels, patch_size, bias=False)
+ self.time_embed = CosmosEmbedding(model_channels, model_channels)
+
+ self.learnable_pos_embed = None
+ if extra_pos_embed_type == "learnable":
+ self.learnable_pos_embed = CosmosLearnablePositionalEmbed(
+ hidden_size=model_channels,
+ max_size=max_size,
+ patch_size=patch_size,
+ )
+
+ self.img_context_proj = None
+ if img_context_dim_in is not None and img_context_dim_in > 0:
+ self.img_context_proj = nn.Sequential(
+ nn.Linear(img_context_dim_in, img_context_dim_out, bias=True),
+ nn.GELU(),
+ )
+
+ # Cross-attention projection for text embeddings (same as transformer)
+ self.crossattn_proj = None
+ if use_crossattn_projection:
+ self.crossattn_proj = nn.Sequential(
+ nn.Linear(crossattn_proj_in_channels, encoder_hidden_states_channels, bias=True),
+ nn.GELU(),
+ )
+
+ # RoPE for both control and base latents
+ self.rope = CosmosRotaryPosEmbed(
+ hidden_size=attention_head_dim, max_size=max_size, patch_size=patch_size, rope_scale=rope_scale
+ )
+
+ self.control_blocks = nn.ModuleList(
+ [
+ CosmosTransformerBlock(
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ cross_attention_dim=text_embed_dim,
+ mlp_ratio=mlp_ratio,
+ adaln_lora_dim=adaln_lora_dim,
+ qk_norm="rms_norm",
+ out_bias=False,
+ img_context=img_context_dim_in is not None and img_context_dim_in > 0,
+ before_proj=(block_idx == 0),
+ after_proj=True,
+ )
+ for block_idx in range(n_controlnet_blocks)
+ ]
+ )
+
+ self.gradient_checkpointing = False
+
+ def _expand_conditioning_scale(self, conditioning_scale: float | list[float]) -> List[float]:
+ if isinstance(conditioning_scale, list):
+ scales = conditioning_scale
+ else:
+ scales = [conditioning_scale] * len(self.control_blocks)
+
+ if len(scales) < len(self.control_blocks):
+ logger.warning(
+ "Received %d control scales, but control network defines %d blocks. "
+ "Scales will be trimmed or repeated to match.",
+ len(scales),
+ len(self.control_blocks),
+ )
+ scales = (scales * len(self.control_blocks))[: len(self.control_blocks)]
+ return scales
+
+ def forward(
+ self,
+ controls_latents: torch.Tensor,
+ latents: torch.Tensor,
+ timestep: torch.Tensor,
+ encoder_hidden_states: Union[Optional[torch.Tensor], Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]],
+ condition_mask: torch.Tensor,
+ conditioning_scale: float | list[float] = 1.0,
+ padding_mask: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ fps: int | None = None,
+ return_dict: bool = True,
+ ) -> Union[CosmosControlNetOutput, Tuple[List[torch.Tensor]]]:
+ """
+ Forward pass for the ControlNet.
+
+ Args:
+ controls_latents: Control signal latents [B, C, T, H, W]
+ latents: Base latents from the noising process [B, C, T, H, W]
+ timestep: Diffusion timestep tensor
+ encoder_hidden_states: Tuple of (text_context, img_context) or text_context
+ condition_mask: Conditioning mask [B, 1, T, H, W]
+ conditioning_scale: Scale factor(s) for control outputs
+ padding_mask: Padding mask [B, 1, H, W] or None
+ attention_mask: Optional attention mask or None
+ fps: Frames per second for RoPE or None
+ return_dict: Whether to return a CosmosControlNetOutput or a tuple
+
+ Returns:
+ CosmosControlNetOutput or tuple of control tensors
+ """
+ B, C, T, H, W = controls_latents.shape
+
+ # 1. Prepare control latents
+ control_hidden_states = controls_latents
+ vace_in_channels = self.config.in_channels - 1
+ if control_hidden_states.shape[1] < vace_in_channels - 1:
+ pad_C = vace_in_channels - 1 - control_hidden_states.shape[1]
+ control_hidden_states = torch.cat(
+ [
+ control_hidden_states,
+ torch.zeros(
+ (B, pad_C, T, H, W), dtype=control_hidden_states.dtype, device=control_hidden_states.device
+ ),
+ ],
+ dim=1,
+ )
+
+ if condition_mask is not None:
+ control_hidden_states = torch.cat([control_hidden_states, condition_mask], dim=1)
+ else:
+ control_hidden_states = torch.cat(
+ [control_hidden_states, torch.zeros_like(controls_latents[:, :1])], dim=1
+ )
+
+ padding_mask_resized = transforms.functional.resize(
+ padding_mask, list(control_hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
+ )
+ control_hidden_states = torch.cat(
+ [control_hidden_states, padding_mask_resized.unsqueeze(2).repeat(B, 1, T, 1, 1)], dim=1
+ )
+
+ # 2. Prepare base latents (same processing as transformer.forward)
+ base_hidden_states = latents
+ if condition_mask is not None:
+ base_hidden_states = torch.cat([base_hidden_states, condition_mask], dim=1)
+
+ base_padding_mask = transforms.functional.resize(
+ padding_mask, list(base_hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
+ )
+ base_hidden_states = torch.cat(
+ [base_hidden_states, base_padding_mask.unsqueeze(2).repeat(B, 1, T, 1, 1)], dim=1
+ )
+
+ # 3. Generate positional embeddings (shared for both)
+ image_rotary_emb = self.rope(control_hidden_states, fps=fps)
+ extra_pos_emb = self.learnable_pos_embed(control_hidden_states) if self.learnable_pos_embed else None
+
+ # 4. Patchify control latents
+ control_hidden_states = self.patch_embed(control_hidden_states)
+ control_hidden_states = control_hidden_states.flatten(1, 3)
+
+ # 5. Patchify base latents
+ p_t, p_h, p_w = self.config.patch_size
+ post_patch_num_frames = T // p_t
+ post_patch_height = H // p_h
+ post_patch_width = W // p_w
+
+ base_hidden_states = self.patch_embed_base(base_hidden_states)
+ base_hidden_states = base_hidden_states.flatten(1, 3)
+
+ # 6. Time embeddings
+ if timestep.ndim == 1:
+ temb, embedded_timestep = self.time_embed(base_hidden_states, timestep)
+ elif timestep.ndim == 5:
+ batch_size, _, num_frames, _, _ = latents.shape
+ assert timestep.shape == (batch_size, 1, num_frames, 1, 1), (
+ f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}"
+ )
+ timestep_flat = timestep.flatten()
+ temb, embedded_timestep = self.time_embed(base_hidden_states, timestep_flat)
+ temb, embedded_timestep = (
+ x.view(batch_size, post_patch_num_frames, 1, 1, -1)
+ .expand(-1, -1, post_patch_height, post_patch_width, -1)
+ .flatten(1, 3)
+ for x in (temb, embedded_timestep)
+ )
+ else:
+ raise ValueError(f"Expected timestep to have shape [B, 1, T, 1, 1] or [T], but got {timestep.shape}")
+
+ # 7. Process encoder hidden states
+ if isinstance(encoder_hidden_states, tuple):
+ text_context, img_context = encoder_hidden_states
+ else:
+ text_context = encoder_hidden_states
+ img_context = None
+
+ # Apply cross-attention projection to text context
+ if self.crossattn_proj is not None:
+ text_context = self.crossattn_proj(text_context)
+
+ # Apply cross-attention projection to image context (if provided)
+ if img_context is not None and self.img_context_proj is not None:
+ img_context = self.img_context_proj(img_context)
+
+ # Combine text and image context into a single tuple
+ if self.config.img_context_dim_in is not None and self.config.img_context_dim_in > 0:
+ processed_encoder_hidden_states = (text_context, img_context)
+ else:
+ processed_encoder_hidden_states = text_context
+
+ # 8. Prepare attention mask
+ if attention_mask is not None:
+ attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) # [B, 1, 1, S]
+
+ # 9. Run control blocks
+ scales = self._expand_conditioning_scale(conditioning_scale)
+ result = []
+ for block_idx, (block, scale) in enumerate(zip(self.control_blocks, scales)):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ control_hidden_states, control_proj = self._gradient_checkpointing_func(
+ block,
+ control_hidden_states,
+ processed_encoder_hidden_states,
+ embedded_timestep,
+ temb,
+ image_rotary_emb,
+ extra_pos_emb,
+ attention_mask,
+ None, # controlnet_residual
+ base_hidden_states,
+ block_idx,
+ )
+ else:
+ control_hidden_states, control_proj = block(
+ hidden_states=control_hidden_states,
+ encoder_hidden_states=processed_encoder_hidden_states,
+ embedded_timestep=embedded_timestep,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ extra_pos_emb=extra_pos_emb,
+ attention_mask=attention_mask,
+ controlnet_residual=None,
+ latents=base_hidden_states,
+ block_idx=block_idx,
+ )
+ result.append(control_proj * scale)
+
+ if not return_dict:
+ return (result,)
+
+ return CosmosControlNetOutput(control_block_samples=result)
diff --git a/src/diffusers/models/controlnets/controlnet_flax.py b/src/diffusers/models/controlnets/controlnet_flax.py
index f7a8b98fa2f0..5bc1a446338e 100644
--- a/src/diffusers/models/controlnets/controlnet_flax.py
+++ b/src/diffusers/models/controlnets/controlnet_flax.py
@@ -11,8 +11,6 @@
# 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.
-from typing import Optional, Tuple, Union
-
import flax
import flax.linen as nn
import jax
@@ -49,7 +47,7 @@ class FlaxControlNetOutput(BaseOutput):
class FlaxControlNetConditioningEmbedding(nn.Module):
conditioning_embedding_channels: int
- block_out_channels: Tuple[int, ...] = (16, 32, 96, 256)
+ block_out_channels: tuple[int, ...] = (16, 32, 96, 256)
dtype: jnp.dtype = jnp.float32
def setup(self) -> None:
@@ -132,15 +130,15 @@ class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin):
The size of the input sample.
in_channels (`int`, *optional*, defaults to 4):
The number of channels in the input sample.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`):
The tuple of downsample blocks to use.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2):
The number of layers per block.
- attention_head_dim (`int` or `Tuple[int]`, *optional*, defaults to 8):
+ attention_head_dim (`int` or `tuple[int]`, *optional*, defaults to 8):
The dimension of the attention heads.
- num_attention_heads (`int` or `Tuple[int]`, *optional*):
+ num_attention_heads (`int` or `tuple[int]`, *optional*):
The number of attention heads.
cross_attention_dim (`int`, *optional*, defaults to 768):
The dimension of the cross attention features.
@@ -157,17 +155,17 @@ class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin):
sample_size: int = 32
in_channels: int = 4
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
)
- only_cross_attention: Union[bool, Tuple[bool, ...]] = False
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280)
+ only_cross_attention: bool | tuple[bool, ...] = False
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280)
layers_per_block: int = 2
- attention_head_dim: Union[int, Tuple[int, ...]] = 8
- num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None
+ attention_head_dim: int | tuple[int, ...] = 8
+ num_attention_heads: int | tuple[int, ...] | None = None
cross_attention_dim: int = 1280
dropout: float = 0.0
use_linear_projection: bool = False
@@ -175,7 +173,7 @@ class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin):
flip_sin_to_cos: bool = True
freq_shift: int = 0
controlnet_conditioning_channel_order: str = "rgb"
- conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256)
+ conditioning_embedding_out_channels: tuple[int, ...] = (16, 32, 96, 256)
def init_weights(self, rng: jax.Array) -> FrozenDict:
# init input tensors
@@ -327,13 +325,13 @@ def setup(self) -> None:
def __call__(
self,
sample: jnp.ndarray,
- timesteps: Union[jnp.ndarray, float, int],
+ timesteps: jnp.ndarray | float | int,
encoder_hidden_states: jnp.ndarray,
controlnet_cond: jnp.ndarray,
conditioning_scale: float = 1.0,
return_dict: bool = True,
train: bool = False,
- ) -> Union[FlaxControlNetOutput, Tuple[Tuple[jnp.ndarray, ...], jnp.ndarray]]:
+ ) -> FlaxControlNetOutput | tuple[tuple[jnp.ndarray, ...], jnp.ndarray]:
r"""
Args:
sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor
diff --git a/src/diffusers/models/controlnets/controlnet_flux.py b/src/diffusers/models/controlnets/controlnet_flux.py
index 639a8ad7390a..56482c299c05 100644
--- a/src/diffusers/models/controlnets/controlnet_flux.py
+++ b/src/diffusers/models/controlnets/controlnet_flux.py
@@ -13,14 +13,18 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, BaseOutput, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import (
+ BaseOutput,
+ apply_lora_scale,
+ logging,
+)
from ..attention import AttentionMixin
from ..controlnets.controlnet import ControlNetConditioningEmbedding, zero_module
from ..embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
@@ -34,8 +38,8 @@
@dataclass
class FluxControlNetOutput(BaseOutput):
- controlnet_block_samples: Tuple[torch.Tensor]
- controlnet_single_block_samples: Tuple[torch.Tensor]
+ controlnet_block_samples: tuple[torch.Tensor]
+ controlnet_single_block_samples: tuple[torch.Tensor]
class FluxControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin):
@@ -53,7 +57,7 @@ def __init__(
joint_attention_dim: int = 4096,
pooled_projection_dim: int = 768,
guidance_embeds: bool = False,
- axes_dims_rope: List[int] = [16, 56, 56],
+ axes_dims_rope: list[int] = [16, 56, 56],
num_mode: int = None,
conditioning_embedding_channels: int = None,
):
@@ -150,6 +154,7 @@ def from_transformer(
return controlnet
+ @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -162,9 +167,9 @@ def forward(
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
guidance: torch.Tensor = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+ ) -> torch.FloatTensor | Transformer2DModelOutput:
"""
The [`FluxTransformer2DModel`] forward method.
@@ -197,20 +202,6 @@ def forward(
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
- if joint_attention_kwargs is not None:
- joint_attention_kwargs = joint_attention_kwargs.copy()
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
hidden_states = self.x_embedder(hidden_states)
if self.input_hint_block is not None:
@@ -323,10 +314,6 @@ def forward(
None if len(controlnet_single_block_samples) == 0 else controlnet_single_block_samples
)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (controlnet_block_samples, controlnet_single_block_samples)
@@ -344,7 +331,7 @@ class FluxMultiControlNetModel(ModelMixin):
compatible with `FluxControlNetModel`.
Args:
- controlnets (`List[FluxControlNetModel]`):
+ controlnets (`list[FluxControlNetModel]`):
Provides additional conditioning to the unet during the denoising process. You must set multiple
`FluxControlNetModel` as a list.
"""
@@ -356,18 +343,18 @@ def __init__(self, controlnets):
def forward(
self,
hidden_states: torch.FloatTensor,
- controlnet_cond: List[torch.tensor],
- controlnet_mode: List[torch.tensor],
- conditioning_scale: List[float],
+ controlnet_cond: list[torch.tensor],
+ controlnet_mode: list[torch.tensor],
+ conditioning_scale: list[float],
encoder_hidden_states: torch.Tensor = None,
pooled_projections: torch.Tensor = None,
timestep: torch.LongTensor = None,
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
guidance: torch.Tensor = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[FluxControlNetOutput, Tuple]:
+ ) -> FluxControlNetOutput | tuple:
# ControlNet-Union with multiple conditions
# only load one ControlNet for saving memories
if len(self.nets) == 1:
diff --git a/src/diffusers/models/controlnets/controlnet_hunyuan.py b/src/diffusers/models/controlnets/controlnet_hunyuan.py
index d17d5692aa40..6ef92d78dd6e 100644
--- a/src/diffusers/models/controlnets/controlnet_hunyuan.py
+++ b/src/diffusers/models/controlnets/controlnet_hunyuan.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
-from typing import Dict, Optional, Union
import torch
from torch import nn
@@ -27,7 +26,7 @@
)
from ..modeling_utils import ModelMixin
from ..transformers.hunyuan_transformer_2d import HunyuanDiTBlock
-from .controlnet import Tuple, zero_module
+from .controlnet import zero_module
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
@@ -35,7 +34,7 @@
@dataclass
class HunyuanControlNetOutput(BaseOutput):
- controlnet_block_samples: Tuple[torch.Tensor]
+ controlnet_block_samples: tuple[torch.Tensor]
class HunyuanDiT2DControlNetModel(ModelMixin, ConfigMixin):
@@ -45,8 +44,8 @@ def __init__(
conditioning_channels: int = 3,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
- patch_size: Optional[int] = None,
+ in_channels: int | None = None,
+ patch_size: int | None = None,
activation_fn: str = "gelu-approximate",
sample_size=32,
hidden_size=1152,
@@ -116,7 +115,7 @@ def __init__(
self.controlnet_blocks.append(controlnet_block)
@property
- def attn_processors(self) -> Dict[str, AttentionProcessor]:
+ def attn_processors(self) -> dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
@@ -125,7 +124,7 @@ def attn_processors(self) -> Dict[str, AttentionProcessor]:
# set recursively
processors = {}
- def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict[str, AttentionProcessor]):
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
@@ -139,7 +138,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors:
return processors
- def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+ def set_attn_processor(self, processor: AttentionProcessor | dict[str, AttentionProcessor]):
r"""
Sets the attention processor to use to compute attention.
@@ -317,7 +316,7 @@ class HunyuanDiT2DMultiControlNetModel(ModelMixin):
designed to be compatible with `HunyuanDiT2DControlNetModel`.
Args:
- controlnets (`List[HunyuanDiT2DControlNetModel]`):
+ controlnets (`list[HunyuanDiT2DControlNetModel]`):
Provides additional conditioning to the unet during the denoising process. You must set multiple
`HunyuanDiT2DControlNetModel` as a list.
"""
diff --git a/src/diffusers/models/controlnets/controlnet_qwenimage.py b/src/diffusers/models/controlnets/controlnet_qwenimage.py
index fa374285eec1..cfe7c159ad89 100644
--- a/src/diffusers/models/controlnets/controlnet_qwenimage.py
+++ b/src/diffusers/models/controlnets/controlnet_qwenimage.py
@@ -13,14 +13,19 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import (
+ BaseOutput,
+ apply_lora_scale,
+ deprecate,
+ logging,
+)
from ..attention import AttentionMixin
from ..cache_utils import CacheMixin
from ..controlnets.controlnet import zero_module
@@ -40,7 +45,7 @@
@dataclass
class QwenImageControlNetOutput(BaseOutput):
- controlnet_block_samples: Tuple[torch.Tensor]
+ controlnet_block_samples: tuple[torch.Tensor]
class QwenImageControlNetModel(
@@ -53,12 +58,12 @@ def __init__(
self,
patch_size: int = 2,
in_channels: int = 64,
- out_channels: Optional[int] = 16,
+ out_channels: int | None = 16,
num_layers: int = 60,
attention_head_dim: int = 128,
num_attention_heads: int = 24,
joint_attention_dim: int = 3584,
- axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+ axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
extra_condition_channels: int = 0, # for controlnet-inpainting
):
super().__init__()
@@ -123,6 +128,7 @@ def from_transformer(
return controlnet
+ @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -131,11 +137,11 @@ def forward(
encoder_hidden_states: torch.Tensor = None,
encoder_hidden_states_mask: torch.Tensor = None,
timestep: torch.LongTensor = None,
- img_shapes: Optional[List[Tuple[int, int, int]]] = None,
- txt_seq_lens: Optional[List[int]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ img_shapes: list[tuple[int, int, int]] | None = None,
+ txt_seq_lens: list[int] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+ ) -> torch.FloatTensor | Transformer2DModelOutput:
"""
The [`QwenImageControlNetModel`] forward method.
@@ -154,9 +160,9 @@ def forward(
(not just contiguous valid tokens followed by padding) since it's applied element-wise in attention.
timestep ( `torch.LongTensor`):
Used to indicate denoising step.
- img_shapes (`List[Tuple[int, int, int]]`, *optional*):
+ img_shapes (`list[tuple[int, int, int]]`, *optional*):
Image shapes for RoPE computation.
- txt_seq_lens (`List[int]`, *optional*):
+ txt_seq_lens (`list[int]`, *optional*):
**Deprecated**. Not needed anymore, we use `encoder_hidden_states` instead to infer text sequence
length.
joint_attention_kwargs (`dict`, *optional*):
@@ -181,20 +187,6 @@ def forward(
standard_warn=False,
)
- if joint_attention_kwargs is not None:
- joint_attention_kwargs = joint_attention_kwargs.copy()
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
hidden_states = self.img_in(hidden_states)
# add
@@ -256,10 +248,6 @@ def forward(
controlnet_block_samples = [sample * conditioning_scale for sample in controlnet_block_samples]
controlnet_block_samples = None if len(controlnet_block_samples) == 0 else controlnet_block_samples
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return controlnet_block_samples
@@ -276,7 +264,7 @@ class QwenImageMultiControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin, F
to be compatible with `QwenImageControlNetModel`.
Args:
- controlnets (`List[QwenImageControlNetModel]`):
+ controlnets (`list[QwenImageControlNetModel]`):
Provides additional conditioning to the unet during the denoising process. You must set multiple
`QwenImageControlNetModel` as a list.
"""
@@ -288,16 +276,16 @@ def __init__(self, controlnets):
def forward(
self,
hidden_states: torch.FloatTensor,
- controlnet_cond: List[torch.tensor],
- conditioning_scale: List[float],
+ controlnet_cond: list[torch.tensor],
+ conditioning_scale: list[float],
encoder_hidden_states: torch.Tensor = None,
encoder_hidden_states_mask: torch.Tensor = None,
timestep: torch.LongTensor = None,
- img_shapes: Optional[List[Tuple[int, int, int]]] = None,
- txt_seq_lens: Optional[List[int]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ img_shapes: list[tuple[int, int, int]] | None = None,
+ txt_seq_lens: list[int] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[QwenImageControlNetOutput, Tuple]:
+ ) -> QwenImageControlNetOutput | tuple:
if txt_seq_lens is not None:
deprecate(
"txt_seq_lens",
diff --git a/src/diffusers/models/controlnets/controlnet_sana.py b/src/diffusers/models/controlnets/controlnet_sana.py
index c71a8b326635..29e5591fa284 100644
--- a/src/diffusers/models/controlnets/controlnet_sana.py
+++ b/src/diffusers/models/controlnets/controlnet_sana.py
@@ -13,14 +13,14 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
from torch import nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, BaseOutput, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import BaseOutput, apply_lora_scale, logging
from ..attention import AttentionMixin
from ..embeddings import PatchEmbed, PixArtAlphaTextProjection
from ..modeling_outputs import Transformer2DModelOutput
@@ -35,7 +35,7 @@
@dataclass
class SanaControlNetOutput(BaseOutput):
- controlnet_block_samples: Tuple[torch.Tensor]
+ controlnet_block_samples: tuple[torch.Tensor]
class SanaControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin):
@@ -47,13 +47,13 @@ class SanaControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMi
def __init__(
self,
in_channels: int = 32,
- out_channels: Optional[int] = 32,
+ out_channels: int | None = 32,
num_attention_heads: int = 70,
attention_head_dim: int = 32,
num_layers: int = 7,
- num_cross_attention_heads: Optional[int] = 20,
- cross_attention_head_dim: Optional[int] = 112,
- cross_attention_dim: Optional[int] = 2240,
+ num_cross_attention_heads: int | None = 20,
+ cross_attention_head_dim: int | None = 112,
+ cross_attention_dim: int | None = 2240,
caption_channels: int = 2304,
mlp_ratio: float = 2.5,
dropout: float = 0.0,
@@ -62,7 +62,7 @@ def __init__(
patch_size: int = 1,
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-6,
- interpolation_scale: Optional[int] = None,
+ interpolation_scale: int | None = None,
) -> None:
super().__init__()
@@ -117,6 +117,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -124,26 +125,11 @@ def forward(
timestep: torch.LongTensor,
controlnet_cond: torch.Tensor,
conditioning_scale: float = 1.0,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput:
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
@@ -218,10 +204,6 @@ def forward(
block_res_sample = controlnet_block(block_res_sample)
controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
controlnet_block_res_samples = [sample * conditioning_scale for sample in controlnet_block_res_samples]
if not return_dict:
diff --git a/src/diffusers/models/controlnets/controlnet_sd3.py b/src/diffusers/models/controlnets/controlnet_sd3.py
index 08b86ff344eb..b8cb97adb41a 100644
--- a/src/diffusers/models/controlnets/controlnet_sd3.py
+++ b/src/diffusers/models/controlnets/controlnet_sd3.py
@@ -14,14 +14,14 @@
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ..attention import AttentionMixin, JointTransformerBlock
from ..attention_processor import Attention, FusedJointAttnProcessor2_0
from ..embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
@@ -36,7 +36,7 @@
@dataclass
class SD3ControlNetOutput(BaseOutput):
- controlnet_block_samples: Tuple[torch.Tensor]
+ controlnet_block_samples: tuple[torch.Tensor]
class SD3ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
@@ -69,7 +69,7 @@ class SD3ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMix
The maximum latent height/width of positional embeddings.
extra_conditioning_channels (`int`, defaults to `0`):
The number of extra channels to use for conditioning for patch embedding.
- dual_attention_layers (`Tuple[int, ...]`, defaults to `()`):
+ dual_attention_layers (`tuple[int, ...]`, defaults to `()`):
The number of dual-stream transformer blocks to use.
qk_norm (`str`, *optional*, defaults to `None`):
The normalization to use for query and key in the attention layer. If `None`, no normalization is used.
@@ -99,9 +99,9 @@ def __init__(
out_channels: int = 16,
pos_embed_max_size: int = 96,
extra_conditioning_channels: int = 0,
- dual_attention_layers: Tuple[int, ...] = (),
- qk_norm: Optional[str] = None,
- pos_embed_type: Optional[str] = "sincos",
+ dual_attention_layers: tuple[int, ...] = (),
+ qk_norm: str | None = None,
+ pos_embed_type: str | None = "sincos",
use_pos_embed: bool = True,
force_zeros_for_pooled_projection: bool = True,
):
@@ -175,7 +175,7 @@ def __init__(
self.gradient_checkpointing = False
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
- def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+ def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -269,6 +269,7 @@ def from_transformer(
return controlnet
+ @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -277,9 +278,9 @@ def forward(
encoder_hidden_states: torch.Tensor = None,
pooled_projections: torch.Tensor = None,
timestep: torch.LongTensor = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ ) -> torch.Tensor | Transformer2DModelOutput:
"""
The [`SD3Transformer2DModel`] forward method.
@@ -308,21 +309,6 @@ def forward(
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
- if joint_attention_kwargs is not None:
- joint_attention_kwargs = joint_attention_kwargs.copy()
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
if self.pos_embed is not None and hidden_states.ndim != 4:
raise ValueError("hidden_states must be 4D when pos_embed is used")
@@ -382,10 +368,6 @@ def forward(
# 6. scaling
controlnet_block_res_samples = [sample * conditioning_scale for sample in controlnet_block_res_samples]
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (controlnet_block_res_samples,)
@@ -400,7 +382,7 @@ class SD3MultiControlNetModel(ModelMixin):
compatible with `SD3ControlNetModel`.
Args:
- controlnets (`List[SD3ControlNetModel]`):
+ controlnets (`list[SD3ControlNetModel]`):
Provides additional conditioning to the unet during the denoising process. You must set multiple
`SD3ControlNetModel` as a list.
"""
@@ -412,14 +394,14 @@ def __init__(self, controlnets):
def forward(
self,
hidden_states: torch.Tensor,
- controlnet_cond: List[torch.tensor],
- conditioning_scale: List[float],
+ controlnet_cond: list[torch.tensor],
+ conditioning_scale: list[float],
pooled_projections: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
timestep: torch.LongTensor = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[SD3ControlNetOutput, Tuple]:
+ ) -> SD3ControlNetOutput | tuple:
for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)):
block_samples = controlnet(
hidden_states=hidden_states,
diff --git a/src/diffusers/models/controlnets/controlnet_sparsectrl.py b/src/diffusers/models/controlnets/controlnet_sparsectrl.py
index 8e7faf2d44b0..7da627fe6dd4 100644
--- a/src/diffusers/models/controlnets/controlnet_sparsectrl.py
+++ b/src/diffusers/models/controlnets/controlnet_sparsectrl.py
@@ -13,7 +13,7 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
from torch import nn
@@ -55,7 +55,7 @@ class SparseControlNetOutput(BaseOutput):
Output can be used to condition the original UNet's middle block activation.
"""
- down_block_res_samples: Tuple[torch.Tensor]
+ down_block_res_samples: tuple[torch.Tensor]
mid_block_res_sample: torch.Tensor
@@ -64,7 +64,7 @@ def __init__(
self,
conditioning_embedding_channels: int,
conditioning_channels: int = 3,
- block_out_channels: Tuple[int, ...] = (16, 32, 96, 256),
+ block_out_channels: tuple[int, ...] = (16, 32, 96, 256),
):
super().__init__()
@@ -110,7 +110,7 @@ class SparseControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOrigina
The frequency shift to apply to the time embedding.
down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
The tuple of downsample blocks to use.
- only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`):
+ only_cross_attention (`bool | tuple[bool]`, defaults to `False`):
block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, defaults to 2):
@@ -128,28 +128,28 @@ class SparseControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOrigina
The epsilon to use for the normalization.
cross_attention_dim (`int`, defaults to 1280):
The dimension of the cross attention features.
- transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
+ transformer_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 1):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
[`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
[`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
- transformer_layers_per_mid_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
+ transformer_layers_per_mid_block (`int` or `tuple[int]`, *optional*, defaults to 1):
The number of transformer layers to use in each layer in the middle block.
- attention_head_dim (`int` or `Tuple[int]`, defaults to 8):
+ attention_head_dim (`int` or `tuple[int]`, defaults to 8):
The dimension of the attention heads.
- num_attention_heads (`int` or `Tuple[int]`, *optional*):
+ num_attention_heads (`int` or `tuple[int]`, *optional*):
The number of heads to use for multi-head attention.
use_linear_projection (`bool`, defaults to `False`):
upcast_attention (`bool`, defaults to `False`):
resnet_time_scale_shift (`str`, defaults to `"default"`):
Time scale shift config for ResNet blocks (see `ResnetBlock2D`). Choose from `default` or `scale_shift`.
- conditioning_embedding_out_channels (`Tuple[int]`, defaults to `(16, 32, 96, 256)`):
+ conditioning_embedding_out_channels (`tuple[int]`, defaults to `(16, 32, 96, 256)`):
The tuple of output channel for each block in the `conditioning_embedding` layer.
global_pool_conditions (`bool`, defaults to `False`):
TODO(Patrick) - unused parameter
controlnet_conditioning_channel_order (`str`, defaults to `rgb`):
motion_max_seq_length (`int`, defaults to `32`):
The maximum sequence length to use in the motion module.
- motion_num_attention_heads (`int` or `Tuple[int]`, defaults to `8`):
+ motion_num_attention_heads (`int` or `tuple[int]`, defaults to `8`):
The number of heads to use in each attention layer of the motion module.
concat_conditioning_mask (`bool`, defaults to `True`):
use_simplified_condition_embedding (`bool`, defaults to `True`):
@@ -164,30 +164,30 @@ def __init__(
conditioning_channels: int = 4,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlockMotion",
"CrossAttnDownBlockMotion",
"CrossAttnDownBlockMotion",
"DownBlockMotion",
),
- only_cross_attention: Union[bool, Tuple[bool]] = False,
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+ only_cross_attention: bool | tuple[bool] = False,
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
layers_per_block: int = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
act_fn: str = "silu",
- norm_num_groups: Optional[int] = 32,
+ norm_num_groups: int | None = 32,
norm_eps: float = 1e-5,
cross_attention_dim: int = 768,
- transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
- transformer_layers_per_mid_block: Optional[Union[int, Tuple[int]]] = None,
- temporal_transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
- attention_head_dim: Union[int, Tuple[int, ...]] = 8,
- num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
+ transformer_layers_per_block: int | tuple[int, ...] = 1,
+ transformer_layers_per_mid_block: int | tuple[int] | None = None,
+ temporal_transformer_layers_per_block: int | tuple[int, ...] = 1,
+ attention_head_dim: int | tuple[int, ...] = 8,
+ num_attention_heads: int | tuple[int, ...] | None = None,
use_linear_projection: bool = False,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
- conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+ conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256),
global_pool_conditions: bool = False,
controlnet_conditioning_channel_order: str = "rgb",
motion_max_seq_length: int = 32,
@@ -389,7 +389,7 @@ def from_unet(
cls,
unet: UNet2DConditionModel,
controlnet_conditioning_channel_order: str = "rgb",
- conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+ conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256),
load_weights_from_unet: bool = True,
conditioning_channels: int = 3,
) -> "SparseControlNetModel":
@@ -465,7 +465,7 @@ def set_default_attn_processor(self):
self.set_attn_processor(processor)
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice
- def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None:
+ def set_attention_slice(self, slice_size: str | int | list[int]) -> None:
r"""
Enable sliced attention computation.
@@ -519,7 +519,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
# Recursively walk through all the children.
# Any children which exposes the set_attention_slice method
# gets the message
- def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]):
if hasattr(module, "set_attention_slice"):
module.set_attention_slice(slice_size.pop())
@@ -533,24 +533,24 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
controlnet_cond: torch.Tensor,
conditioning_scale: float = 1.0,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- conditioning_mask: Optional[torch.Tensor] = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ conditioning_mask: torch.Tensor | None = None,
guess_mode: bool = False,
return_dict: bool = True,
- ) -> Union[SparseControlNetOutput, Tuple[Tuple[torch.Tensor, ...], torch.Tensor]]:
+ ) -> SparseControlNetOutput | tuple[tuple[torch.Tensor, ...], torch.Tensor]:
"""
The [`SparseControlNetModel`] forward method.
Args:
sample (`torch.Tensor`):
The noisy input tensor.
- timestep (`Union[torch.Tensor, float, int]`):
+ timestep (`torch.Tensor | float | int`):
The number of timesteps to denoise an input.
encoder_hidden_states (`torch.Tensor`):
The encoder hidden states.
diff --git a/src/diffusers/models/controlnets/controlnet_union.py b/src/diffusers/models/controlnets/controlnet_union.py
index b4ee6536ca2f..8e434ba1f250 100644
--- a/src/diffusers/models/controlnets/controlnet_union.py
+++ b/src/diffusers/models/controlnets/controlnet_union.py
@@ -11,7 +11,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.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
from torch import nn
@@ -94,7 +94,7 @@ class ControlNetUnionModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginal
The frequency shift to apply to the time embedding.
down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
The tuple of downsample blocks to use.
- only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`):
+ only_cross_attention (`bool | tuple[bool]`, defaults to `False`):
block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, defaults to 2):
@@ -112,7 +112,7 @@ class ControlNetUnionModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginal
The epsilon to use for the normalization.
cross_attention_dim (`int`, defaults to 1280):
The dimension of the cross attention features.
- transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
+ transformer_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 1):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
[`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
[`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
@@ -122,7 +122,7 @@ class ControlNetUnionModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginal
encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
- attention_head_dim (`Union[int, Tuple[int]]`, defaults to 8):
+ attention_head_dim (`int | tuple[int]`, defaults to 8):
The dimension of the attention heads.
use_linear_projection (`bool`, defaults to `False`):
class_embed_type (`str`, *optional*, defaults to `None`):
@@ -156,36 +156,36 @@ def __init__(
conditioning_channels: int = 3,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- only_cross_attention: Union[bool, Tuple[bool]] = False,
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+ only_cross_attention: bool | tuple[bool] = False,
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
layers_per_block: int = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
act_fn: str = "silu",
- norm_num_groups: Optional[int] = 32,
+ norm_num_groups: int | None = 32,
norm_eps: float = 1e-5,
cross_attention_dim: int = 1280,
- transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
- encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
- attention_head_dim: Union[int, Tuple[int, ...]] = 8,
- num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
+ transformer_layers_per_block: int | tuple[int, ...] = 1,
+ encoder_hid_dim: int | None = None,
+ encoder_hid_dim_type: str | None = None,
+ attention_head_dim: int | tuple[int, ...] = 8,
+ num_attention_heads: int | tuple[int, ...] | None = None,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
- addition_time_embed_dim: Optional[int] = None,
- num_class_embeds: Optional[int] = None,
+ class_embed_type: str | None = None,
+ addition_embed_type: str | None = None,
+ addition_time_embed_dim: int | None = None,
+ num_class_embeds: int | None = None,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
- projection_class_embeddings_input_dim: Optional[int] = None,
+ projection_class_embeddings_input_dim: int | None = None,
controlnet_conditioning_channel_order: str = "rgb",
- conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (48, 96, 192, 384),
+ conditioning_embedding_out_channels: tuple[int, ...] | None = (48, 96, 192, 384),
global_pool_conditions: bool = False,
addition_embed_type_num_heads: int = 64,
num_control_type: int = 6,
@@ -390,7 +390,7 @@ def from_unet(
cls,
unet: UNet2DConditionModel,
controlnet_conditioning_channel_order: str = "rgb",
- conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+ conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256),
load_weights_from_unet: bool = True,
):
r"""
@@ -472,7 +472,7 @@ def set_default_attn_processor(self):
self.set_attn_processor(processor)
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice
- def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None:
+ def set_attention_slice(self, slice_size: str | int | list[int]) -> None:
r"""
Enable sliced attention computation.
@@ -526,7 +526,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
# Recursively walk through all the children.
# Any children which exposes the set_attention_slice method
# gets the message
- def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]):
if hasattr(module, "set_attention_slice"):
module.set_attention_slice(slice_size.pop())
@@ -540,37 +540,37 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- controlnet_cond: List[torch.Tensor],
+ controlnet_cond: list[torch.Tensor],
control_type: torch.Tensor,
- control_type_idx: List[int],
- conditioning_scale: Union[float, List[float]] = 1.0,
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ control_type_idx: list[int],
+ conditioning_scale: float | list[float] = 1.0,
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
from_multi: bool = False,
guess_mode: bool = False,
return_dict: bool = True,
- ) -> Union[ControlNetOutput, Tuple[Tuple[torch.Tensor, ...], torch.Tensor]]:
+ ) -> ControlNetOutput | tuple[tuple[torch.Tensor, ...], torch.Tensor]:
"""
The [`ControlNetUnionModel`] forward method.
Args:
sample (`torch.Tensor`):
The noisy input tensor.
- timestep (`Union[torch.Tensor, float, int]`):
+ timestep (`torch.Tensor | float | int`):
The number of timesteps to denoise an input.
encoder_hidden_states (`torch.Tensor`):
The encoder hidden states.
- controlnet_cond (`List[torch.Tensor]`):
+ controlnet_cond (`list[torch.Tensor]`):
The conditional input tensors.
control_type (`torch.Tensor`):
A tensor of shape `(batch, num_control_type)` with values `0` or `1` depending on whether the control
type is used.
- control_type_idx (`List[int]`):
+ control_type_idx (`list[int]`):
The indices of `control_type`.
conditioning_scale (`float`, defaults to `1.0`):
The scale factor for ControlNet outputs.
diff --git a/src/diffusers/models/controlnets/controlnet_xs.py b/src/diffusers/models/controlnets/controlnet_xs.py
index 119492b0fac4..6221d4878de9 100644
--- a/src/diffusers/models/controlnets/controlnet_xs.py
+++ b/src/diffusers/models/controlnets/controlnet_xs.py
@@ -13,7 +13,7 @@
# limitations under the License.
from dataclasses import dataclass
from math import gcd
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
from torch import Tensor, nn
@@ -71,8 +71,8 @@ def __init__(
resnets: nn.ModuleList,
base_to_ctrl: nn.ModuleList,
ctrl_to_base: nn.ModuleList,
- attentions: Optional[nn.ModuleList] = None,
- downsampler: Optional[nn.Conv2d] = None,
+ attentions: nn.ModuleList | None = None,
+ downsampler: nn.Conv2d | None = None,
):
super().__init__()
self.resnets = resnets
@@ -107,14 +107,14 @@ def get_down_block_adapter(
ctrl_in_channels: int,
ctrl_out_channels: int,
temb_channels: int,
- max_norm_num_groups: Optional[int] = 32,
+ max_norm_num_groups: int | None = 32,
has_crossattn=True,
- transformer_layers_per_block: Optional[Union[int, Tuple[int]]] = 1,
- num_attention_heads: Optional[int] = 1,
- cross_attention_dim: Optional[int] = 1024,
+ transformer_layers_per_block: int | tuple[int] | None = 1,
+ num_attention_heads: int | None = 1,
+ cross_attention_dim: int | None = 1024,
add_downsample: bool = True,
- upcast_attention: Optional[bool] = False,
- use_linear_projection: Optional[bool] = True,
+ upcast_attention: bool | None = False,
+ use_linear_projection: bool | None = True,
):
num_layers = 2 # only support sd + sdxl
@@ -195,11 +195,11 @@ def get_down_block_adapter(
def get_mid_block_adapter(
base_channels: int,
ctrl_channels: int,
- temb_channels: Optional[int] = None,
- max_norm_num_groups: Optional[int] = 32,
+ temb_channels: int | None = None,
+ max_norm_num_groups: int | None = 32,
transformer_layers_per_block: int = 1,
- num_attention_heads: Optional[int] = 1,
- cross_attention_dim: Optional[int] = 1024,
+ num_attention_heads: int | None = 1,
+ cross_attention_dim: int | None = 1024,
upcast_attention: bool = False,
use_linear_projection: bool = True,
):
@@ -230,7 +230,7 @@ def get_mid_block_adapter(
def get_up_block_adapter(
out_channels: int,
prev_output_channel: int,
- ctrl_skip_channels: List[int],
+ ctrl_skip_channels: list[int],
):
ctrl_to_base = []
num_layers = 3 # only support sd + sdxl
@@ -278,7 +278,7 @@ class ControlNetXSAdapter(ModelMixin, AttentionMixin, ConfigMixin):
The tuple of downsample blocks to use.
sample_size (`int`, defaults to 96):
Height and width of input/output sample.
- transformer_layers_per_block (`Union[int, Tuple[int]]`, defaults to 1):
+ transformer_layers_per_block (`int | tuple[int]`, defaults to 1):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
[`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
upcast_attention (`bool`, defaults to `True`):
@@ -293,21 +293,21 @@ def __init__(
self,
conditioning_channels: int = 3,
conditioning_channel_order: str = "rgb",
- conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256),
+ conditioning_embedding_out_channels: tuple[int] = (16, 32, 96, 256),
time_embedding_mix: float = 1.0,
learn_time_embedding: bool = False,
- num_attention_heads: Union[int, Tuple[int]] = 4,
- block_out_channels: Tuple[int, ...] = (4, 8, 16, 16),
- base_block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+ num_attention_heads: int | tuple[int] = 4,
+ block_out_channels: tuple[int] = (4, 8, 16, 16),
+ base_block_out_channels: tuple[int] = (320, 640, 1280, 1280),
cross_attention_dim: int = 1024,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- sample_size: Optional[int] = 96,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ sample_size: int | None = 96,
+ transformer_layers_per_block: int | tuple[int] = 1,
upcast_attention: bool = True,
max_norm_num_groups: int = 32,
use_linear_projection: bool = True,
@@ -429,14 +429,14 @@ def __init__(
def from_unet(
cls,
unet: UNet2DConditionModel,
- size_ratio: Optional[float] = None,
- block_out_channels: Optional[List[int]] = None,
- num_attention_heads: Optional[List[int]] = None,
+ size_ratio: float | None = None,
+ block_out_channels: list[int] | None = None,
+ num_attention_heads: list[int] | None = None,
learn_time_embedding: bool = False,
time_embedding_mix: int = 1.0,
conditioning_channels: int = 3,
conditioning_channel_order: str = "rgb",
- conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256),
+ conditioning_embedding_out_channels: tuple[int] = (16, 32, 96, 256),
):
r"""
Instantiate a [`ControlNetXSAdapter`] from a [`UNet2DConditionModel`].
@@ -447,9 +447,9 @@ def from_unet(
size_ratio (float, *optional*, defaults to `None`):
When given, block_out_channels is set to a fraction of the base model's block_out_channels. Either this
or `block_out_channels` must be given.
- block_out_channels (`List[int]`, *optional*, defaults to `None`):
+ block_out_channels (`list[int]`, *optional*, defaults to `None`):
Down blocks output channels in control model. Either this or `size_ratio` must be given.
- num_attention_heads (`List[int]`, *optional*, defaults to `None`):
+ num_attention_heads (`list[int]`, *optional*, defaults to `None`):
The dimension of the attention heads. The naming seems a bit confusing and it is, see
https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 for why.
learn_time_embedding (`bool`, defaults to `False`):
@@ -461,7 +461,7 @@ def from_unet(
Number of channels of conditioning input (e.g. an image)
conditioning_channel_order (`str`, defaults to `"rgb"`):
The channel order of conditional image. Will convert to `rgb` if it's `bgr`.
- conditioning_embedding_out_channels (`Tuple[int]`, defaults to `(16, 32, 96, 256)`):
+ conditioning_embedding_out_channels (`tuple[int]`, defaults to `(16, 32, 96, 256)`):
The tuple of output channel for each block in the `controlnet_cond_embedding` layer.
"""
@@ -528,38 +528,33 @@ class UNetControlNetXSModel(ModelMixin, AttentionMixin, ConfigMixin):
def __init__(
self,
# unet configs
- sample_size: Optional[int] = 96,
- down_block_types: Tuple[str, ...] = (
+ sample_size: int | None = 96,
+ down_block_types: tuple[str] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- up_block_types: Tuple[str, ...] = (
- "UpBlock2D",
- "CrossAttnUpBlock2D",
- "CrossAttnUpBlock2D",
- "CrossAttnUpBlock2D",
- ),
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
- norm_num_groups: Optional[int] = 32,
- cross_attention_dim: Union[int, Tuple[int]] = 1024,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
- num_attention_heads: Union[int, Tuple[int]] = 8,
- addition_embed_type: Optional[str] = None,
- addition_time_embed_dim: Optional[int] = None,
+ up_block_types: tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
+ block_out_channels: tuple[int] = (320, 640, 1280, 1280),
+ norm_num_groups: int | None = 32,
+ cross_attention_dim: int | tuple[int] = 1024,
+ transformer_layers_per_block: int | tuple[int] = 1,
+ num_attention_heads: int | tuple[int] = 8,
+ addition_embed_type: str | None = None,
+ addition_time_embed_dim: int | None = None,
upcast_attention: bool = True,
use_linear_projection: bool = True,
- time_cond_proj_dim: Optional[int] = None,
- projection_class_embeddings_input_dim: Optional[int] = None,
+ time_cond_proj_dim: int | None = None,
+ projection_class_embeddings_input_dim: int | None = None,
# additional controlnet configs
time_embedding_mix: float = 1.0,
ctrl_conditioning_channels: int = 3,
- ctrl_conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256),
+ ctrl_conditioning_embedding_out_channels: tuple[int] = (16, 32, 96, 256),
ctrl_conditioning_channel_order: str = "rgb",
ctrl_learn_time_embedding: bool = False,
- ctrl_block_out_channels: Tuple[int, ...] = (4, 8, 16, 16),
- ctrl_num_attention_heads: Union[int, Tuple[int]] = 4,
+ ctrl_block_out_channels: tuple[int] = (4, 8, 16, 16),
+ ctrl_num_attention_heads: int | tuple[int] = 4,
ctrl_max_norm_num_groups: int = 32,
):
super().__init__()
@@ -724,11 +719,11 @@ def __init__(
def from_unet(
cls,
unet: UNet2DConditionModel,
- controlnet: Optional[ControlNetXSAdapter] = None,
- size_ratio: Optional[float] = None,
- ctrl_block_out_channels: Optional[List[float]] = None,
- time_embedding_mix: Optional[float] = None,
- ctrl_optional_kwargs: Optional[Dict] = None,
+ controlnet: ControlNetXSAdapter | None = None,
+ size_ratio: float | None = None,
+ ctrl_block_out_channels: list[float] | None = None,
+ time_embedding_mix: float | None = None,
+ ctrl_optional_kwargs: dict | None = None,
):
r"""
Instantiate a [`UNetControlNetXSModel`] from a [`UNet2DConditionModel`] and an optional [`ControlNetXSAdapter`]
@@ -742,7 +737,7 @@ def from_unet(
adapter will be created.
size_ratio (float, *optional*, defaults to `None`):
Used to construct the controlnet if none is given. See [`ControlNetXSAdapter.from_unet`] for details.
- ctrl_block_out_channels (`List[int]`, *optional*, defaults to `None`):
+ ctrl_block_out_channels (`list[int]`, *optional*, defaults to `None`):
Used to construct the controlnet if none is given. See [`ControlNetXSAdapter.from_unet`] for details,
where this parameter is called `block_out_channels`.
time_embedding_mix (`float`, *optional*, defaults to None):
@@ -953,25 +948,25 @@ def unfuse_qkv_projections(self):
def forward(
self,
sample: Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- controlnet_cond: Optional[torch.Tensor] = None,
- conditioning_scale: Optional[float] = 1.0,
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+ controlnet_cond: torch.Tensor | None = None,
+ conditioning_scale: float | None = 1.0,
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
return_dict: bool = True,
apply_control: bool = True,
- ) -> Union[ControlNetXSOutput, Tuple]:
+ ) -> ControlNetXSOutput | tuple:
"""
The [`ControlNetXSModel`] forward method.
Args:
sample (`Tensor`):
The noisy input tensor.
- timestep (`Union[torch.Tensor, float, int]`):
+ timestep (`torch.Tensor | float | int`):
The number of timesteps to denoise an input.
encoder_hidden_states (`torch.Tensor`):
The encoder hidden states.
@@ -1166,13 +1161,13 @@ def __init__(
norm_num_groups: int = 32,
ctrl_max_norm_num_groups: int = 32,
has_crossattn=True,
- transformer_layers_per_block: Optional[Union[int, Tuple[int]]] = 1,
- base_num_attention_heads: Optional[int] = 1,
- ctrl_num_attention_heads: Optional[int] = 1,
- cross_attention_dim: Optional[int] = 1024,
+ transformer_layers_per_block: int | tuple[int] | None = 1,
+ base_num_attention_heads: int | None = 1,
+ ctrl_num_attention_heads: int | None = 1,
+ cross_attention_dim: int | None = 1024,
add_downsample: bool = True,
- upcast_attention: Optional[bool] = False,
- use_linear_projection: Optional[bool] = True,
+ upcast_attention: bool | None = False,
+ use_linear_projection: bool | None = True,
):
super().__init__()
base_resnets = []
@@ -1361,14 +1356,14 @@ def forward(
self,
hidden_states_base: Tensor,
temb: Tensor,
- encoder_hidden_states: Optional[Tensor] = None,
- hidden_states_ctrl: Optional[Tensor] = None,
- conditioning_scale: Optional[float] = 1.0,
- attention_mask: Optional[Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[Tensor] = None,
+ encoder_hidden_states: Tensor | None = None,
+ hidden_states_ctrl: Tensor | None = None,
+ conditioning_scale: float | None = 1.0,
+ attention_mask: Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: Tensor | None = None,
apply_control: bool = True,
- ) -> Tuple[Tensor, Tensor, Tuple[Tensor, ...], Tuple[Tensor, ...]]:
+ ) -> tuple[Tensor, Tensor, tuple[Tensor, ...], tuple[Tensor, ...]]:
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
@@ -1455,15 +1450,15 @@ def __init__(
self,
base_channels: int,
ctrl_channels: int,
- temb_channels: Optional[int] = None,
+ temb_channels: int | None = None,
norm_num_groups: int = 32,
ctrl_max_norm_num_groups: int = 32,
transformer_layers_per_block: int = 1,
- base_num_attention_heads: Optional[int] = 1,
- ctrl_num_attention_heads: Optional[int] = 1,
- cross_attention_dim: Optional[int] = 1024,
+ base_num_attention_heads: int | None = 1,
+ ctrl_num_attention_heads: int | None = 1,
+ cross_attention_dim: int | None = 1024,
upcast_attention: bool = False,
- use_linear_projection: Optional[bool] = True,
+ use_linear_projection: bool | None = True,
):
super().__init__()
@@ -1568,13 +1563,13 @@ def forward(
hidden_states_base: Tensor,
temb: Tensor,
encoder_hidden_states: Tensor,
- hidden_states_ctrl: Optional[Tensor] = None,
- conditioning_scale: Optional[float] = 1.0,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- attention_mask: Optional[Tensor] = None,
- encoder_attention_mask: Optional[Tensor] = None,
+ hidden_states_ctrl: Tensor | None = None,
+ conditioning_scale: float | None = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ attention_mask: Tensor | None = None,
+ encoder_attention_mask: Tensor | None = None,
apply_control: bool = True,
- ) -> Tuple[Tensor, Tensor]:
+ ) -> tuple[Tensor, Tensor]:
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
@@ -1606,17 +1601,17 @@ def __init__(
in_channels: int,
out_channels: int,
prev_output_channel: int,
- ctrl_skip_channels: List[int],
+ ctrl_skip_channels: list[int],
temb_channels: int,
norm_num_groups: int = 32,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
has_crossattn=True,
transformer_layers_per_block: int = 1,
num_attention_heads: int = 1,
cross_attention_dim: int = 1024,
add_upsample: bool = True,
upcast_attention: bool = False,
- use_linear_projection: Optional[bool] = True,
+ use_linear_projection: bool | None = True,
):
super().__init__()
resnets = []
@@ -1751,15 +1746,15 @@ def freeze_base_params(self) -> None:
def forward(
self,
hidden_states: Tensor,
- res_hidden_states_tuple_base: Tuple[Tensor, ...],
- res_hidden_states_tuple_ctrl: Tuple[Tensor, ...],
+ res_hidden_states_tuple_base: tuple[Tensor, ...],
+ res_hidden_states_tuple_ctrl: tuple[Tensor, ...],
temb: Tensor,
- encoder_hidden_states: Optional[Tensor] = None,
- conditioning_scale: Optional[float] = 1.0,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- attention_mask: Optional[Tensor] = None,
- upsample_size: Optional[int] = None,
- encoder_attention_mask: Optional[Tensor] = None,
+ encoder_hidden_states: Tensor | None = None,
+ conditioning_scale: float | None = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ attention_mask: Tensor | None = None,
+ upsample_size: int | None = None,
+ encoder_attention_mask: Tensor | None = None,
apply_control: bool = True,
) -> Tensor:
if cross_attention_kwargs is not None:
diff --git a/src/diffusers/models/controlnets/controlnet_z_image.py b/src/diffusers/models/controlnets/controlnet_z_image.py
index 3f79ec925419..85fa0d365547 100644
--- a/src/diffusers/models/controlnets/controlnet_z_image.py
+++ b/src/diffusers/models/controlnets/controlnet_z_image.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import List, Literal, Optional, Tuple
+from typing import Literal
import torch
import torch.nn as nn
@@ -94,9 +94,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- freqs_cis: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ freqs_cis: torch.Tensor | None = None,
) -> torch.Tensor:
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
@@ -234,10 +234,10 @@ def forward(
x: torch.Tensor,
attn_mask: torch.Tensor,
freqs_cis: torch.Tensor,
- adaln_input: Optional[torch.Tensor] = None,
- noise_mask: Optional[torch.Tensor] = None,
- adaln_noisy: Optional[torch.Tensor] = None,
- adaln_clean: Optional[torch.Tensor] = None,
+ adaln_input: torch.Tensor | None = None,
+ noise_mask: torch.Tensor | None = None,
+ adaln_noisy: torch.Tensor | None = None,
+ adaln_clean: torch.Tensor | None = None,
):
if self.modulation:
seq_len = x.shape[1]
@@ -291,8 +291,8 @@ class RopeEmbedder:
def __init__(
self,
theta: float = 256.0,
- axes_dims: List[int] = (16, 56, 56),
- axes_lens: List[int] = (64, 128, 128),
+ axes_dims: list[int] = (16, 56, 56),
+ axes_lens: list[int] = (64, 128, 128),
):
self.theta = theta
self.axes_dims = axes_dims
@@ -301,7 +301,7 @@ def __init__(
self.freqs_cis = None
@staticmethod
- def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0):
+ def precompute_freqs_cis(dim: list[int], end: list[int], theta: float = 256.0):
with torch.device("cpu"):
freqs_cis = []
for i, (d, e) in enumerate(zip(dim, end)):
@@ -389,7 +389,7 @@ def forward(
x: torch.Tensor,
attn_mask: torch.Tensor,
freqs_cis: torch.Tensor,
- adaln_input: Optional[torch.Tensor] = None,
+ adaln_input: torch.Tensor | None = None,
):
# Control
if self.block_id == 0:
@@ -435,10 +435,10 @@ class ZImageControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrigi
@register_to_config
def __init__(
self,
- control_layers_places: List[int] = None,
- control_refiner_layers_places: List[int] = None,
+ control_layers_places: list[int] = None,
+ control_refiner_layers_places: list[int] = None,
control_in_dim=None,
- add_control_noise_refiner: Optional[Literal["control_layers", "control_noise_refiner"]] = None,
+ add_control_noise_refiner: Literal["control_layers", "control_noise_refiner"] | None = None,
all_patch_size=(2,),
all_f_patch_size=(1,),
dim=3840,
@@ -505,15 +505,15 @@ def __init__(
]
)
- self.t_scale: Optional[float] = None
- self.t_embedder: Optional[TimestepEmbedder] = None
- self.all_x_embedder: Optional[nn.ModuleDict] = None
- self.cap_embedder: Optional[nn.Sequential] = None
- self.rope_embedder: Optional[RopeEmbedder] = None
- self.noise_refiner: Optional[nn.ModuleList] = None
- self.context_refiner: Optional[nn.ModuleList] = None
- self.x_pad_token: Optional[nn.Parameter] = None
- self.cap_pad_token: Optional[nn.Parameter] = None
+ self.t_scale: float | None = None
+ self.t_embedder: TimestepEmbedder | None = None
+ self.all_x_embedder: nn.ModuleDict | None = None
+ self.cap_embedder: nn.Sequential | None = None
+ self.rope_embedder: RopeEmbedder | None = None
+ self.noise_refiner: nn.ModuleList | None = None
+ self.context_refiner: nn.ModuleList | None = None
+ self.x_pad_token: nn.Parameter | None = None
+ self.cap_pad_token: nn.Parameter | None = None
@classmethod
def from_transformer(cls, controlnet, transformer):
@@ -551,10 +551,10 @@ def _patchify_image(self, image: torch.Tensor, patch_size: int, f_patch_size: in
def _pad_with_ids(
self,
feat: torch.Tensor,
- pos_grid_size: Tuple,
- pos_start: Tuple,
+ pos_grid_size: tuple,
+ pos_start: tuple,
device: torch.device,
- noise_mask_val: Optional[int] = None,
+ noise_mask_val: int | None = None,
):
"""Pad feature to SEQ_MULTI_OF, create position IDs and pad mask."""
ori_len = len(feat)
@@ -587,7 +587,7 @@ def _pad_with_ids(
# Copied from diffusers.models.transformers.transformer_z_image.ZImageTransformer2DModel.patchify_and_embed
def patchify_and_embed(
- self, all_image: List[torch.Tensor], all_cap_feats: List[torch.Tensor], patch_size: int, f_patch_size: int
+ self, all_image: list[torch.Tensor], all_cap_feats: list[torch.Tensor], patch_size: int, f_patch_size: int
):
"""Patchify for basic mode: single image per batch item."""
device = all_image[0].device
@@ -625,7 +625,7 @@ def patchify_and_embed(
def patchify(
self,
- all_image: List[torch.Tensor],
+ all_image: list[torch.Tensor],
patch_size: int,
f_patch_size: int,
):
@@ -653,10 +653,10 @@ def patchify(
def forward(
self,
- x: List[torch.Tensor],
+ x: list[torch.Tensor],
t,
- cap_feats: List[torch.Tensor],
- control_context: List[torch.Tensor],
+ cap_feats: list[torch.Tensor],
+ control_context: list[torch.Tensor],
conditioning_scale: float = 1.0,
patch_size=2,
f_patch_size=1,
diff --git a/src/diffusers/models/controlnets/multicontrolnet.py b/src/diffusers/models/controlnets/multicontrolnet.py
index 87a952294997..705c59c0f925 100644
--- a/src/diffusers/models/controlnets/multicontrolnet.py
+++ b/src/diffusers/models/controlnets/multicontrolnet.py
@@ -1,5 +1,5 @@
import os
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from torch import nn
@@ -20,30 +20,30 @@ class MultiControlNetModel(ModelMixin):
compatible with `ControlNetModel`.
Args:
- controlnets (`List[ControlNetModel]`):
+ controlnets (`list[ControlNetModel]`):
Provides additional conditioning to the unet during the denoising process. You must set multiple
`ControlNetModel` as a list.
"""
- def __init__(self, controlnets: Union[List[ControlNetModel], Tuple[ControlNetModel]]):
+ def __init__(self, controlnets: list[ControlNetModel] | tuple[ControlNetModel]):
super().__init__()
self.nets = nn.ModuleList(controlnets)
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- controlnet_cond: List[torch.tensor],
- conditioning_scale: List[float],
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ controlnet_cond: list[torch.tensor],
+ conditioning_scale: list[float],
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guess_mode: bool = False,
return_dict: bool = True,
- ) -> Union[ControlNetOutput, Tuple]:
+ ) -> ControlNetOutput | tuple:
for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)):
down_samples, mid_sample = controlnet(
sample=sample,
@@ -74,11 +74,11 @@ def forward(
def save_pretrained(
self,
- save_directory: Union[str, os.PathLike],
+ save_directory: str | os.PathLike,
is_main_process: bool = True,
save_function: Callable = None,
safe_serialization: bool = True,
- variant: Optional[str] = None,
+ variant: str | None = None,
):
"""
Save a model and its configuration file to a directory, so that it can be re-loaded using the
@@ -111,7 +111,7 @@ def save_pretrained(
)
@classmethod
- def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]], **kwargs):
+ def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwargs):
r"""
Instantiate a pretrained MultiControlNet model from multiple pre-trained controlnet models.
@@ -134,7 +134,7 @@ def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]
Override the default `torch.dtype` and load the model under this dtype.
output_loading_info(`bool`, *optional*, defaults to `False`):
Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
- device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+ device_map (`str` or `dict[str, int | str | torch.device]`, *optional*):
A map that specifies where each submodule should go. It doesn't need to be refined to each
parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
same device.
diff --git a/src/diffusers/models/controlnets/multicontrolnet_union.py b/src/diffusers/models/controlnets/multicontrolnet_union.py
index d5506dc186e3..98552f99623a 100644
--- a/src/diffusers/models/controlnets/multicontrolnet_union.py
+++ b/src/diffusers/models/controlnets/multicontrolnet_union.py
@@ -1,5 +1,5 @@
import os
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from torch import nn
@@ -21,32 +21,32 @@ class MultiControlNetUnionModel(ModelMixin):
be compatible with `ControlNetUnionModel`.
Args:
- controlnets (`List[ControlNetUnionModel]`):
+ controlnets (`list[ControlNetUnionModel]`):
Provides additional conditioning to the unet during the denoising process. You must set multiple
`ControlNetUnionModel` as a list.
"""
- def __init__(self, controlnets: Union[List[ControlNetUnionModel], Tuple[ControlNetUnionModel]]):
+ def __init__(self, controlnets: list[ControlNetUnionModel] | tuple[ControlNetUnionModel]):
super().__init__()
self.nets = nn.ModuleList(controlnets)
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- controlnet_cond: List[torch.tensor],
- control_type: List[torch.Tensor],
- control_type_idx: List[List[int]],
- conditioning_scale: List[float],
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ controlnet_cond: list[torch.tensor],
+ control_type: list[torch.Tensor],
+ control_type_idx: list[list[int]],
+ conditioning_scale: list[float],
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guess_mode: bool = False,
return_dict: bool = True,
- ) -> Union[ControlNetOutput, Tuple]:
+ ) -> ControlNetOutput | tuple:
down_block_res_samples, mid_block_res_sample = None, None
for i, (image, ctype, ctype_idx, scale, controlnet) in enumerate(
zip(controlnet_cond, control_type, control_type_idx, conditioning_scale, self.nets)
@@ -86,11 +86,11 @@ def forward(
# Copied from diffusers.models.controlnets.multicontrolnet.MultiControlNetModel.save_pretrained with ControlNet->ControlNetUnion
def save_pretrained(
self,
- save_directory: Union[str, os.PathLike],
+ save_directory: str | os.PathLike,
is_main_process: bool = True,
save_function: Callable = None,
safe_serialization: bool = True,
- variant: Optional[str] = None,
+ variant: str | None = None,
):
"""
Save a model and its configuration file to a directory, so that it can be re-loaded using the
@@ -124,7 +124,7 @@ def save_pretrained(
@classmethod
# Copied from diffusers.models.controlnets.multicontrolnet.MultiControlNetModel.from_pretrained with ControlNet->ControlNetUnion
- def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]], **kwargs):
+ def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwargs):
r"""
Instantiate a pretrained MultiControlNetUnion model from multiple pre-trained controlnet models.
@@ -147,7 +147,7 @@ def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]
Override the default `torch.dtype` and load the model under this dtype.
output_loading_info(`bool`, *optional*, defaults to `False`):
Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
- device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+ device_map (`str` or `dict[str, int | str | torch.device]`, *optional*):
A map that specifies where each submodule should go. It doesn't need to be refined to each
parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
same device.
diff --git a/src/diffusers/models/downsampling.py b/src/diffusers/models/downsampling.py
index 505816422b2a..871c0ed7ddf7 100644
--- a/src/diffusers/models/downsampling.py
+++ b/src/diffusers/models/downsampling.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple
-
import torch
import torch.nn as nn
import torch.nn.functional as F
@@ -43,7 +41,7 @@ def __init__(
self,
channels: int,
use_conv: bool = False,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
padding: int = 1,
name: str = "conv",
):
@@ -86,7 +84,7 @@ def __init__(
self,
channels: int,
use_conv: bool = False,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
padding: int = 1,
name: str = "conv",
kernel_size=3,
@@ -165,10 +163,10 @@ class FirDownsample2D(nn.Module):
def __init__(
self,
- channels: Optional[int] = None,
- out_channels: Optional[int] = None,
+ channels: int | None = None,
+ out_channels: int | None = None,
use_conv: bool = False,
- fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1),
+ fir_kernel: tuple[int, int, int, int] = (1, 3, 3, 1),
):
super().__init__()
out_channels = out_channels if out_channels else channels
@@ -181,8 +179,8 @@ def __init__(
def _downsample_2d(
self,
hidden_states: torch.Tensor,
- weight: Optional[torch.Tensor] = None,
- kernel: Optional[torch.Tensor] = None,
+ weight: torch.Tensor | None = None,
+ kernel: torch.Tensor | None = None,
factor: int = 2,
gain: float = 1,
) -> torch.Tensor:
@@ -355,7 +353,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
def downsample_2d(
hidden_states: torch.Tensor,
- kernel: Optional[torch.Tensor] = None,
+ kernel: torch.Tensor | None = None,
factor: int = 2,
gain: float = 1,
) -> torch.Tensor:
diff --git a/src/diffusers/models/embeddings.py b/src/diffusers/models/embeddings.py
index 37fc412adcc3..bcd192c1f166 100644
--- a/src/diffusers/models/embeddings.py
+++ b/src/diffusers/models/embeddings.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import math
-from typing import List, Optional, Tuple, Union
import numpy as np
import torch
@@ -80,11 +79,11 @@ def get_timestep_embedding(
def get_3d_sincos_pos_embed(
embed_dim: int,
- spatial_size: Union[int, Tuple[int, int]],
+ spatial_size: int | tuple[int, int],
temporal_size: int,
spatial_interpolation_scale: float = 1.0,
temporal_interpolation_scale: float = 1.0,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
output_type: str = "np",
) -> torch.Tensor:
r"""
@@ -93,7 +92,7 @@ def get_3d_sincos_pos_embed(
Args:
embed_dim (`int`):
The embedding dimension of inputs. It must be divisible by 16.
- spatial_size (`int` or `Tuple[int, int]`):
+ spatial_size (`int` or `tuple[int, int]`):
The spatial dimension of positional embeddings. If an integer is provided, the same size is applied to both
spatial dimensions (height and width).
temporal_size (`int`):
@@ -154,7 +153,7 @@ def get_3d_sincos_pos_embed(
def _get_3d_sincos_pos_embed_np(
embed_dim: int,
- spatial_size: Union[int, Tuple[int, int]],
+ spatial_size: int | tuple[int, int],
temporal_size: int,
spatial_interpolation_scale: float = 1.0,
temporal_interpolation_scale: float = 1.0,
@@ -165,7 +164,7 @@ def _get_3d_sincos_pos_embed_np(
Args:
embed_dim (`int`):
The embedding dimension of inputs. It must be divisible by 16.
- spatial_size (`int` or `Tuple[int, int]`):
+ spatial_size (`int` or `tuple[int, int]`):
The spatial dimension of positional embeddings. If an integer is provided, the same size is applied to both
spatial dimensions (height and width).
temporal_size (`int`):
@@ -225,7 +224,7 @@ def get_2d_sincos_pos_embed(
extra_tokens=0,
interpolation_scale=1.0,
base_size=16,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
output_type: str = "np",
):
"""
@@ -609,10 +608,10 @@ def forward(self, x, freqs_cis):
Patchifies and embeds the input tensor(s).
Args:
- x (List[torch.Tensor] | torch.Tensor): The input tensor(s) to be patchified and embedded.
+ x (list[torch.Tensor] | torch.Tensor): The input tensor(s) to be patchified and embedded.
Returns:
- Tuple[torch.Tensor, torch.Tensor, List[Tuple[int, int]], torch.Tensor]: A tuple containing the patchified
+ tuple[torch.Tensor, torch.Tensor, list[tuple[int, int]], torch.Tensor]: A tuple containing the patchified
and embedded tensor(s), the mask indicating the valid patches, the original image size(s), and the
frequency tensor(s).
"""
@@ -642,7 +641,7 @@ class CogVideoXPatchEmbed(nn.Module):
def __init__(
self,
patch_size: int = 2,
- patch_size_t: Optional[int] = None,
+ patch_size_t: int | None = None,
in_channels: int = 16,
embed_dim: int = 1920,
text_embed_dim: int = 4096,
@@ -689,7 +688,7 @@ def __init__(
self.register_buffer("pos_embedding", pos_embedding, persistent=persistent)
def _get_positional_embeddings(
- self, sample_height: int, sample_width: int, sample_frames: int, device: Optional[torch.device] = None
+ self, sample_height: int, sample_width: int, sample_frames: int, device: torch.device | None = None
) -> torch.Tensor:
post_patch_height = sample_height // self.patch_size
post_patch_width = sample_width // self.patch_size
@@ -836,18 +835,18 @@ def get_3d_rotary_pos_embed(
theta: int = 10000,
use_real: bool = True,
grid_type: str = "linspace",
- max_size: Optional[Tuple[int, int]] = None,
- device: Optional[torch.device] = None,
-) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ max_size: tuple[int, int] | None = None,
+ device: torch.device | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""
RoPE for video tokens with 3D structure.
Args:
embed_dim: (`int`):
The embedding dimension size, corresponding to hidden_size_head.
- crops_coords (`Tuple[int]`):
+ crops_coords (`tuple[int]`):
The top-left and bottom-right coordinates of the crop.
- grid_size (`Tuple[int]`):
+ grid_size (`tuple[int]`):
The grid size of the spatial positional embedding (height, width).
temporal_size (`int`):
The size of the temporal dimension.
@@ -934,10 +933,10 @@ def get_3d_rotary_pos_embed_allegro(
crops_coords,
grid_size,
temporal_size,
- interpolation_scale: Tuple[float, float, float] = (1.0, 1.0, 1.0),
+ interpolation_scale: tuple[float, float, float] = (1.0, 1.0, 1.0),
theta: int = 10000,
- device: Optional[torch.device] = None,
-) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ device: torch.device | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
# TODO(aryan): docs
start, stop = crops_coords
grid_size_h, grid_size_w = grid_size
@@ -973,7 +972,7 @@ def get_3d_rotary_pos_embed_allegro(
def get_2d_rotary_pos_embed(
- embed_dim, crops_coords, grid_size, use_real=True, device: Optional[torch.device] = None, output_type: str = "np"
+ embed_dim, crops_coords, grid_size, use_real=True, device: torch.device | None = None, output_type: str = "np"
):
"""
RoPE for image tokens with 2d structure.
@@ -981,9 +980,9 @@ def get_2d_rotary_pos_embed(
Args:
embed_dim: (`int`):
The embedding dimension size
- crops_coords (`Tuple[int]`)
+ crops_coords (`tuple[int]`)
The top-left and bottom-right coordinates of the crop.
- grid_size (`Tuple[int]`):
+ grid_size (`tuple[int]`):
The grid size of the positional embedding.
use_real (`bool`):
If True, return real part and imaginary part separately. Otherwise, return complex numbers.
@@ -1029,9 +1028,9 @@ def _get_2d_rotary_pos_embed_np(embed_dim, crops_coords, grid_size, use_real=Tru
Args:
embed_dim: (`int`):
The embedding dimension size
- crops_coords (`Tuple[int]`)
+ crops_coords (`tuple[int]`)
The top-left and bottom-right coordinates of the crop.
- grid_size (`Tuple[int]`):
+ grid_size (`tuple[int]`):
The grid size of the positional embedding.
use_real (`bool`):
If True, return real part and imaginary part separately. Otherwise, return complex numbers.
@@ -1119,7 +1118,7 @@ def get_2d_rotary_pos_embed_lumina(embed_dim, len_h, len_w, linear_factor=1.0, n
def get_1d_rotary_pos_embed(
dim: int,
- pos: Union[np.ndarray, int],
+ pos: np.ndarray | int,
theta: float = 10000.0,
use_real=False,
linear_factor=1.0,
@@ -1186,11 +1185,11 @@ def get_1d_rotary_pos_embed(
def apply_rotary_emb(
x: torch.Tensor,
- freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+ freqs_cis: torch.Tensor | tuple[torch.Tensor],
use_real: bool = True,
use_real_unbind_dim: int = -1,
sequence_dim: int = 2,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
"""
Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
@@ -1200,10 +1199,10 @@ def apply_rotary_emb(
Args:
x (`torch.Tensor`):
Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
- freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+ freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
Returns:
- Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+ tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings.
"""
if use_real:
cos, sin = freqs_cis # [S, D]
@@ -1266,7 +1265,7 @@ def __init__(
time_embed_dim: int,
act_fn: str = "silu",
out_dim: int = None,
- post_act_fn: Optional[str] = None,
+ post_act_fn: str | None = None,
cond_proj_dim=None,
sample_proj_bias=True,
):
@@ -1816,7 +1815,7 @@ def forward(
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
encoder_attention_mask: torch.Tensor,
- hidden_dtype: Optional[torch.dtype] = None,
+ hidden_dtype: torch.dtype | None = None,
):
time_proj = self.time_proj(timestep)
time_emb = self.timestep_embedder(time_proj.to(dtype=hidden_dtype))
@@ -1961,7 +1960,7 @@ def __init__(
self,
num_attention_heads: int,
embed_dim: int,
- output_dim: Optional[int] = None,
+ output_dim: int | None = None,
) -> None:
super().__init__()
@@ -2543,7 +2542,7 @@ def __init__(
self.time_proj = Timesteps(timestep_in_dim, timestep_flip_sin_to_cos, timestep_freq_shift)
self.time_embedding = TimestepEmbedding(timestep_in_dim, hidden_dim, act_fn="silu")
- def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Forward pass.
Args:
@@ -2552,7 +2551,7 @@ def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> Tuple[torch.Tensor
timestep (`torch.Tensor`):
Timestep in denoising process.
Returns:
- `Tuple`[`torch.Tensor`, `torch.Tensor`]: The pair (latents, timestep_emb).
+ `tuple`[`torch.Tensor`, `torch.Tensor`]: The pair (latents, timestep_emb).
"""
timestep_emb = self.time_proj(timestep).to(dtype=x.dtype)
timestep_emb = self.time_embedding(timestep_emb)
@@ -2572,7 +2571,7 @@ def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> Tuple[torch.Tensor
class MultiIPAdapterImageProjection(nn.Module):
- def __init__(self, IPAdapterImageProjectionLayers: Union[List[nn.Module], Tuple[nn.Module]]):
+ def __init__(self, IPAdapterImageProjectionLayers: list[nn.Module] | tuple[nn.Module]):
super().__init__()
self.image_projection_layers = nn.ModuleList(IPAdapterImageProjectionLayers)
@@ -2581,7 +2580,7 @@ def num_ip_adapters(self) -> int:
"""Number of IP-Adapters loaded."""
return len(self.image_projection_layers)
- def forward(self, image_embeds: List[torch.Tensor]):
+ def forward(self, image_embeds: list[torch.Tensor]):
projected_image_embeds = []
# currently, we accept `image_embeds` as
diff --git a/src/diffusers/models/lora.py b/src/diffusers/models/lora.py
index 85d61d6d7cdf..489a0f0abea9 100644
--- a/src/diffusers/models/lora.py
+++ b/src/diffusers/models/lora.py
@@ -21,8 +21,6 @@
# ----------------------------------------------------------------#
###################################################################
-from typing import Optional, Tuple, Union
-
import torch
import torch.nn.functional as F
from torch import nn
@@ -198,9 +196,9 @@ def __init__(
in_features: int,
out_features: int,
rank: int = 4,
- network_alpha: Optional[float] = None,
- device: Optional[Union[torch.device, str]] = None,
- dtype: Optional[torch.dtype] = None,
+ network_alpha: float | None = None,
+ device: torch.device | str | None = None,
+ dtype: torch.dtype | None = None,
):
super().__init__()
@@ -260,10 +258,10 @@ def __init__(
in_features: int,
out_features: int,
rank: int = 4,
- kernel_size: Union[int, Tuple[int, int]] = (1, 1),
- stride: Union[int, Tuple[int, int]] = (1, 1),
- padding: Union[int, Tuple[int, int], str] = 0,
- network_alpha: Optional[float] = None,
+ kernel_size: int | tuple[int, int] = (1, 1),
+ stride: int | tuple[int, int] = (1, 1),
+ padding: int | tuple[int, int] | str = 0,
+ network_alpha: float | None = None,
):
super().__init__()
@@ -301,14 +299,14 @@ class LoRACompatibleConv(nn.Conv2d):
A convolutional layer that can be used with LoRA.
"""
- def __init__(self, *args, lora_layer: Optional[LoRAConv2dLayer] = None, **kwargs):
+ def __init__(self, *args, lora_layer: LoRAConv2dLayer | None = None, **kwargs):
deprecation_message = "Use of `LoRACompatibleConv` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`."
deprecate("LoRACompatibleConv", "1.0.0", deprecation_message)
super().__init__(*args, **kwargs)
self.lora_layer = lora_layer
- def set_lora_layer(self, lora_layer: Optional[LoRAConv2dLayer]):
+ def set_lora_layer(self, lora_layer: LoRAConv2dLayer | None):
deprecation_message = "Use of `set_lora_layer()` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`."
deprecate("set_lora_layer", "1.0.0", deprecation_message)
@@ -388,14 +386,14 @@ class LoRACompatibleLinear(nn.Linear):
A Linear layer that can be used with LoRA.
"""
- def __init__(self, *args, lora_layer: Optional[LoRALinearLayer] = None, **kwargs):
+ def __init__(self, *args, lora_layer: LoRALinearLayer | None = None, **kwargs):
deprecation_message = "Use of `LoRACompatibleLinear` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`."
deprecate("LoRACompatibleLinear", "1.0.0", deprecation_message)
super().__init__(*args, **kwargs)
self.lora_layer = lora_layer
- def set_lora_layer(self, lora_layer: Optional[LoRALinearLayer]):
+ def set_lora_layer(self, lora_layer: LoRALinearLayer | None):
deprecation_message = "Use of `set_lora_layer()` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`."
deprecate("set_lora_layer", "1.0.0", deprecation_message)
self.lora_layer = lora_layer
diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py
index 6d2e8df9c286..04642ad5d401 100644
--- a/src/diffusers/models/model_loading_utils.py
+++ b/src/diffusers/models/model_loading_utils.py
@@ -22,7 +22,6 @@
from collections import OrderedDict, defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
-from typing import Dict, List, Optional, Union
from zipfile import is_zipfile
import safetensors
@@ -136,7 +135,7 @@ def _fetch_remapped_cls_from_config(config, old_class):
return old_class
-def _determine_param_device(param_name: str, device_map: Optional[Dict[str, Union[int, str, torch.device]]]):
+def _determine_param_device(param_name: str, device_map: dict[str, int | str | torch.device] | None):
"""
Find the device of param_name from the device_map.
"""
@@ -154,10 +153,10 @@ def _determine_param_device(param_name: str, device_map: Optional[Dict[str, Unio
def load_state_dict(
- checkpoint_file: Union[str, os.PathLike],
- dduf_entries: Optional[Dict[str, DDUFEntry]] = None,
+ checkpoint_file: str | os.PathLike,
+ dduf_entries: dict[str, DDUFEntry] | None = None,
disable_mmap: bool = False,
- map_location: Union[str, torch.device] = "cpu",
+ map_location: str | torch.device = "cpu",
):
"""
Reads a checkpoint file, returning properly formatted errors if they arise.
@@ -214,17 +213,17 @@ def load_state_dict(
def load_model_dict_into_meta(
model,
state_dict: OrderedDict,
- dtype: Optional[Union[str, torch.dtype]] = None,
- model_name_or_path: Optional[str] = None,
- hf_quantizer: Optional[DiffusersQuantizer] = None,
- keep_in_fp32_modules: Optional[List] = None,
- device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
- unexpected_keys: Optional[List[str]] = None,
- offload_folder: Optional[Union[str, os.PathLike]] = None,
- offload_index: Optional[Dict] = None,
- state_dict_index: Optional[Dict] = None,
- state_dict_folder: Optional[Union[str, os.PathLike]] = None,
-) -> List[str]:
+ dtype: str | torch.dtype | None = None,
+ model_name_or_path: str | None = None,
+ hf_quantizer: DiffusersQuantizer | None = None,
+ keep_in_fp32_modules: list | None = None,
+ device_map: dict[str, int | str | torch.device] | None = None,
+ unexpected_keys: list[str] | None = None,
+ offload_folder: str | os.PathLike | None = None,
+ offload_index: dict | None = None,
+ state_dict_index: dict | None = None,
+ state_dict_folder: str | os.PathLike | None = None,
+) -> list[str]:
"""
This is somewhat similar to `_load_state_dict_into_model`, but deals with a model that has some or all of its
params on a `meta` device. It replaces the model params with the data from the `state_dict`
@@ -474,7 +473,7 @@ def _find_mismatched_keys(
def _load_state_dict_into_model(
model_to_load, state_dict: OrderedDict, assign_to_params_buffers: bool = False
-) -> List[str]:
+) -> list[str]:
# Convert old format to new format if needed from a PyTorch state_dict
# copy state_dict so _load_from_state_dict can modify it
state_dict = state_dict.copy()
@@ -513,7 +512,7 @@ def _fetch_index_file(
revision,
user_agent,
commit_hash,
- dduf_entries: Optional[Dict[str, DDUFEntry]] = None,
+ dduf_entries: dict[str, DDUFEntry] | None = None,
):
if is_local:
index_file = Path(
@@ -563,7 +562,7 @@ def _fetch_index_file_legacy(
revision,
user_agent,
commit_hash,
- dduf_entries: Optional[Dict[str, DDUFEntry]] = None,
+ dduf_entries: dict[str, DDUFEntry] | None = None,
):
if is_local:
index_file = Path(
@@ -722,7 +721,7 @@ def _expand_device_map(device_map, param_names):
# Adapted from: https://github.com/huggingface/transformers/blob/0687d481e2c71544501ef9cb3eef795a6e79b1de/src/transformers/modeling_utils.py#L5859
def _caching_allocator_warmup(
- model, expanded_device_map: Dict[str, torch.device], dtype: torch.dtype, hf_quantizer: Optional[DiffusersQuantizer]
+ model, expanded_device_map: dict[str, torch.device], dtype: torch.dtype, hf_quantizer: DiffusersQuantizer | None
) -> None:
"""
This function warm-ups the caching allocator based on the size of the model tensors that will reside on each
diff --git a/src/diffusers/models/modeling_flax_utils.py b/src/diffusers/models/modeling_flax_utils.py
index 3f060993190f..9f62bd7199e0 100644
--- a/src/diffusers/models/modeling_flax_utils.py
+++ b/src/diffusers/models/modeling_flax_utils.py
@@ -15,7 +15,7 @@
import os
from pickle import UnpicklingError
-from typing import Any, Dict, Union
+from typing import Any
import jax
import jax.numpy as jnp
@@ -68,7 +68,7 @@ def _from_config(cls, config, **kwargs):
"""
return cls(config, **kwargs)
- def _cast_floating_to(self, params: Union[Dict, FrozenDict], dtype: jnp.dtype, mask: Any = None) -> Any:
+ def _cast_floating_to(self, params: dict | FrozenDict, dtype: jnp.dtype, mask: Any = None) -> Any:
"""
Helper method to cast floating-point values of given parameter `PyTree` to given `dtype`.
"""
@@ -92,7 +92,7 @@ def conditional_cast(param):
return unflatten_dict(flat_params)
- def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None):
+ def to_bf16(self, params: dict | FrozenDict, mask: Any = None):
r"""
Cast the floating-point `params` to `jax.numpy.bfloat16`. This returns a new `params` tree and does not cast
the `params` in place.
@@ -101,9 +101,9 @@ def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None):
half-precision training or to save weights in bfloat16 for inference in order to save memory and improve speed.
Arguments:
- params (`Union[Dict, FrozenDict]`):
+ params (`dict | FrozenDict`):
A `PyTree` of model parameters.
- mask (`Union[Dict, FrozenDict]`):
+ mask (`dict | FrozenDict`):
A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True`
for params you want to cast, and `False` for those you want to skip.
@@ -131,15 +131,15 @@ def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None):
```"""
return self._cast_floating_to(params, jnp.bfloat16, mask)
- def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None):
+ def to_fp32(self, params: dict | FrozenDict, mask: Any = None):
r"""
Cast the floating-point `params` to `jax.numpy.float32`. This method can be used to explicitly convert the
model parameters to fp32 precision. This returns a new `params` tree and does not cast the `params` in place.
Arguments:
- params (`Union[Dict, FrozenDict]`):
+ params (`dict | FrozenDict`):
A `PyTree` of model parameters.
- mask (`Union[Dict, FrozenDict]`):
+ mask (`dict | FrozenDict`):
A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True`
for params you want to cast, and `False` for those you want to skip.
@@ -158,7 +158,7 @@ def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None):
```"""
return self._cast_floating_to(params, jnp.float32, mask)
- def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None):
+ def to_fp16(self, params: dict | FrozenDict, mask: Any = None):
r"""
Cast the floating-point `params` to `jax.numpy.float16`. This returns a new `params` tree and does not cast the
`params` in place.
@@ -167,9 +167,9 @@ def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None):
half-precision training or to save weights in float16 for inference in order to save memory and improve speed.
Arguments:
- params (`Union[Dict, FrozenDict]`):
+ params (`dict | FrozenDict`):
A `PyTree` of model parameters.
- mask (`Union[Dict, FrozenDict]`):
+ mask (`dict | FrozenDict`):
A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True`
for params you want to cast, and `False` for those you want to skip.
@@ -197,14 +197,14 @@ def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None):
```"""
return self._cast_floating_to(params, jnp.float16, mask)
- def init_weights(self, rng: jax.Array) -> Dict:
+ def init_weights(self, rng: jax.Array) -> dict:
raise NotImplementedError(f"init_weights method has to be implemented for {self}")
@classmethod
@validate_hf_hub_args
def from_pretrained(
cls,
- pretrained_model_name_or_path: Union[str, os.PathLike],
+ pretrained_model_name_or_path: str | os.PathLike,
dtype: jnp.dtype = jnp.float32,
*model_args,
**kwargs,
@@ -233,14 +233,14 @@ def from_pretrained(
model_args (sequence of positional arguments, *optional*):
All remaining positional arguments are passed to the underlying model's `__init__` method.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
local_files_only(`bool`, *optional*, defaults to `False`):
@@ -493,8 +493,8 @@ def from_pretrained(
def save_pretrained(
self,
- save_directory: Union[str, os.PathLike],
- params: Union[Dict, FrozenDict],
+ save_directory: str | os.PathLike,
+ params: dict | FrozenDict,
is_main_process: bool = True,
push_to_hub: bool = False,
**kwargs,
@@ -506,7 +506,7 @@ def save_pretrained(
Arguments:
save_directory (`str` or `os.PathLike`):
Directory to save a model and its configuration file to. Will be created if it doesn't exist.
- params (`Union[Dict, FrozenDict]`):
+ params (`dict | FrozenDict`):
A `PyTree` of model parameters.
is_main_process (`bool`, *optional*, defaults to `True`):
Whether the process calling this is the main process or not. Useful during distributed training and you
@@ -516,7 +516,7 @@ def save_pretrained(
Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
namespace).
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
"""
if os.path.isfile(save_directory):
diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py
index 63e50af61771..0901840679e3 100644
--- a/src/diffusers/models/modeling_utils.py
+++ b/src/diffusers/models/modeling_utils.py
@@ -27,7 +27,7 @@
from contextlib import ExitStack, contextmanager
from functools import wraps
from pathlib import Path
-from typing import Any, Callable, ContextManager, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, ContextManager, Type
import safetensors
import torch
@@ -81,7 +81,7 @@ class ContextManagers:
in the `fastcore` library.
"""
- def __init__(self, context_managers: List[ContextManager]):
+ def __init__(self, context_managers: list[ContextManager]):
self.context_managers = context_managers
self.stack = ExitStack()
@@ -143,7 +143,7 @@ def get_parameter_device(parameter: torch.nn.Module) -> torch.device:
except StopIteration:
# For torch.nn.DataParallel compatibility in PyTorch 1.5
- def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
+ def find_tensor_attributes(module: torch.nn.Module) -> list[tuple[str, Tensor]]:
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
return tuples
@@ -191,7 +191,7 @@ def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype:
return last_dtype
# For nn.DataParallel compatibility in PyTorch > 1.5
- def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]:
+ def find_tensor_attributes(module: nn.Module) -> list[tuple[str, Tensor]]:
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
return tuples
@@ -280,7 +280,7 @@ def is_gradient_checkpointing(self) -> bool:
"""
return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
- def enable_gradient_checkpointing(self, gradient_checkpointing_func: Optional[Callable] = None) -> None:
+ def enable_gradient_checkpointing(self, gradient_checkpointing_func: Callable | None = None) -> None:
"""
Activates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or
*checkpoint activations* in other frameworks).
@@ -349,7 +349,7 @@ def disable_npu_flash_attention(self) -> None:
self.set_use_npu_flash_attention(False)
def set_use_xla_flash_attention(
- self, use_xla_flash_attention: bool, partition_spec: Optional[Callable] = None, **kwargs
+ self, use_xla_flash_attention: bool, partition_spec: Callable | None = None, **kwargs
) -> None:
# Recursively walk through all the children.
# Any children which exposes the set_use_xla_flash_attention method
@@ -365,7 +365,7 @@ def fn_recursive_set_flash_attention(module: torch.nn.Module):
if isinstance(module, torch.nn.Module):
fn_recursive_set_flash_attention(module)
- def enable_xla_flash_attention(self, partition_spec: Optional[Callable] = None, **kwargs):
+ def enable_xla_flash_attention(self, partition_spec: Callable | None = None, **kwargs):
r"""
Enable the flash attention pallals kernel for torch_xla.
"""
@@ -377,9 +377,7 @@ def disable_xla_flash_attention(self):
"""
self.set_use_xla_flash_attention(False)
- def set_use_memory_efficient_attention_xformers(
- self, valid: bool, attention_op: Optional[Callable] = None
- ) -> None:
+ def set_use_memory_efficient_attention_xformers(self, valid: bool, attention_op: Callable | None = None) -> None:
# Recursively walk through all the children.
# Any children which exposes the set_use_memory_efficient_attention_xformers method
# gets the message
@@ -394,7 +392,7 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
if isinstance(module, torch.nn.Module):
fn_recursive_set_mem_eff(module)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None) -> None:
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None) -> None:
r"""
Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/).
@@ -435,9 +433,9 @@ def disable_xformers_memory_efficient_attention(self) -> None:
def enable_layerwise_casting(
self,
storage_dtype: torch.dtype = torch.float8_e4m3fn,
- compute_dtype: Optional[torch.dtype] = None,
- skip_modules_pattern: Optional[Tuple[str, ...]] = None,
- skip_modules_classes: Optional[Tuple[Type[torch.nn.Module], ...]] = None,
+ compute_dtype: torch.dtype | None = None,
+ skip_modules_pattern: tuple[str, ...] | None = None,
+ skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None,
non_blocking: bool = False,
) -> None:
r"""
@@ -473,11 +471,11 @@ def enable_layerwise_casting(
The dtype to which the model should be cast for storage.
compute_dtype (`torch.dtype`):
The dtype to which the model weights should be cast during the forward pass.
- skip_modules_pattern (`Tuple[str, ...]`, *optional*):
+ skip_modules_pattern (`tuple[str, ...]`, *optional*):
A list of patterns to match the names of the modules to skip during the layerwise casting process. If
set to `None`, default skip patterns are used to ignore certain internal layers of modules and PEFT
layers.
- skip_modules_classes (`Tuple[Type[torch.nn.Module], ...]`, *optional*):
+ skip_modules_classes (`tuple[Type[torch.nn.Module], ...]`, *optional*):
A list of module classes to skip during the layerwise casting process.
non_blocking (`bool`, *optional*, defaults to `False`):
If `True`, the weight casting operations are non-blocking.
@@ -522,14 +520,14 @@ def enable_group_offload(
onload_device: torch.device,
offload_device: torch.device = torch.device("cpu"),
offload_type: str = "block_level",
- num_blocks_per_group: Optional[int] = None,
+ num_blocks_per_group: int | None = None,
non_blocking: bool = False,
use_stream: bool = False,
record_stream: bool = False,
low_cpu_mem_usage=False,
- offload_to_disk_path: Optional[str] = None,
- block_modules: Optional[str] = None,
- exclude_kwargs: Optional[str] = None,
+ offload_to_disk_path: str | None = None,
+ block_modules: str | None = None,
+ exclude_kwargs: str | None = None,
) -> None:
r"""
Activates group offloading for the current model.
@@ -668,12 +666,12 @@ def reset_attention_backend(self) -> None:
def save_pretrained(
self,
- save_directory: Union[str, os.PathLike],
+ save_directory: str | os.PathLike,
is_main_process: bool = True,
- save_function: Optional[Callable] = None,
+ save_function: Callable | None = None,
safe_serialization: bool = True,
- variant: Optional[str] = None,
- max_shard_size: Union[int, str] = "10GB",
+ variant: str | None = None,
+ max_shard_size: int | str = "10GB",
push_to_hub: bool = False,
**kwargs,
):
@@ -707,7 +705,7 @@ def save_pretrained(
Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
namespace).
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
"""
if os.path.isfile(save_directory):
@@ -835,7 +833,7 @@ def dequantize(self):
@classmethod
@validate_hf_hub_args
- def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs) -> Self:
+ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs) -> Self:
r"""
Instantiate a pretrained PyTorch model from a pretrained model configuration.
@@ -851,7 +849,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
- A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
with [`~ModelMixin.save_pretrained`].
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
torch_dtype (`torch.dtype`, *optional*):
@@ -859,7 +857,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info (`bool`, *optional*, defaults to `False`):
@@ -881,7 +879,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
Mirror source to resolve accessibility issues if you're downloading a model in China. We do not
guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
information.
- device_map (`Union[int, str, torch.device]` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+ device_map (`int | str | torch.device` or `dict[str, int | str | torch.device]`, *optional*):
A map that specifies where each submodule should go. It doesn't need to be defined for each
parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the
same device. Defaults to `None`, meaning that the model will be loaded on CPU.
@@ -983,9 +981,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
variant = kwargs.pop("variant", None)
use_safetensors = kwargs.pop("use_safetensors", None)
quantization_config = kwargs.pop("quantization_config", None)
- dduf_entries: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_entries", None)
+ dduf_entries: dict[str, DDUFEntry] | None = kwargs.pop("dduf_entries", None)
disable_mmap = kwargs.pop("disable_mmap", False)
- parallel_config: Optional[Union[ParallelConfig, ContextParallelConfig]] = kwargs.pop("parallel_config", None)
+ parallel_config: ParallelConfig | ContextParallelConfig | None = kwargs.pop("parallel_config", None)
is_parallel_loading_enabled = HF_ENABLE_PARALLEL_LOADING
if is_parallel_loading_enabled and not low_cpu_mem_usage:
@@ -1510,8 +1508,8 @@ def compile_repeated_blocks(self, *args, **kwargs):
def enable_parallelism(
self,
*,
- config: Union[ParallelConfig, ContextParallelConfig],
- cp_plan: Optional[Dict[str, ContextParallelModelPlan]] = None,
+ config: ParallelConfig | ContextParallelConfig,
+ cp_plan: dict[str, ContextParallelModelPlan] | None = None,
):
logger.warning(
"`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning."
@@ -1599,20 +1597,20 @@ def _load_pretrained_model(
cls,
model,
state_dict: OrderedDict,
- resolved_model_file: List[str],
- pretrained_model_name_or_path: Union[str, os.PathLike],
- loaded_keys: List[str],
+ resolved_model_file: list[str],
+ pretrained_model_name_or_path: str | os.PathLike,
+ loaded_keys: list[str],
ignore_mismatched_sizes: bool = False,
assign_to_params_buffers: bool = False,
- hf_quantizer: Optional[DiffusersQuantizer] = None,
+ hf_quantizer: DiffusersQuantizer | None = None,
low_cpu_mem_usage: bool = True,
- dtype: Optional[Union[str, torch.dtype]] = None,
- keep_in_fp32_modules: Optional[List[str]] = None,
- device_map: Union[str, int, torch.device, Dict[str, Union[int, str, torch.device]]] = None,
- offload_state_dict: Optional[bool] = None,
- offload_folder: Optional[Union[str, os.PathLike]] = None,
- dduf_entries: Optional[Dict[str, DDUFEntry]] = None,
- is_parallel_loading_enabled: Optional[bool] = False,
+ dtype: str | torch.dtype | None = None,
+ keep_in_fp32_modules: list[str] | None = None,
+ device_map: str | int | torch.device | dict[str, str | int | torch.device] = None,
+ offload_state_dict: bool | None = None,
+ offload_folder: str | os.PathLike | None = None,
+ dduf_entries: dict[str, DDUFEntry] | None = None,
+ is_parallel_loading_enabled: bool | None = False,
disable_mmap: bool = False,
):
model_state_dict = model.state_dict()
@@ -1776,7 +1774,7 @@ def _get_no_split_modules(self, device_map: str):
The device map value. Options are ["auto", "balanced", "balanced_low_0", "sequential"]
Returns:
- `List[str]`: List of modules that should not be split
+ `list[str]`: list of modules that should not be split
"""
_no_split_modules = set()
modules_to_check = [self]
@@ -1997,7 +1995,7 @@ class LegacyModelMixin(ModelMixin):
@classmethod
@validate_hf_hub_args
- def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs):
# To prevent dependency import problem.
from .model_loading_utils import _fetch_remapped_cls_from_config
diff --git a/src/diffusers/models/normalization.py b/src/diffusers/models/normalization.py
index ae2a6298f5f7..84ffb67bfd6a 100644
--- a/src/diffusers/models/normalization.py
+++ b/src/diffusers/models/normalization.py
@@ -14,7 +14,6 @@
# limitations under the License.
import numbers
-from typing import Dict, Optional, Tuple
import torch
import torch.nn as nn
@@ -41,8 +40,8 @@ class AdaLayerNorm(nn.Module):
def __init__(
self,
embedding_dim: int,
- num_embeddings: Optional[int] = None,
- output_dim: Optional[int] = None,
+ num_embeddings: int | None = None,
+ output_dim: int | None = None,
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-5,
chunk_dim: int = 0,
@@ -62,7 +61,7 @@ def __init__(
self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
def forward(
- self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None
+ self, x: torch.Tensor, timestep: torch.Tensor | None = None, temb: torch.Tensor | None = None
) -> torch.Tensor:
if self.emb is not None:
temb = self.emb(timestep)
@@ -116,8 +115,8 @@ def __init__(self, embedding_dim: int, norm_type: str = "layer_norm", bias: bool
def forward(
self,
hidden_states: torch.Tensor,
- emb: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, ...]:
+ emb: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, ...]:
emb = self.linear(self.silu(emb))
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2 = emb.chunk(
9, dim=1
@@ -137,7 +136,7 @@ class AdaLayerNormZero(nn.Module):
num_embeddings (`int`): The size of the embeddings dictionary.
"""
- def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, norm_type="layer_norm", bias=True):
+ def __init__(self, embedding_dim: int, num_embeddings: int | None = None, norm_type="layer_norm", bias=True):
super().__init__()
if num_embeddings is not None:
self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim)
@@ -158,11 +157,11 @@ def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, nor
def forward(
self,
x: torch.Tensor,
- timestep: Optional[torch.Tensor] = None,
- class_labels: Optional[torch.LongTensor] = None,
- hidden_dtype: Optional[torch.dtype] = None,
- emb: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ timestep: torch.Tensor | None = None,
+ class_labels: torch.LongTensor | None = None,
+ hidden_dtype: torch.dtype | None = None,
+ emb: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
if self.emb is not None:
emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
emb = self.linear(self.silu(emb))
@@ -195,8 +194,8 @@ def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
def forward(
self,
x: torch.Tensor,
- emb: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ emb: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
emb = self.linear(self.silu(emb))
shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
@@ -224,8 +223,8 @@ def __init__(self, embedding_dim: int, norm_eps: float, norm_elementwise_affine:
def forward(
self,
x: torch.Tensor,
- emb: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ emb: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
emb = self.linear(self.silu(emb))
scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
x = self.norm(x) * (1 + scale_msa[:, None])
@@ -257,10 +256,10 @@ def __init__(self, embedding_dim: int, use_additional_conditions: bool = False):
def forward(
self,
timestep: torch.Tensor,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- batch_size: Optional[int] = None,
- hidden_dtype: Optional[torch.dtype] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ batch_size: int | None = None,
+ hidden_dtype: torch.dtype | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
# No modulation happening here.
added_cond_kwargs = added_cond_kwargs or {"resolution": None, "aspect_ratio": None}
embedded_timestep = self.emb(timestep, **added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_dtype)
@@ -280,7 +279,7 @@ class AdaGroupNorm(nn.Module):
"""
def __init__(
- self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: Optional[str] = None, eps: float = 1e-5
+ self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: str | None = None, eps: float = 1e-5
):
super().__init__()
self.num_groups = num_groups
@@ -366,7 +365,7 @@ def __init__(
eps=1e-5,
bias=True,
norm_type="layer_norm",
- out_dim: Optional[int] = None,
+ out_dim: int | None = None,
):
super().__init__()
@@ -422,8 +421,8 @@ def forward(
self,
x: torch.Tensor,
context: torch.Tensor,
- emb: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ emb: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
emb = self.linear(self.silu(emb))
(
shift_msa,
@@ -463,7 +462,7 @@ def __init__(
def forward(
self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
shift, scale, gate, enc_shift, enc_scale, enc_gate = self.linear(self.silu(temb)).chunk(6, dim=1)
hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
encoder_hidden_states = self.norm(encoder_hidden_states) * (1 + enc_scale)[:, None, :] + enc_shift[:, None, :]
@@ -632,7 +631,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
def get_normalization(
norm_type: str = "batch_norm",
- num_features: Optional[int] = None,
+ num_features: int | None = None,
eps: float = 1e-5,
elementwise_affine: bool = True,
bias: bool = True,
diff --git a/src/diffusers/models/resnet.py b/src/diffusers/models/resnet.py
index a98cb491146b..df793b534ebb 100644
--- a/src/diffusers/models/resnet.py
+++ b/src/diffusers/models/resnet.py
@@ -14,7 +14,6 @@
# limitations under the License.
from functools import partial
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -75,21 +74,21 @@ def __init__(
self,
*,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
conv_shortcut: bool = False,
dropout: float = 0.0,
temb_channels: int = 512,
groups: int = 32,
- groups_out: Optional[int] = None,
+ groups_out: int | None = None,
eps: float = 1e-6,
non_linearity: str = "swish",
time_embedding_norm: str = "ada_group", # ada_group, spatial
output_scale_factor: float = 1.0,
- use_in_shortcut: Optional[bool] = None,
+ use_in_shortcut: bool | None = None,
up: bool = False,
down: bool = False,
conv_shortcut_bias: bool = True,
- conv_2d_out_channels: Optional[int] = None,
+ conv_2d_out_channels: int | None = None,
):
super().__init__()
self.in_channels = in_channels
@@ -221,24 +220,24 @@ def __init__(
self,
*,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
conv_shortcut: bool = False,
dropout: float = 0.0,
temb_channels: int = 512,
groups: int = 32,
- groups_out: Optional[int] = None,
+ groups_out: int | None = None,
pre_norm: bool = True,
eps: float = 1e-6,
non_linearity: str = "swish",
skip_time_act: bool = False,
time_embedding_norm: str = "default", # default, scale_shift,
- kernel: Optional[torch.Tensor] = None,
+ kernel: torch.Tensor | None = None,
output_scale_factor: float = 1.0,
- use_in_shortcut: Optional[bool] = None,
+ use_in_shortcut: bool | None = None,
up: bool = False,
down: bool = False,
conv_shortcut_bias: bool = True,
- conv_2d_out_channels: Optional[int] = None,
+ conv_2d_out_channels: int | None = None,
):
super().__init__()
if time_embedding_norm == "ada_group":
@@ -406,7 +405,7 @@ def __init__(
self,
inp_channels: int,
out_channels: int,
- kernel_size: Union[int, Tuple[int, int]],
+ kernel_size: int | tuple[int, int],
n_groups: int = 8,
activation: str = "mish",
):
@@ -443,7 +442,7 @@ def __init__(
inp_channels: int,
out_channels: int,
embed_dim: int,
- kernel_size: Union[int, Tuple[int, int]] = 5,
+ kernel_size: int | tuple[int, int] = 5,
activation: str = "mish",
):
super().__init__()
@@ -487,7 +486,7 @@ class TemporalConvLayer(nn.Module):
def __init__(
self,
in_dim: int,
- out_dim: Optional[int] = None,
+ out_dim: int | None = None,
dropout: float = 0.0,
norm_num_groups: int = 32,
):
@@ -559,7 +558,7 @@ class TemporalResnetBlock(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
temb_channels: int = 512,
eps: float = 1e-6,
):
@@ -658,10 +657,10 @@ class SpatioTemporalResBlock(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
temb_channels: int = 512,
eps: float = 1e-6,
- temporal_eps: Optional[float] = None,
+ temporal_eps: float | None = None,
merge_factor: float = 0.5,
merge_strategy="learned_with_images",
switch_spatial_to_temporal_mix: bool = False,
@@ -691,8 +690,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- image_only_indicator: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ image_only_indicator: torch.Tensor | None = None,
):
num_frames = image_only_indicator.shape[-1]
hidden_states = self.spatial_res_block(hidden_states, temb)
@@ -790,7 +789,7 @@ def forward(
self,
x_spatial: torch.Tensor,
x_temporal: torch.Tensor,
- image_only_indicator: Optional[torch.Tensor] = None,
+ image_only_indicator: torch.Tensor | None = None,
) -> torch.Tensor:
alpha = self.get_alpha(image_only_indicator, x_spatial.ndim)
alpha = alpha.to(x_spatial.dtype)
diff --git a/src/diffusers/models/transformers/auraflow_transformer_2d.py b/src/diffusers/models/transformers/auraflow_transformer_2d.py
index e3732662e408..3fa4df738784 100644
--- a/src/diffusers/models/transformers/auraflow_transformer_2d.py
+++ b/src/diffusers/models/transformers/auraflow_transformer_2d.py
@@ -13,7 +13,7 @@
# limitations under the License.
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import AttentionMixin
from ..attention_processor import (
@@ -172,7 +172,7 @@ def forward(
self,
hidden_states: torch.FloatTensor,
temb: torch.FloatTensor,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
) -> torch.Tensor:
residual = hidden_states
attention_kwargs = attention_kwargs or {}
@@ -241,8 +241,8 @@ def forward(
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor,
temb: torch.FloatTensor,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
residual = hidden_states
residual_context = encoder_hidden_states
attention_kwargs = attention_kwargs or {}
@@ -397,29 +397,15 @@ def unfuse_qkv_projections(self):
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
timestep: torch.LongTensor = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
height, width = hidden_states.shape[-2:]
# Apply patch embedding, timestep embedding, and project the caption embeddings.
@@ -486,10 +472,6 @@ def forward(
shape=(hidden_states.shape[0], out_channels, height * patch_size, width * patch_size)
)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/cogvideox_transformer_3d.py b/src/diffusers/models/transformers/cogvideox_transformer_3d.py
index 14b38cd46c52..4b8beeeb6fe3 100644
--- a/src/diffusers/models/transformers/cogvideox_transformer_3d.py
+++ b/src/diffusers/models/transformers/cogvideox_transformer_3d.py
@@ -13,14 +13,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
from torch import nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import Attention, AttentionMixin, FeedForward
from ..attention_processor import CogVideoXAttnProcessor2_0, FusedCogVideoXAttnProcessor2_0
@@ -83,7 +83,7 @@ def __init__(
norm_elementwise_affine: bool = True,
norm_eps: float = 1e-5,
final_dropout: bool = True,
- ff_inner_dim: Optional[int] = None,
+ ff_inner_dim: int | None = None,
ff_bias: bool = True,
attention_out_bias: bool = True,
):
@@ -120,9 +120,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_length = encoder_hidden_states.size(1)
attention_kwargs = attention_kwargs or {}
@@ -223,11 +223,11 @@ def __init__(
num_attention_heads: int = 30,
attention_head_dim: int = 64,
in_channels: int = 16,
- out_channels: Optional[int] = 16,
+ out_channels: int | None = 16,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
time_embed_dim: int = 512,
- ofs_embed_dim: Optional[int] = None,
+ ofs_embed_dim: int | None = None,
text_embed_dim: int = 4096,
num_layers: int = 30,
dropout: float = 0.0,
@@ -236,7 +236,7 @@ def __init__(
sample_height: int = 60,
sample_frames: int = 49,
patch_size: int = 2,
- patch_size_t: Optional[int] = None,
+ patch_size_t: int | None = None,
temporal_compression_ratio: int = 4,
max_text_seq_length: int = 226,
activation_fn: str = "gelu-approximate",
@@ -363,32 +363,18 @@ def unfuse_qkv_projections(self):
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- timestep: Union[int, float, torch.LongTensor],
- timestep_cond: Optional[torch.Tensor] = None,
- ofs: Optional[Union[int, float, torch.LongTensor]] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ timestep: int | float | torch.LongTensor,
+ timestep_cond: torch.Tensor | None = None,
+ ofs: int | float | torch.LongTensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_frames, channels, height, width = hidden_states.shape
# 1. Time embedding
@@ -454,10 +440,6 @@ def forward(
)
output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
diff --git a/src/diffusers/models/transformers/consisid_transformer_3d.py b/src/diffusers/models/transformers/consisid_transformer_3d.py
index be20b0a3eacf..64a58e394366 100644
--- a/src/diffusers/models/transformers/consisid_transformer_3d.py
+++ b/src/diffusers/models/transformers/consisid_transformer_3d.py
@@ -13,14 +13,14 @@
# limitations under the License.
import math
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
from torch import nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import Attention, AttentionMixin, FeedForward
from ..attention_processor import CogVideoXAttnProcessor2_0
@@ -34,7 +34,7 @@
class PerceiverAttention(nn.Module):
- def __init__(self, dim: int, dim_head: int = 64, heads: int = 8, kv_dim: Optional[int] = None):
+ def __init__(self, dim: int, dim_head: int = 64, heads: int = 8, kv_dim: int | None = None):
super().__init__()
self.scale = dim_head**-0.5
@@ -152,7 +152,7 @@ def __init__(
nn.Linear(vit_dim, vit_dim * num_id_token),
)
- def forward(self, id_embeds: torch.Tensor, vit_hidden_states: List[torch.Tensor]) -> torch.Tensor:
+ def forward(self, id_embeds: torch.Tensor, vit_hidden_states: list[torch.Tensor]) -> torch.Tensor:
# Repeat latent queries for the batch size
latents = self.latents.repeat(id_embeds.size(0), 1, 1)
@@ -277,7 +277,7 @@ def __init__(
norm_elementwise_affine: bool = True,
norm_eps: float = 1e-5,
final_dropout: bool = True,
- ff_inner_dim: Optional[int] = None,
+ ff_inner_dim: int | None = None,
ff_bias: bool = True,
attention_out_bias: bool = True,
):
@@ -314,8 +314,8 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_length = encoder_hidden_states.size(1)
# norm & modulate
@@ -465,7 +465,7 @@ def __init__(
num_attention_heads: int = 30,
attention_head_dim: int = 64,
in_channels: int = 16,
- out_channels: Optional[int] = 16,
+ out_channels: int | None = 16,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
time_embed_dim: int = 512,
@@ -620,33 +620,19 @@ def _init_face_inputs(self):
]
)
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- timestep: Union[int, float, torch.LongTensor],
- timestep_cond: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- id_cond: Optional[torch.Tensor] = None,
- id_vit_hidden: Optional[torch.Tensor] = None,
+ timestep: int | float | torch.LongTensor,
+ timestep_cond: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ id_cond: torch.Tensor | None = None,
+ id_vit_hidden: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
# fuse clip and insightface
valid_face_emb = None
if self.is_train_face:
@@ -720,10 +706,6 @@ def forward(
output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)
output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
diff --git a/src/diffusers/models/transformers/dit_transformer_2d.py b/src/diffusers/models/transformers/dit_transformer_2d.py
index 68f6f769436e..3d10c278cdbb 100644
--- a/src/diffusers/models/transformers/dit_transformer_2d.py
+++ b/src/diffusers/models/transformers/dit_transformer_2d.py
@@ -11,7 +11,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.
-from typing import Any, Dict, Optional
+from typing import Any
import torch
import torch.nn.functional as F
@@ -74,7 +74,7 @@ def __init__(
num_attention_heads: int = 16,
attention_head_dim: int = 72,
in_channels: int = 4,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 28,
dropout: float = 0.0,
norm_num_groups: int = 32,
@@ -82,7 +82,7 @@ def __init__(
sample_size: int = 32,
patch_size: int = 2,
activation_fn: str = "gelu-approximate",
- num_embeds_ada_norm: Optional[int] = 1000,
+ num_embeds_ada_norm: int | None = 1000,
upcast_attention: bool = False,
norm_type: str = "ada_norm_zero",
norm_elementwise_affine: bool = False,
@@ -148,9 +148,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- timestep: Optional[torch.LongTensor] = None,
- class_labels: Optional[torch.LongTensor] = None,
- cross_attention_kwargs: Dict[str, Any] = None,
+ timestep: torch.LongTensor | None = None,
+ class_labels: torch.LongTensor | None = None,
+ cross_attention_kwargs: dict[str, Any] = None,
return_dict: bool = True,
):
"""
@@ -164,7 +164,7 @@ def forward(
class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
`AdaLayerZeroNorm`.
- cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
+ cross_attention_kwargs ( `dict[str, Any]`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
diff --git a/src/diffusers/models/transformers/dual_transformer_2d.py b/src/diffusers/models/transformers/dual_transformer_2d.py
index 24eed2168229..c25c6e9c4227 100644
--- a/src/diffusers/models/transformers/dual_transformer_2d.py
+++ b/src/diffusers/models/transformers/dual_transformer_2d.py
@@ -11,8 +11,6 @@
# 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.
-from typing import Optional
-
from torch import nn
from ..modeling_outputs import Transformer2DModelOutput
@@ -50,16 +48,16 @@ def __init__(
self,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
+ in_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
norm_num_groups: int = 32,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
- sample_size: Optional[int] = None,
- num_vector_embeds: Optional[int] = None,
+ sample_size: int | None = None,
+ num_vector_embeds: int | None = None,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
):
super().__init__()
self.transformers = nn.ModuleList(
diff --git a/src/diffusers/models/transformers/hunyuan_transformer_2d.py b/src/diffusers/models/transformers/hunyuan_transformer_2d.py
index cecb675b32b7..a25aa99fb8b9 100644
--- a/src/diffusers/models/transformers/hunyuan_transformer_2d.py
+++ b/src/diffusers/models/transformers/hunyuan_transformer_2d.py
@@ -11,8 +11,6 @@
# 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.
-from typing import Optional
-
import torch
from torch import nn
@@ -98,7 +96,7 @@ def __init__(
norm_elementwise_affine: bool = True,
norm_eps: float = 1e-6,
final_dropout: bool = False,
- ff_inner_dim: Optional[int] = None,
+ ff_inner_dim: int | None = None,
ff_bias: bool = True,
skip: bool = False,
qk_norm: bool = True,
@@ -158,7 +156,7 @@ def __init__(
self._chunk_dim = 0
# Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward
- def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+ def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0):
# Sets chunk feed-forward
self._chunk_size = chunk_size
self._chunk_dim = dim
@@ -166,8 +164,8 @@ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
image_rotary_emb=None,
skip=None,
) -> torch.Tensor:
@@ -252,8 +250,8 @@ def __init__(
self,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
- patch_size: Optional[int] = None,
+ in_channels: int | None = None,
+ patch_size: int | None = None,
activation_fn: str = "gelu-approximate",
sample_size=32,
hidden_size=1152,
@@ -469,7 +467,7 @@ def forward(
return Transformer2DModelOutput(sample=output)
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
- def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+ def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
diff --git a/src/diffusers/models/transformers/latte_transformer_3d.py b/src/diffusers/models/transformers/latte_transformer_3d.py
index 990c90512e39..32e97aff8fb7 100644
--- a/src/diffusers/models/transformers/latte_transformer_3d.py
+++ b/src/diffusers/models/transformers/latte_transformer_3d.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional
-
import torch
from torch import nn
@@ -73,16 +71,16 @@ def __init__(
self,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
- out_channels: Optional[int] = None,
+ in_channels: int | None = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
sample_size: int = 64,
- patch_size: Optional[int] = None,
+ patch_size: int | None = None,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
norm_type: str = "layer_norm",
norm_elementwise_affine: bool = True,
norm_eps: float = 1e-5,
@@ -168,9 +166,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- timestep: Optional[torch.LongTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ timestep: torch.LongTensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
enable_temporal_attentions: bool = True,
return_dict: bool = True,
):
diff --git a/src/diffusers/models/transformers/lumina_nextdit2d.py b/src/diffusers/models/transformers/lumina_nextdit2d.py
index bed5e69c2d36..46a6753b4cb1 100644
--- a/src/diffusers/models/transformers/lumina_nextdit2d.py
+++ b/src/diffusers/models/transformers/lumina_nextdit2d.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -123,7 +123,7 @@ def forward(
encoder_hidden_states: torch.Tensor,
encoder_mask: torch.Tensor,
temb: torch.Tensor,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
) -> torch.Tensor:
"""
Perform a forward pass through the LuminaNextDiTBlock.
@@ -135,7 +135,7 @@ def forward(
encoder_hidden_states: (`torch.Tensor`): The hidden_states of text prompt are processed by Gemma encoder.
encoder_mask (`torch.Tensor`): The hidden_states of text prompt attention mask.
temb (`torch.Tensor`): Timestep embedding with text prompt embedding.
- cross_attention_kwargs (`Dict[str, Any]`): kwargs for cross attention.
+ cross_attention_kwargs (`dict[str, Any]`): kwargs for cross attention.
"""
residual = hidden_states
@@ -227,19 +227,19 @@ class LuminaNextDiT2DModel(ModelMixin, ConfigMixin):
def __init__(
self,
sample_size: int = 128,
- patch_size: Optional[int] = 2,
- in_channels: Optional[int] = 4,
- hidden_size: Optional[int] = 2304,
- num_layers: Optional[int] = 32,
- num_attention_heads: Optional[int] = 32,
- num_kv_heads: Optional[int] = None,
- multiple_of: Optional[int] = 256,
- ffn_dim_multiplier: Optional[float] = None,
- norm_eps: Optional[float] = 1e-5,
- learn_sigma: Optional[bool] = True,
- qk_norm: Optional[bool] = True,
- cross_attention_dim: Optional[int] = 2048,
- scaling_factor: Optional[float] = 1.0,
+ patch_size: int | None = 2,
+ in_channels: int | None = 4,
+ hidden_size: int | None = 2304,
+ num_layers: int | None = 32,
+ num_attention_heads: int | None = 32,
+ num_kv_heads: int | None = None,
+ multiple_of: int | None = 256,
+ ffn_dim_multiplier: float | None = None,
+ norm_eps: float | None = 1e-5,
+ learn_sigma: bool | None = True,
+ qk_norm: bool | None = True,
+ cross_attention_dim: int | None = 2048,
+ scaling_factor: float | None = 1.0,
) -> None:
super().__init__()
self.sample_size = sample_size
@@ -295,9 +295,9 @@ def forward(
encoder_hidden_states: torch.Tensor,
encoder_mask: torch.Tensor,
image_rotary_emb: torch.Tensor,
- cross_attention_kwargs: Dict[str, Any] = None,
+ cross_attention_kwargs: dict[str, Any] = None,
return_dict=True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
"""
Forward pass of LuminaNextDiT.
diff --git a/src/diffusers/models/transformers/pixart_transformer_2d.py b/src/diffusers/models/transformers/pixart_transformer_2d.py
index 072670ee0c30..2476668ba307 100644
--- a/src/diffusers/models/transformers/pixart_transformer_2d.py
+++ b/src/diffusers/models/transformers/pixart_transformer_2d.py
@@ -11,7 +11,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.
-from typing import Any, Dict, Optional
+from typing import Any
import torch
from torch import nn
@@ -87,24 +87,24 @@ def __init__(
num_attention_heads: int = 16,
attention_head_dim: int = 72,
in_channels: int = 4,
- out_channels: Optional[int] = 8,
+ out_channels: int | None = 8,
num_layers: int = 28,
dropout: float = 0.0,
norm_num_groups: int = 32,
- cross_attention_dim: Optional[int] = 1152,
+ cross_attention_dim: int | None = 1152,
attention_bias: bool = True,
sample_size: int = 128,
patch_size: int = 2,
activation_fn: str = "gelu-approximate",
- num_embeds_ada_norm: Optional[int] = 1000,
+ num_embeds_ada_norm: int | None = 1000,
upcast_attention: bool = False,
norm_type: str = "ada_norm_single",
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-6,
- interpolation_scale: Optional[int] = None,
- use_additional_conditions: Optional[bool] = None,
- caption_channels: Optional[int] = None,
- attention_type: Optional[str] = "default",
+ interpolation_scale: int | None = None,
+ use_additional_conditions: bool | None = None,
+ caption_channels: int | None = None,
+ attention_type: str | None = "default",
):
super().__init__()
@@ -227,12 +227,12 @@ def unfuse_qkv_projections(self):
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- timestep: Optional[torch.LongTensor] = None,
- added_cond_kwargs: Dict[str, torch.Tensor] = None,
- cross_attention_kwargs: Dict[str, Any] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ timestep: torch.LongTensor | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] = None,
+ cross_attention_kwargs: dict[str, Any] = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
return_dict: bool = True,
):
"""
@@ -246,8 +246,8 @@ def forward(
self-attention.
timestep (`torch.LongTensor`, *optional*):
Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
- added_cond_kwargs: (`Dict[str, Any]`, *optional*): Additional conditions to be used as inputs.
- cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
+ added_cond_kwargs: (`dict[str, Any]`, *optional*): Additional conditions to be used as inputs.
+ cross_attention_kwargs ( `dict[str, Any]`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
diff --git a/src/diffusers/models/transformers/prior_transformer.py b/src/diffusers/models/transformers/prior_transformer.py
index 757bb436040f..ace2b529c4f2 100644
--- a/src/diffusers/models/transformers/prior_transformer.py
+++ b/src/diffusers/models/transformers/prior_transformer.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import Optional, Union
import torch
import torch.nn.functional as F
@@ -80,13 +79,13 @@ def __init__(
additional_embeddings=4,
dropout: float = 0.0,
time_embed_act_fn: str = "silu",
- norm_in_type: Optional[str] = None, # layer
- embedding_proj_norm_type: Optional[str] = None, # layer
- encoder_hid_proj_type: Optional[str] = "linear", # linear
- added_emb_type: Optional[str] = "prd", # prd
- time_embed_dim: Optional[int] = None,
- embedding_proj_dim: Optional[int] = None,
- clip_embed_dim: Optional[int] = None,
+ norm_in_type: str | None = None, # layer
+ embedding_proj_norm_type: str | None = None, # layer
+ encoder_hid_proj_type: str | None = "linear", # linear
+ added_emb_type: str | None = "prd", # prd
+ time_embed_dim: int | None = None,
+ embedding_proj_dim: int | None = None,
+ clip_embed_dim: int | None = None,
):
super().__init__()
self.num_attention_heads = num_attention_heads
@@ -184,10 +183,10 @@ def set_default_attn_processor(self):
def forward(
self,
hidden_states,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
proj_embedding: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.BoolTensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.BoolTensor | None = None,
return_dict: bool = True,
):
"""
diff --git a/src/diffusers/models/transformers/sana_transformer.py b/src/diffusers/models/transformers/sana_transformer.py
index d45badb1b121..ff078dc695d7 100644
--- a/src/diffusers/models/transformers/sana_transformer.py
+++ b/src/diffusers/models/transformers/sana_transformer.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn.functional as F
@@ -20,7 +20,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ..attention import AttentionMixin
from ..attention_processor import (
Attention,
@@ -41,7 +41,7 @@ def __init__(
in_channels: int,
out_channels: int,
expand_ratio: float = 4,
- norm_type: Optional[str] = None,
+ norm_type: str | None = None,
residual_connection: bool = True,
) -> None:
super().__init__()
@@ -132,8 +132,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
@@ -196,15 +196,15 @@ def __init__(
num_attention_heads: int = 70,
attention_head_dim: int = 32,
dropout: float = 0.0,
- num_cross_attention_heads: Optional[int] = 20,
- cross_attention_head_dim: Optional[int] = 112,
- cross_attention_dim: Optional[int] = 2240,
+ num_cross_attention_heads: int | None = 20,
+ cross_attention_head_dim: int | None = 112,
+ cross_attention_dim: int | None = 2240,
attention_bias: bool = True,
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-6,
attention_out_bias: bool = True,
mlp_ratio: float = 2.5,
- qk_norm: Optional[str] = None,
+ qk_norm: str | None = None,
) -> None:
super().__init__()
@@ -246,10 +246,10 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- timestep: Optional[torch.LongTensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ timestep: torch.LongTensor | None = None,
height: int = None,
width: int = None,
) -> torch.Tensor:
@@ -340,13 +340,13 @@ class SanaTransformer2DModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapte
def __init__(
self,
in_channels: int = 32,
- out_channels: Optional[int] = 32,
+ out_channels: int | None = 32,
num_attention_heads: int = 70,
attention_head_dim: int = 32,
num_layers: int = 20,
- num_cross_attention_heads: Optional[int] = 20,
- cross_attention_head_dim: Optional[int] = 112,
- cross_attention_dim: Optional[int] = 2240,
+ num_cross_attention_heads: int | None = 20,
+ cross_attention_head_dim: int | None = 112,
+ cross_attention_dim: int | None = 2240,
caption_channels: int = 2304,
mlp_ratio: float = 2.5,
dropout: float = 0.0,
@@ -355,10 +355,10 @@ def __init__(
patch_size: int = 1,
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-6,
- interpolation_scale: Optional[int] = None,
+ interpolation_scale: int | None = None,
guidance_embeds: bool = False,
guidance_embeds_scale: float = 0.1,
- qk_norm: Optional[str] = None,
+ qk_norm: str | None = None,
timestep_scale: float = 1.0,
) -> None:
super().__init__()
@@ -414,33 +414,19 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep: torch.Tensor,
- guidance: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_block_samples: Optional[Tuple[torch.Tensor]] = None,
+ guidance: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ controlnet_block_samples: tuple[torch.Tensor] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput:
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
@@ -527,10 +513,6 @@ def forward(
hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4)
output = hidden_states.reshape(batch_size, -1, post_patch_height * p, post_patch_width * p)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/stable_audio_transformer.py b/src/diffusers/models/transformers/stable_audio_transformer.py
index 2c3b6b5df91d..f4974926ec72 100644
--- a/src/diffusers/models/transformers/stable_audio_transformer.py
+++ b/src/diffusers/models/transformers/stable_audio_transformer.py
@@ -13,8 +13,6 @@
# limitations under the License.
-from typing import Optional, Union
-
import numpy as np
import torch
import torch.nn as nn
@@ -87,10 +85,10 @@ def __init__(
num_key_value_attention_heads: int,
attention_head_dim: int,
dropout=0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
upcast_attention: bool = False,
norm_eps: float = 1e-5,
- ff_inner_dim: Optional[int] = None,
+ ff_inner_dim: int | None = None,
):
super().__init__()
# Define 3 blocks. Each block has its own normalization layer.
@@ -138,7 +136,7 @@ def __init__(
self._chunk_size = None
self._chunk_dim = 0
- def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+ def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0):
# Sets chunk feed-forward
self._chunk_size = chunk_size
self._chunk_dim = dim
@@ -146,10 +144,10 @@ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
def forward(
self,
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- rotary_embedding: Optional[torch.FloatTensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ rotary_embedding: torch.FloatTensor | None = None,
) -> torch.Tensor:
# Notice that normalization is always applied before the real computation in the following blocks.
# 0. Self-Attention
@@ -289,9 +287,9 @@ def forward(
global_hidden_states: torch.FloatTensor = None,
rotary_embedding: torch.FloatTensor = None,
return_dict: bool = True,
- attention_mask: Optional[torch.LongTensor] = None,
- encoder_attention_mask: Optional[torch.LongTensor] = None,
- ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+ attention_mask: torch.LongTensor | None = None,
+ encoder_attention_mask: torch.LongTensor | None = None,
+ ) -> torch.FloatTensor | Transformer2DModelOutput:
"""
The [`StableAudioDiTModel`] forward method.
diff --git a/src/diffusers/models/transformers/t5_film_transformer.py b/src/diffusers/models/transformers/t5_film_transformer.py
index 7a9608735e32..1ae2b1e3fedb 100644
--- a/src/diffusers/models/transformers/t5_film_transformer.py
+++ b/src/diffusers/models/transformers/t5_film_transformer.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import math
-from typing import Optional, Tuple
import torch
from torch import nn
@@ -196,12 +195,12 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- conditioning_emb: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ conditioning_emb: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
encoder_decoder_position_bias=None,
- ) -> Tuple[torch.Tensor]:
+ ) -> tuple[torch.Tensor]:
hidden_states = self.layer[0](
hidden_states,
conditioning_emb=conditioning_emb,
@@ -250,8 +249,8 @@ def __init__(self, d_model: int, d_kv: int, num_heads: int, dropout_rate: float)
def forward(
self,
hidden_states: torch.Tensor,
- conditioning_emb: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ conditioning_emb: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
# pre_self_attention_layer_norm
normed_hidden_states = self.layer_norm(hidden_states)
@@ -293,8 +292,8 @@ def __init__(self, d_model: int, d_kv: int, num_heads: int, dropout_rate: float,
def forward(
self,
hidden_states: torch.Tensor,
- key_value_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ key_value_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
normed_hidden_states = self.layer_norm(hidden_states)
attention_output = self.attention(
@@ -328,7 +327,7 @@ def __init__(self, d_model: int, d_ff: int, dropout_rate: float, layer_norm_epsi
self.layer_norm = T5LayerNorm(d_model, eps=layer_norm_epsilon)
self.dropout = nn.Dropout(dropout_rate)
- def forward(self, hidden_states: torch.Tensor, conditioning_emb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, conditioning_emb: torch.Tensor | None = None) -> torch.Tensor:
forwarded_states = self.layer_norm(hidden_states)
if conditioning_emb is not None:
forwarded_states = self.film(forwarded_states, conditioning_emb)
diff --git a/src/diffusers/models/transformers/transformer_2d.py b/src/diffusers/models/transformers/transformer_2d.py
index 67fe9a33109b..12f89201d752 100644
--- a/src/diffusers/models/transformers/transformer_2d.py
+++ b/src/diffusers/models/transformers/transformer_2d.py
@@ -11,7 +11,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.
-from typing import Any, Dict, Optional
+from typing import Any
import torch
import torch.nn.functional as F
@@ -73,18 +73,18 @@ def __init__(
self,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
- out_channels: Optional[int] = None,
+ in_channels: int | None = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
norm_num_groups: int = 32,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
- sample_size: Optional[int] = None,
- num_vector_embeds: Optional[int] = None,
- patch_size: Optional[int] = None,
+ sample_size: int | None = None,
+ num_vector_embeds: int | None = None,
+ patch_size: int | None = None,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
use_linear_projection: bool = False,
only_cross_attention: bool = False,
double_self_attention: bool = False,
@@ -95,7 +95,7 @@ def __init__(
attention_type: str = "default",
caption_channels: int = None,
interpolation_scale: float = None,
- use_additional_conditions: Optional[bool] = None,
+ use_additional_conditions: bool | None = None,
):
super().__init__()
@@ -324,13 +324,13 @@ def _init_patched_inputs(self, norm_type):
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- timestep: Optional[torch.LongTensor] = None,
- added_cond_kwargs: Dict[str, torch.Tensor] = None,
- class_labels: Optional[torch.LongTensor] = None,
- cross_attention_kwargs: Dict[str, Any] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ timestep: torch.LongTensor | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] = None,
+ class_labels: torch.LongTensor | None = None,
+ cross_attention_kwargs: dict[str, Any] = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
return_dict: bool = True,
):
"""
@@ -347,7 +347,7 @@ def forward(
class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
`AdaLayerZeroNorm`.
- cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
+ cross_attention_kwargs ( `dict[str, Any]`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
diff --git a/src/diffusers/models/transformers/transformer_allegro.py b/src/diffusers/models/transformers/transformer_allegro.py
index 5fa59a71d977..934e2787674f 100644
--- a/src/diffusers/models/transformers/transformer_allegro.py
+++ b/src/diffusers/models/transformers/transformer_allegro.py
@@ -13,8 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple
-
import torch
import torch.nn as nn
import torch.nn.functional as F
@@ -69,7 +67,7 @@ def __init__(
num_attention_heads: int,
attention_head_dim: int,
dropout=0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
activation_fn: str = "geglu",
attention_bias: bool = False,
norm_elementwise_affine: bool = True,
@@ -117,10 +115,10 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- temb: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ temb: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
image_rotary_emb=None,
) -> torch.Tensor:
# 0. Self-Attention
@@ -309,9 +307,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep: torch.LongTensor,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
return_dict: bool = True,
):
batch_size, num_channels, num_frames, height, width = hidden_states.shape
diff --git a/src/diffusers/models/transformers/transformer_bria.py b/src/diffusers/models/transformers/transformer_bria.py
index d54679306e64..99b7bbfd64cf 100644
--- a/src/diffusers/models/transformers/transformer_bria.py
+++ b/src/diffusers/models/transformers/transformer_bria.py
@@ -1,5 +1,5 @@
import inspect
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -8,7 +8,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import AttentionModuleMixin, FeedForward
from ..attention_dispatch import dispatch_attention_fn
@@ -54,7 +54,7 @@ def _get_qkv_projections(attn: "BriaAttention", hidden_states, encoder_hidden_st
def get_1d_rotary_pos_embed(
dim: int,
- pos: Union[np.ndarray, int],
+ pos: np.ndarray | int,
theta: float = 10000.0,
use_real=False,
linear_factor=1.0,
@@ -131,8 +131,8 @@ def __call__(
attn: "BriaAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
attn, hidden_states, encoder_hidden_states
@@ -198,12 +198,12 @@ def __init__(
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
- added_kv_proj_dim: Optional[int] = None,
- added_proj_bias: Optional[bool] = True,
+ added_kv_proj_dim: int | None = None,
+ added_proj_bias: bool | None = True,
out_bias: bool = True,
eps: float = 1e-5,
out_dim: int = None,
- context_pre_only: Optional[bool] = None,
+ context_pre_only: bool | None = None,
pre_only: bool = False,
elementwise_affine: bool = True,
processor=None,
@@ -248,9 +248,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -266,7 +266,7 @@ def forward(
class BriaEmbedND(torch.nn.Module):
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
- def __init__(self, theta: int, axes_dim: List[int]):
+ def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -334,7 +334,7 @@ def forward(self, timestep, dtype):
class BriaPosEmbed(torch.nn.Module):
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
- def __init__(self, theta: int, axes_dim: List[int]):
+ def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -395,9 +395,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -476,9 +476,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_len = encoder_hidden_states.shape[1]
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
@@ -537,7 +537,7 @@ def __init__(
joint_attention_dim: int = 4096,
pooled_projection_dim: int = None,
guidance_embeds: bool = False,
- axes_dims_rope: List[int] = [16, 56, 56],
+ axes_dims_rope: list[int] = [16, 56, 56],
rope_theta=10000,
time_theta=10000,
):
@@ -581,6 +581,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -590,11 +591,11 @@ def forward(
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
guidance: torch.Tensor = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
controlnet_block_samples=None,
controlnet_single_block_samples=None,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
"""
The [`BriaTransformer2DModel`] forward method.
@@ -621,20 +622,6 @@ def forward(
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
hidden_states = self.x_embedder(hidden_states)
timestep = timestep.to(hidden_states.dtype)
@@ -715,10 +702,6 @@ def forward(
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_bria_fibo.py b/src/diffusers/models/transformers/transformer_bria_fibo.py
index e10bfddcbc86..7ddbccfa47c5 100644
--- a/src/diffusers/models/transformers/transformer_bria_fibo.py
+++ b/src/diffusers/models/transformers/transformer_bria_fibo.py
@@ -8,7 +8,7 @@
#
# See the license for further details.
import inspect
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -22,10 +22,8 @@
from ...models.modeling_utils import ModelMixin
from ...models.transformers.transformer_bria import BriaAttnProcessor
from ...utils import (
- USE_PEFT_BACKEND,
+ apply_lora_scale,
logging,
- scale_lora_layers,
- unscale_lora_layers,
)
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import AttentionModuleMixin, FeedForward
@@ -80,8 +78,8 @@ def __call__(
attn: "BriaFiboAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
attn, hidden_states, encoder_hidden_states
@@ -146,12 +144,12 @@ def __init__(
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
- added_kv_proj_dim: Optional[int] = None,
- added_proj_bias: Optional[bool] = True,
+ added_kv_proj_dim: int | None = None,
+ added_proj_bias: bool | None = True,
out_bias: bool = True,
eps: float = 1e-5,
out_dim: int = None,
- context_pre_only: Optional[bool] = None,
+ context_pre_only: bool | None = None,
pre_only: bool = False,
elementwise_affine: bool = True,
processor=None,
@@ -196,9 +194,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -214,7 +212,7 @@ def forward(
class BriaFiboEmbedND(torch.nn.Module):
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
- def __init__(self, theta: int, axes_dim: List[int]):
+ def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -272,8 +270,8 @@ def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
) -> torch.Tensor:
residual = hidden_states
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
@@ -339,9 +337,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -458,7 +456,7 @@ def __init__(
joint_attention_dim: int = 4096,
pooled_projection_dim: int = None,
guidance_embeds: bool = False,
- axes_dims_rope: List[int] = [16, 56, 56],
+ axes_dims_rope: list[int] = [16, 56, 56],
rope_theta=10000,
time_theta=10000,
text_encoder_dim: int = 2048,
@@ -510,6 +508,7 @@ def __init__(
]
self.caption_projection = nn.ModuleList(caption_projection)
+ @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -520,9 +519,9 @@ def forward(
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
guidance: torch.Tensor = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+ ) -> torch.FloatTensor | Transformer2DModelOutput:
"""
Args:
@@ -545,20 +544,7 @@ def forward(
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
- if joint_attention_kwargs is not None:
- joint_attention_kwargs = joint_attention_kwargs.copy()
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
hidden_states = self.x_embedder(hidden_states)
timestep = timestep.to(hidden_states.dtype)
@@ -645,10 +631,6 @@ def forward(
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_chroma.py b/src/diffusers/models/transformers/transformer_chroma.py
index 2ef3643dafbd..d7cc96d018b3 100644
--- a/src/diffusers/models/transformers/transformer_chroma.py
+++ b/src/diffusers/models/transformers/transformer_chroma.py
@@ -13,7 +13,7 @@
# limitations under the License.
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, deprecate, logging
from ...utils.import_utils import is_torch_npu_available
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import AttentionMixin, FeedForward
@@ -45,7 +45,7 @@ class ChromaAdaLayerNormZeroPruned(nn.Module):
num_embeddings (`int`): The size of the embeddings dictionary.
"""
- def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, norm_type="layer_norm", bias=True):
+ def __init__(self, embedding_dim: int, num_embeddings: int | None = None, norm_type="layer_norm", bias=True):
super().__init__()
if num_embeddings is not None:
self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim)
@@ -64,11 +64,11 @@ def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, nor
def forward(
self,
x: torch.Tensor,
- timestep: Optional[torch.Tensor] = None,
- class_labels: Optional[torch.LongTensor] = None,
- hidden_dtype: Optional[torch.dtype] = None,
- emb: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ timestep: torch.Tensor | None = None,
+ class_labels: torch.LongTensor | None = None,
+ hidden_dtype: torch.dtype | None = None,
+ emb: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
if self.emb is not None:
emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.flatten(1, 2).chunk(6, dim=1)
@@ -98,8 +98,8 @@ def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
def forward(
self,
x: torch.Tensor,
- emb: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ emb: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
shift_msa, scale_msa, gate_msa = emb.flatten(1, 2).chunk(3, dim=1)
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return x, gate_msa
@@ -243,9 +243,9 @@ def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
) -> torch.Tensor:
residual = hidden_states
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
@@ -309,10 +309,10 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
temb_img, temb_txt = temb[:, :6], temb[:, 6:]
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb_img)
@@ -399,7 +399,7 @@ class ChromaTransformer2DModel(
joint_attention_dim (`int`, defaults to `4096`):
The number of dimensions to use for the joint attention (embedding/channel dimension of
`encoder_hidden_states`).
- axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`):
The dimensions to use for the rotary positional embeddings.
"""
@@ -413,13 +413,13 @@ def __init__(
self,
patch_size: int = 1,
in_channels: int = 64,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 19,
num_single_layers: int = 38,
attention_head_dim: int = 128,
num_attention_heads: int = 24,
joint_attention_dim: int = 4096,
- axes_dims_rope: Tuple[int, ...] = (16, 56, 56),
+ axes_dims_rope: tuple[int, ...] = (16, 56, 56),
approximator_num_channels: int = 64,
approximator_hidden_dim: int = 5120,
approximator_layers: int = 5,
@@ -473,6 +473,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -481,12 +482,12 @@ def forward(
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
attention_mask: torch.Tensor = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
controlnet_block_samples=None,
controlnet_single_block_samples=None,
return_dict: bool = True,
controlnet_blocks_repeat: bool = False,
- ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ ) -> torch.Tensor | Transformer2DModelOutput:
"""
The [`FluxTransformer2DModel`] forward method.
@@ -511,20 +512,6 @@ def forward(
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
- if joint_attention_kwargs is not None:
- joint_attention_kwargs = joint_attention_kwargs.copy()
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
hidden_states = self.x_embedder(hidden_states)
@@ -631,10 +618,6 @@ def forward(
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_chronoedit.py b/src/diffusers/models/transformers/transformer_chronoedit.py
index 79828b6464f4..25eb6f87a93a 100644
--- a/src/diffusers/models/transformers/transformer_chronoedit.py
+++ b/src/diffusers/models/transformers/transformer_chronoedit.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, deprecate, logging
from ...utils.torch_utils import maybe_allow_in_graph
from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
@@ -43,7 +43,7 @@ def _get_qkv_projections(attn: "WanAttention", hidden_states: torch.Tensor, enco
encoder_hidden_states = hidden_states
if attn.fused_projections:
- if attn.cross_attention_dim_head is None:
+ if not attn.is_cross_attention:
# In self-attention layers, we can fuse the entire QKV projection into a single linear
query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
else:
@@ -82,9 +82,9 @@ def __call__(
self,
attn: "WanAttention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
encoder_hidden_states_img = None
if attn.add_k_proj is not None:
@@ -188,8 +188,8 @@ def __init__(
dim_head: int = 64,
eps: float = 1e-5,
dropout: float = 0.0,
- added_kv_proj_dim: Optional[int] = None,
- cross_attention_dim_head: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
+ cross_attention_dim_head: int | None = None,
processor=None,
is_cross_attention=None,
):
@@ -219,7 +219,10 @@ def __init__(
self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True)
self.norm_added_k = torch.nn.RMSNorm(dim_head * heads, eps=eps)
- self.is_cross_attention = cross_attention_dim_head is not None
+ if is_cross_attention is not None:
+ self.is_cross_attention = is_cross_attention
+ else:
+ self.is_cross_attention = cross_attention_dim_head is not None
self.set_processor(processor)
@@ -227,7 +230,7 @@ def fuse_projections(self):
if getattr(self, "fused_projections", False):
return
- if self.cross_attention_dim_head is None:
+ if not self.is_cross_attention:
concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])
concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])
out_features, in_features = concatenated_weights.shape
@@ -275,9 +278,9 @@ def unfuse_projections(self):
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs)
@@ -316,8 +319,8 @@ def __init__(
time_freq_dim: int,
time_proj_dim: int,
text_embed_dim: int,
- image_embed_dim: Optional[int] = None,
- pos_embed_seq_len: Optional[int] = None,
+ image_embed_dim: int | None = None,
+ pos_embed_seq_len: int | None = None,
):
super().__init__()
@@ -335,8 +338,8 @@ def forward(
self,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
- timestep_seq_len: Optional[int] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
+ timestep_seq_len: int | None = None,
):
timestep = self.timesteps_proj(timestep)
if timestep_seq_len is not None:
@@ -359,7 +362,7 @@ class ChronoEditRotaryPosEmbed(nn.Module):
def __init__(
self,
attention_head_dim: int,
- patch_size: Tuple[int, int, int],
+ patch_size: tuple[int, int, int],
max_seq_len: int,
theta: float = 10000.0,
temporal_skip_len: int = 8,
@@ -438,7 +441,7 @@ def __init__(
qk_norm: str = "rms_norm_across_heads",
cross_attn_norm: bool = False,
eps: float = 1e-6,
- added_kv_proj_dim: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
):
super().__init__()
@@ -524,7 +527,7 @@ class ChronoEditTransformer3DModel(
A Transformer model for video-like data used in the ChronoEdit model.
Args:
- patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`):
+ patch_size (`tuple[int]`, defaults to `(1, 2, 2)`):
3D patch dimensions for video embedding (t_patch, h_patch, w_patch).
num_attention_heads (`int`, defaults to `40`):
Fixed length for text embeddings.
@@ -542,7 +545,7 @@ class ChronoEditTransformer3DModel(
Intermediate dimension in feed-forward network.
num_layers (`int`, defaults to `40`):
The number of layers of transformer blocks to use.
- window_size (`Tuple[int]`, defaults to `(-1, -1)`):
+ window_size (`tuple[int]`, defaults to `(-1, -1)`):
Window size for local attention (-1 indicates global attention).
cross_attn_norm (`bool`, defaults to `True`):
Enable cross-attention normalization.
@@ -581,7 +584,7 @@ class ChronoEditTransformer3DModel(
@register_to_config
def __init__(
self,
- patch_size: Tuple[int] = (1, 2, 2),
+ patch_size: tuple[int] = (1, 2, 2),
num_attention_heads: int = 40,
attention_head_dim: int = 128,
in_channels: int = 16,
@@ -591,12 +594,12 @@ def __init__(
ffn_dim: int = 13824,
num_layers: int = 40,
cross_attn_norm: bool = True,
- qk_norm: Optional[str] = "rms_norm_across_heads",
+ qk_norm: str | None = "rms_norm_across_heads",
eps: float = 1e-6,
- image_dim: Optional[int] = None,
- added_kv_proj_dim: Optional[int] = None,
+ image_dim: int | None = None,
+ added_kv_proj_dim: int | None = None,
rope_max_seq_len: int = 1024,
- pos_embed_seq_len: Optional[int] = None,
+ pos_embed_seq_len: int | None = None,
rope_temporal_skip_len: int = 8,
) -> None:
super().__init__()
@@ -638,30 +641,16 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> torch.Tensor | dict[str, torch.Tensor]:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.config.patch_size
post_patch_num_frames = num_frames // p_t
@@ -729,10 +718,6 @@ def forward(
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_cogview3plus.py b/src/diffusers/models/transformers/transformer_cogview3plus.py
index e48290fb39d4..ad6a442acbcc 100644
--- a/src/diffusers/models/transformers/transformer_cogview3plus.py
+++ b/src/diffusers/models/transformers/transformer_cogview3plus.py
@@ -13,8 +13,6 @@
# limitations under the License.
-from typing import Tuple, Union
-
import torch
import torch.nn as nn
@@ -79,7 +77,7 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
emb: torch.Tensor,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_length = encoder_hidden_states.size(1)
# norm & modulate
@@ -233,7 +231,7 @@ def forward(
target_size: torch.Tensor,
crop_coords: torch.Tensor,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
"""
The [`CogView3PlusTransformer2DModel`] forward method.
diff --git a/src/diffusers/models/transformers/transformer_cogview4.py b/src/diffusers/models/transformers/transformer_cogview4.py
index 64e9a538a7c2..308e0e6cccaf 100644
--- a/src/diffusers/models/transformers/transformer_cogview4.py
+++ b/src/diffusers/models/transformers/transformer_cogview4.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -20,7 +20,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import FeedForward
from ..attention_processor import Attention
@@ -73,7 +73,7 @@ def __init__(self, embedding_dim: int, dim: int) -> None:
def forward(
self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
dtype = hidden_states.dtype
norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype)
@@ -129,9 +129,9 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
dtype = encoder_hidden_states.dtype
batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
@@ -211,14 +211,12 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- latent_attn_mask: Optional[torch.Tensor] = None,
- text_attn_mask: Optional[torch.Tensor] = None,
- batch_flag: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[
- Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]]
- ] = None,
+ latent_attn_mask: torch.Tensor | None = None,
+ text_attn_mask: torch.Tensor | None = None,
+ batch_flag: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None,
**kwargs,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
"""
Args:
attn (`Attention`):
@@ -238,10 +236,10 @@ def __call__(
Values from 0 to n-1 indicating which samples belong to the same batch. Samples with the same
batch_flag are packed together. Example: [0, 1, 1, 2, 2] means sample 0 forms batch0, samples 1-2 form
batch1, and samples 3-4 form batch2. If None, no packing is used.
- image_rotary_emb (`Tuple[torch.Tensor, torch.Tensor]` or `list[Tuple[torch.Tensor, torch.Tensor]]`, *optional*):
+ image_rotary_emb (`tuple[torch.Tensor, torch.Tensor]` or `list[tuple[torch.Tensor, torch.Tensor]]`, *optional*):
The rotary embedding for the image part of the input.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`: The processed hidden states for both image and text streams.
+ `tuple[torch.Tensor, torch.Tensor]`: The processed hidden states for both image and text streams.
"""
# Get dimensions and device info
@@ -488,13 +486,11 @@ def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[
- Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]]
- ] = None,
- attention_mask: Optional[Dict[str, torch.Tensor]] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ temb: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None,
+ attention_mask: dict[str, torch.Tensor] | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# 1. Timestep conditioning
(
norm_hidden_states,
@@ -537,7 +533,7 @@ def forward(
class CogView4RotaryPosEmbed(nn.Module):
- def __init__(self, dim: int, patch_size: int, rope_axes_dim: Tuple[int, int], theta: float = 10000.0) -> None:
+ def __init__(self, dim: int, patch_size: int, rope_axes_dim: tuple[int, int], theta: float = 10000.0) -> None:
super().__init__()
self.dim = dim
@@ -545,7 +541,7 @@ def __init__(self, dim: int, patch_size: int, rope_axes_dim: Tuple[int, int], th
self.rope_axes_dim = rope_axes_dim
self.theta = theta
- def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
batch_size, num_channels, height, width = hidden_states.shape
height, width = height // self.patch_size, width // self.patch_size
@@ -666,7 +662,7 @@ def __init__(
condition_dim: int = 256,
pos_embed_max_size: int = 128,
sample_size: int = 128,
- rope_axes_dim: Tuple[int, int] = (256, 256),
+ rope_axes_dim: tuple[int, int] = (256, 256),
):
super().__init__()
@@ -703,6 +699,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -711,28 +708,11 @@ def forward(
original_size: torch.Tensor,
target_size: torch.Tensor,
crop_coords: torch.Tensor,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[
- Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]]
- ] = None,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None,
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_channels, height, width = hidden_states.shape
# 1. RoPE
@@ -779,10 +759,6 @@ def forward(
hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p)
output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
diff --git a/src/diffusers/models/transformers/transformer_cosmos.py b/src/diffusers/models/transformers/transformer_cosmos.py
index 2b0c2667072b..46746a19a678 100644
--- a/src/diffusers/models/transformers/transformer_cosmos.py
+++ b/src/diffusers/models/transformers/transformer_cosmos.py
@@ -12,17 +12,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple
-
import numpy as np
import torch
import torch.nn as nn
-import torch.nn.functional as F
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin
from ...utils import is_torchvision_available
from ..attention import FeedForward
+from ..attention_dispatch import dispatch_attention_fn
from ..attention_processor import Attention
from ..embeddings import Timesteps
from ..modeling_outputs import Transformer2DModelOutput
@@ -36,7 +34,7 @@
class CosmosPatchEmbed(nn.Module):
def __init__(
- self, in_channels: int, out_channels: int, patch_size: Tuple[int, int, int], bias: bool = True
+ self, in_channels: int, out_channels: int, patch_size: tuple[int, int, int], bias: bool = True
) -> None:
super().__init__()
self.patch_size = patch_size
@@ -94,7 +92,7 @@ def __init__(self, in_features: int, hidden_features: int) -> None:
self.linear_2 = nn.Linear(hidden_features, 2 * in_features, bias=False)
def forward(
- self, hidden_states: torch.Tensor, embedded_timestep: torch.Tensor, temb: Optional[torch.Tensor] = None
+ self, hidden_states: torch.Tensor, embedded_timestep: torch.Tensor, temb: torch.Tensor | None = None
) -> torch.Tensor:
embedded_timestep = self.activation(embedded_timestep)
embedded_timestep = self.linear_1(embedded_timestep)
@@ -114,7 +112,7 @@ def forward(
class CosmosAdaLayerNormZero(nn.Module):
- def __init__(self, in_features: int, hidden_features: Optional[int] = None) -> None:
+ def __init__(self, in_features: int, hidden_features: int | None = None) -> None:
super().__init__()
self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6)
@@ -131,7 +129,7 @@ def forward(
self,
hidden_states: torch.Tensor,
embedded_timestep: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
) -> torch.Tensor:
embedded_timestep = self.activation(embedded_timestep)
embedded_timestep = self.linear_1(embedded_timestep)
@@ -152,16 +150,16 @@ def forward(
class CosmosAttnProcessor2_0:
def __init__(self):
- if not hasattr(F, "scaled_dot_product_attention"):
+ if not hasattr(torch.nn.functional, "scaled_dot_product_attention"):
raise ImportError("CosmosAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
# 1. QKV projections
if encoder_hidden_states is None:
@@ -191,7 +189,6 @@ def __call__(
query_idx = torch.tensor(query.size(3), device=query.device)
key_idx = torch.tensor(key.size(3), device=key.device)
value_idx = torch.tensor(value.size(3), device=value.device)
-
else:
query_idx = query.size(3)
key_idx = key.size(3)
@@ -200,18 +197,148 @@ def __call__(
value = value.repeat_interleave(query_idx // value_idx, dim=3)
# 5. Attention
- hidden_states = F.scaled_dot_product_attention(
- query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ hidden_states = dispatch_attention_fn(
+ query.transpose(1, 2),
+ key.transpose(1, 2),
+ value.transpose(1, 2),
+ attn_mask=attention_mask,
+ dropout_p=0.0,
+ is_causal=False,
)
- hidden_states = hidden_states.transpose(1, 2).flatten(2, 3).type_as(query)
-
- # 6. Output projection
+ hidden_states = hidden_states.flatten(2, 3).type_as(query)
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
+class CosmosAttnProcessor2_5:
+ def __init__(self):
+ if not hasattr(torch.nn.functional, "scaled_dot_product_attention"):
+ raise ImportError("CosmosAttnProcessor2_5 requires PyTorch 2.0. Please upgrade PyTorch to 2.0 or newer.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: tuple[torch.Tensor, torch.Tensor],
+ image_rotary_emb=None,
+ ) -> torch.Tensor:
+ if not isinstance(encoder_hidden_states, tuple):
+ raise ValueError("Expected encoder_hidden_states as (text_context, img_context) tuple.")
+
+ text_context, img_context = encoder_hidden_states if encoder_hidden_states else (None, None)
+ text_mask, img_mask = attention_mask if attention_mask else (None, None)
+
+ if text_context is None:
+ text_context = hidden_states
+
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(text_context)
+ value = attn.to_v(text_context)
+
+ query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+ key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+ value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+
+ query = attn.norm_q(query)
+ key = attn.norm_k(key)
+
+ if image_rotary_emb is not None:
+ from ..embeddings import apply_rotary_emb
+
+ query = apply_rotary_emb(query, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
+ key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
+
+ if torch.onnx.is_in_onnx_export():
+ query_idx = torch.tensor(query.size(3), device=query.device)
+ key_idx = torch.tensor(key.size(3), device=key.device)
+ value_idx = torch.tensor(value.size(3), device=value.device)
+ else:
+ query_idx = query.size(3)
+ key_idx = key.size(3)
+ value_idx = value.size(3)
+ key = key.repeat_interleave(query_idx // key_idx, dim=3)
+ value = value.repeat_interleave(query_idx // value_idx, dim=3)
+
+ attn_out = dispatch_attention_fn(
+ query.transpose(1, 2),
+ key.transpose(1, 2),
+ value.transpose(1, 2),
+ attn_mask=text_mask,
+ dropout_p=0.0,
+ is_causal=False,
+ )
+ attn_out = attn_out.flatten(2, 3).type_as(query)
+
+ if img_context is not None:
+ q_img = attn.q_img(hidden_states)
+ k_img = attn.k_img(img_context)
+ v_img = attn.v_img(img_context)
+
+ batch_size = hidden_states.shape[0]
+ dim_head = attn.out_dim // attn.heads
+
+ q_img = q_img.view(batch_size, -1, attn.heads, dim_head).transpose(1, 2)
+ k_img = k_img.view(batch_size, -1, attn.heads, dim_head).transpose(1, 2)
+ v_img = v_img.view(batch_size, -1, attn.heads, dim_head).transpose(1, 2)
+
+ q_img = attn.q_img_norm(q_img)
+ k_img = attn.k_img_norm(k_img)
+
+ q_img_idx = q_img.size(3)
+ k_img_idx = k_img.size(3)
+ v_img_idx = v_img.size(3)
+ k_img = k_img.repeat_interleave(q_img_idx // k_img_idx, dim=3)
+ v_img = v_img.repeat_interleave(q_img_idx // v_img_idx, dim=3)
+
+ img_out = dispatch_attention_fn(
+ q_img.transpose(1, 2),
+ k_img.transpose(1, 2),
+ v_img.transpose(1, 2),
+ attn_mask=img_mask,
+ dropout_p=0.0,
+ is_causal=False,
+ )
+ img_out = img_out.flatten(2, 3).type_as(q_img)
+ hidden_states = attn_out + img_out
+ else:
+ hidden_states = attn_out
+
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+ return hidden_states
+
+
+class CosmosAttention(Attention):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ # add parameters for image q/k/v
+ inner_dim = self.heads * self.to_q.out_features // self.heads
+ self.q_img = nn.Linear(self.query_dim, inner_dim, bias=False)
+ self.k_img = nn.Linear(self.query_dim, inner_dim, bias=False)
+ self.v_img = nn.Linear(self.query_dim, inner_dim, bias=False)
+ self.q_img_norm = RMSNorm(self.to_q.out_features // self.heads, eps=1e-6, elementwise_affine=True)
+ self.k_img_norm = RMSNorm(self.to_k.out_features // self.heads, eps=1e-6, elementwise_affine=True)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None = None,
+ **cross_attention_kwargs,
+ ) -> torch.Tensor:
+ return super().forward(
+ hidden_states=hidden_states,
+ # NOTE: type-hint in base class can be ignored
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ **cross_attention_kwargs,
+ )
+
+
class CosmosTransformerBlock(nn.Module):
def __init__(
self,
@@ -222,12 +349,16 @@ def __init__(
adaln_lora_dim: int = 256,
qk_norm: str = "rms_norm",
out_bias: bool = False,
+ img_context: bool = False,
+ before_proj: bool = False,
+ after_proj: bool = False,
) -> None:
super().__init__()
hidden_size = num_attention_heads * attention_head_dim
self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
+ self.img_context = img_context
self.attn1 = Attention(
query_dim=hidden_size,
cross_attention_dim=None,
@@ -240,30 +371,56 @@ def __init__(
)
self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
- self.attn2 = Attention(
- query_dim=hidden_size,
- cross_attention_dim=cross_attention_dim,
- heads=num_attention_heads,
- dim_head=attention_head_dim,
- qk_norm=qk_norm,
- elementwise_affine=True,
- out_bias=out_bias,
- processor=CosmosAttnProcessor2_0(),
- )
+ if img_context:
+ self.attn2 = CosmosAttention(
+ query_dim=hidden_size,
+ cross_attention_dim=cross_attention_dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ qk_norm=qk_norm,
+ elementwise_affine=True,
+ out_bias=out_bias,
+ processor=CosmosAttnProcessor2_5(),
+ )
+ else:
+ self.attn2 = Attention(
+ query_dim=hidden_size,
+ cross_attention_dim=cross_attention_dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ qk_norm=qk_norm,
+ elementwise_affine=True,
+ out_bias=out_bias,
+ processor=CosmosAttnProcessor2_0(),
+ )
self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu", bias=out_bias)
+ # NOTE: zero conv for CosmosControlNet
+ self.before_proj = None
+ self.after_proj = None
+ if before_proj:
+ self.before_proj = nn.Linear(hidden_size, hidden_size)
+ if after_proj:
+ self.after_proj = nn.Linear(hidden_size, hidden_size)
+
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor | None | tuple[torch.Tensor | None, torch.Tensor | None],
embedded_timestep: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
- extra_pos_emb: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
+ temb: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
+ extra_pos_emb: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ controlnet_residual: torch.Tensor | None = None,
+ latents: torch.Tensor | None = None,
+ block_idx: int | None = None,
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+ if self.before_proj is not None:
+ hidden_states = self.before_proj(hidden_states) + latents
+
if extra_pos_emb is not None:
hidden_states = hidden_states + extra_pos_emb
@@ -284,6 +441,16 @@ def forward(
ff_output = self.ff(norm_hidden_states)
hidden_states = hidden_states + gate * ff_output
+ if controlnet_residual is not None:
+ assert self.after_proj is None
+ # NOTE: this is assumed to be scaled by the controlnet
+ hidden_states += controlnet_residual
+
+ if self.after_proj is not None:
+ assert controlnet_residual is None
+ hs_proj = self.after_proj(hidden_states)
+ return hidden_states, hs_proj
+
return hidden_states
@@ -291,10 +458,10 @@ class CosmosRotaryPosEmbed(nn.Module):
def __init__(
self,
hidden_size: int,
- max_size: Tuple[int, int, int] = (128, 240, 240),
- patch_size: Tuple[int, int, int] = (1, 2, 2),
+ max_size: tuple[int, int, int] = (128, 240, 240),
+ patch_size: tuple[int, int, int] = (1, 2, 2),
base_fps: int = 24,
- rope_scale: Tuple[float, float, float] = (2.0, 1.0, 1.0),
+ rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0),
) -> None:
super().__init__()
@@ -310,7 +477,7 @@ def __init__(
self.w_ntk_factor = rope_scale[2] ** (self.dim_w / (self.dim_w - 2))
self.t_ntk_factor = rope_scale[0] ** (self.dim_t / (self.dim_t - 2))
- def forward(self, hidden_states: torch.Tensor, fps: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+ def forward(self, hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]
device = hidden_states.device
@@ -355,8 +522,8 @@ class CosmosLearnablePositionalEmbed(nn.Module):
def __init__(
self,
hidden_size: int,
- max_size: Tuple[int, int, int],
- patch_size: Tuple[int, int, int],
+ max_size: tuple[int, int, int],
+ patch_size: tuple[int, int, int],
eps: float = 1e-6,
) -> None:
super().__init__()
@@ -405,17 +572,28 @@ class CosmosTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
Input dimension of text embeddings from the text encoder.
adaln_lora_dim (`int`, defaults to `256`):
The hidden dimension of the Adaptive LayerNorm LoRA layer.
- max_size (`Tuple[int, int, int]`, defaults to `(128, 240, 240)`):
+ max_size (`tuple[int, int, int]`, defaults to `(128, 240, 240)`):
The maximum size of the input latent tensors in the temporal, height, and width dimensions.
- patch_size (`Tuple[int, int, int]`, defaults to `(1, 2, 2)`):
+ patch_size (`tuple[int, int, int]`, defaults to `(1, 2, 2)`):
The patch size to use for patchifying the input latent tensors in the temporal, height, and width
dimensions.
- rope_scale (`Tuple[float, float, float]`, defaults to `(2.0, 1.0, 1.0)`):
+ rope_scale (`tuple[float, float, float]`, defaults to `(2.0, 1.0, 1.0)`):
The scaling factor to use for RoPE in the temporal, height, and width dimensions.
concat_padding_mask (`bool`, defaults to `True`):
Whether to concatenate the padding mask to the input latent tensors.
extra_pos_embed_type (`str`, *optional*, defaults to `learnable`):
The type of extra positional embeddings to use. Can be one of `None` or `learnable`.
+ controlnet_block_every_n (`int`, *optional*):
+ Interval between transformer blocks that should receive control residuals (for example, `7` to inject after
+ every seventh block). Required for Cosmos Transfer2.5.
+ img_context_dim_in (`int`, *optional*):
+ The dimension of the input image context feature vector, i.e. it is the D in [B, N, D].
+ img_context_num_tokens (`int`):
+ The number of tokens in the image context feature vector, i.e. it is the N in [B, N, D]. If
+ `img_context_dim_in` is not provided, then this parameter is ignored.
+ img_context_dim_out (`int`):
+ The output dimension of the image context projection layer. If `img_context_dim_in` is not provided, then
+ this parameter is ignored.
"""
_supports_gradient_checkpointing = True
@@ -434,14 +612,18 @@ def __init__(
mlp_ratio: float = 4.0,
text_embed_dim: int = 1024,
adaln_lora_dim: int = 256,
- max_size: Tuple[int, int, int] = (128, 240, 240),
- patch_size: Tuple[int, int, int] = (1, 2, 2),
- rope_scale: Tuple[float, float, float] = (2.0, 1.0, 1.0),
+ max_size: tuple[int, int, int] = (128, 240, 240),
+ patch_size: tuple[int, int, int] = (1, 2, 2),
+ rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0),
concat_padding_mask: bool = True,
- extra_pos_embed_type: Optional[str] = "learnable",
+ extra_pos_embed_type: str | None = "learnable",
use_crossattn_projection: bool = False,
crossattn_proj_in_channels: int = 1024,
encoder_hidden_states_channels: int = 1024,
+ controlnet_block_every_n: int | None = None,
+ img_context_dim_in: int | None = None,
+ img_context_num_tokens: int = 256,
+ img_context_dim_out: int = 2048,
) -> None:
super().__init__()
hidden_size = num_attention_heads * attention_head_dim
@@ -477,6 +659,7 @@ def __init__(
adaln_lora_dim=adaln_lora_dim,
qk_norm="rms_norm",
out_bias=False,
+ img_context=self.config.img_context_dim_in is not None and self.config.img_context_dim_in > 0,
)
for _ in range(num_layers)
]
@@ -496,17 +679,24 @@ def __init__(
self.gradient_checkpointing = False
+ if self.config.img_context_dim_in:
+ self.img_context_proj = nn.Sequential(
+ nn.Linear(self.config.img_context_dim_in, self.config.img_context_dim_out, bias=True),
+ nn.GELU(),
+ )
+
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- fps: Optional[int] = None,
- condition_mask: Optional[torch.Tensor] = None,
- padding_mask: Optional[torch.Tensor] = None,
+ block_controlnet_hidden_states: list[torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ fps: int | None = None,
+ condition_mask: torch.Tensor | None = None,
+ padding_mask: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> torch.Tensor:
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
# 1. Concatenate padding mask if needed & prepare attention mask
@@ -514,11 +704,11 @@ def forward(
hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
if self.config.concat_padding_mask:
- padding_mask = transforms.functional.resize(
+ padding_mask_resized = transforms.functional.resize(
padding_mask, list(hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
)
hidden_states = torch.cat(
- [hidden_states, padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1
+ [hidden_states, padding_mask_resized.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1
)
if attention_mask is not None:
@@ -554,36 +744,59 @@ def forward(
for x in (temb, embedded_timestep)
) # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C]
else:
- assert False
+ raise ValueError(f"Expected timestep to have shape [B, 1, T, 1, 1] or [T], but got {timestep.shape}")
+ # 5. Process encoder hidden states
+ text_context, img_context = (
+ encoder_hidden_states if isinstance(encoder_hidden_states, tuple) else (encoder_hidden_states, None)
+ )
if self.config.use_crossattn_projection:
- encoder_hidden_states = self.crossattn_proj(encoder_hidden_states)
+ text_context = self.crossattn_proj(text_context)
+
+ if img_context is not None and self.config.img_context_dim_in:
+ img_context = self.img_context_proj(img_context)
- # 5. Transformer blocks
- for block in self.transformer_blocks:
+ processed_encoder_hidden_states = (
+ (text_context, img_context) if isinstance(encoder_hidden_states, tuple) else text_context
+ )
+
+ # 6. Build controlnet block index map
+ controlnet_block_index_map = {}
+ if block_controlnet_hidden_states is not None:
+ n_blocks = len(self.transformer_blocks)
+ controlnet_block_index_map = {
+ block_idx: block_controlnet_hidden_states[idx]
+ for idx, block_idx in list(enumerate(range(0, n_blocks, self.config.controlnet_block_every_n)))
+ }
+
+ # 7. Transformer blocks
+ for block_idx, block in enumerate(self.transformer_blocks):
+ controlnet_residual = controlnet_block_index_map.get(block_idx)
if torch.is_grad_enabled() and self.gradient_checkpointing:
hidden_states = self._gradient_checkpointing_func(
block,
hidden_states,
- encoder_hidden_states,
+ processed_encoder_hidden_states,
embedded_timestep,
temb,
image_rotary_emb,
extra_pos_emb,
attention_mask,
+ controlnet_residual,
)
else:
hidden_states = block(
- hidden_states=hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- embedded_timestep=embedded_timestep,
- temb=temb,
- image_rotary_emb=image_rotary_emb,
- extra_pos_emb=extra_pos_emb,
- attention_mask=attention_mask,
+ hidden_states,
+ processed_encoder_hidden_states,
+ embedded_timestep,
+ temb,
+ image_rotary_emb,
+ extra_pos_emb,
+ attention_mask,
+ controlnet_residual,
)
- # 6. Output norm & projection & unpatchify
+ # 8. Output norm & projection & unpatchify
hidden_states = self.norm_out(hidden_states, embedded_timestep, temb)
hidden_states = self.proj_out(hidden_states)
hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
diff --git a/src/diffusers/models/transformers/transformer_easyanimate.py b/src/diffusers/models/transformers/transformer_easyanimate.py
index 545fa29730db..a665d420c230 100755
--- a/src/diffusers/models/transformers/transformer_easyanimate.py
+++ b/src/diffusers/models/transformers/transformer_easyanimate.py
@@ -13,8 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import torch
import torch.nn.functional as F
from torch import nn
@@ -58,7 +56,7 @@ def __init__(
def forward(
self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
shift, scale, gate, enc_shift, enc_scale, enc_gate = self.linear(self.silu(temb)).chunk(6, dim=1)
hidden_states = self.norm(hidden_states) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
encoder_hidden_states = self.norm(encoder_hidden_states) * (1 + enc_scale.unsqueeze(1)) + enc_shift.unsqueeze(
@@ -68,7 +66,7 @@ def forward(
class EasyAnimateRotaryPosEmbed(nn.Module):
- def __init__(self, patch_size: int, rope_dim: List[int]) -> None:
+ def __init__(self, patch_size: int, rope_dim: list[int]) -> None:
super().__init__()
self.patch_size = patch_size
@@ -128,8 +126,8 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
if attn.add_q_proj is None and encoder_hidden_states is not None:
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
@@ -220,7 +218,7 @@ def __init__(
norm_elementwise_affine: bool = True,
norm_eps: float = 1e-6,
final_dropout: bool = True,
- ff_inner_dim: Optional[int] = None,
+ ff_inner_dim: int | None = None,
ff_bias: bool = True,
qk_norm: bool = True,
after_norm: bool = False,
@@ -280,8 +278,8 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# 1. Attention
norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
hidden_states, encoder_hidden_states, temb
@@ -375,9 +373,9 @@ def __init__(
self,
num_attention_heads: int = 48,
attention_head_dim: int = 64,
- in_channels: Optional[int] = None,
- out_channels: Optional[int] = None,
- patch_size: Optional[int] = None,
+ in_channels: int | None = None,
+ out_channels: int | None = None,
+ patch_size: int | None = None,
sample_width: int = 90,
sample_height: int = 60,
activation_fn: str = "gelu-approximate",
@@ -464,13 +462,13 @@ def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.Tensor,
- timestep_cond: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_hidden_states_t5: Optional[torch.Tensor] = None,
- inpaint_latents: Optional[torch.Tensor] = None,
- control_latents: Optional[torch.Tensor] = None,
+ timestep_cond: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_hidden_states_t5: torch.Tensor | None = None,
+ inpaint_latents: torch.Tensor | None = None,
+ control_latents: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, channels, video_length, height, width = hidden_states.size()
p = self.config.patch_size
post_patch_height = height // p
diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py
index 37b995867fd0..78a77ebcfea9 100644
--- a/src/diffusers/models/transformers/transformer_flux.py
+++ b/src/diffusers/models/transformers/transformer_flux.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -22,7 +22,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
@@ -85,8 +85,8 @@ def __call__(
attn: "FluxAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
attn, hidden_states, encoder_hidden_states
@@ -185,10 +185,10 @@ def __call__(
attn: "FluxAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
- ip_hidden_states: Optional[List[torch.Tensor]] = None,
- ip_adapter_masks: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
+ ip_hidden_states: list[torch.Tensor] | None = None,
+ ip_adapter_masks: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size = hidden_states.shape[0]
@@ -286,12 +286,12 @@ def __init__(
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
- added_kv_proj_dim: Optional[int] = None,
- added_proj_bias: Optional[bool] = True,
+ added_kv_proj_dim: int | None = None,
+ added_proj_bias: bool | None = True,
out_bias: bool = True,
eps: float = 1e-5,
out_dim: int = None,
- context_pre_only: Optional[bool] = None,
+ context_pre_only: bool | None = None,
pre_only: bool = False,
elementwise_affine: bool = True,
processor=None,
@@ -336,9 +336,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -379,9 +379,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_len = encoder_hidden_states.shape[1]
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
@@ -439,9 +439,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -493,7 +493,7 @@ def forward(
class FluxPosEmbed(nn.Module):
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
- def __init__(self, theta: int, axes_dim: List[int]):
+ def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -558,7 +558,7 @@ class FluxTransformer2DModel(
The number of dimensions to use for the pooled projection.
guidance_embeds (`bool`, defaults to `False`):
Whether to use guidance embeddings for guidance-distilled variant of the model.
- axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`):
The dimensions to use for the rotary positional embeddings.
"""
@@ -581,7 +581,7 @@ def __init__(
self,
patch_size: int = 1,
in_channels: int = 64,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 19,
num_single_layers: int = 38,
attention_head_dim: int = 128,
@@ -589,7 +589,7 @@ def __init__(
joint_attention_dim: int = 4096,
pooled_projection_dim: int = 768,
guidance_embeds: bool = False,
- axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+ axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
):
super().__init__()
self.out_channels = out_channels or in_channels
@@ -634,6 +634,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -643,12 +644,12 @@ def forward(
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
guidance: torch.Tensor = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
controlnet_block_samples=None,
controlnet_single_block_samples=None,
return_dict: bool = True,
controlnet_blocks_repeat: bool = False,
- ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ ) -> torch.Tensor | Transformer2DModelOutput:
"""
The [`FluxTransformer2DModel`] forward method.
@@ -675,20 +676,6 @@ def forward(
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
- if joint_attention_kwargs is not None:
- joint_attention_kwargs = joint_attention_kwargs.copy()
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
hidden_states = self.x_embedder(hidden_states)
@@ -785,10 +772,6 @@ def forward(
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_flux2.py b/src/diffusers/models/transformers/transformer_flux2.py
index 9cadfcefc497..f77498c74fc1 100644
--- a/src/diffusers/models/transformers/transformer_flux2.py
+++ b/src/diffusers/models/transformers/transformer_flux2.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
from ..attention import AttentionMixin, AttentionModuleMixin
from ..attention_dispatch import dispatch_attention_fn
@@ -90,9 +90,9 @@ class Flux2FeedForward(nn.Module):
def __init__(
self,
dim: int,
- dim_out: Optional[int] = None,
+ dim_out: int | None = None,
mult: float = 3.0,
- inner_dim: Optional[int] = None,
+ inner_dim: int | None = None,
bias: bool = False,
):
super().__init__()
@@ -125,8 +125,8 @@ def __call__(
attn: "Flux2Attention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
attn, hidden_states, encoder_hidden_states
@@ -192,8 +192,8 @@ def __init__(
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
- added_kv_proj_dim: Optional[int] = None,
- added_proj_bias: Optional[bool] = True,
+ added_kv_proj_dim: int | None = None,
+ added_proj_bias: bool | None = True,
out_bias: bool = True,
eps: float = 1e-5,
out_dim: int = None,
@@ -241,9 +241,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -268,8 +268,8 @@ def __call__(
self,
attn: "Flux2ParallelSelfAttention",
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
# Parallel in (QKV + MLP in) projection
hidden_states = attn.to_qkv_mlp_proj(hidden_states)
@@ -376,8 +376,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -423,20 +423,20 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor],
- temb_mod_params: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ encoder_hidden_states: torch.Tensor | None,
+ temb_mod: torch.Tensor,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
split_hidden_states: bool = False,
- text_seq_len: Optional[int] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ text_seq_len: int | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already
# concatenated
if encoder_hidden_states is not None:
text_seq_len = encoder_hidden_states.shape[1]
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
- mod_shift, mod_scale, mod_gate = temb_mod_params
+ mod_shift, mod_scale, mod_gate = Flux2Modulation.split(temb_mod, 1)[0]
norm_hidden_states = self.norm(hidden_states)
norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift
@@ -498,16 +498,18 @@ def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- temb_mod_params_img: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...],
- temb_mod_params_txt: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...],
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ temb_mod_img: torch.Tensor,
+ temb_mod_txt: torch.Tensor,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
joint_attention_kwargs = joint_attention_kwargs or {}
# Modulation parameters shape: [1, 1, self.dim]
- (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = temb_mod_params_img
- (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = temb_mod_params_txt
+ (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = Flux2Modulation.split(temb_mod_img, 2)
+ (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = Flux2Modulation.split(
+ temb_mod_txt, 2
+ )
# Img stream
norm_hidden_states = self.norm1(hidden_states)
@@ -554,7 +556,7 @@ def forward(
class Flux2PosEmbed(nn.Module):
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
- def __init__(self, theta: int, axes_dim: List[int]):
+ def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -627,15 +629,19 @@ def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False):
self.linear = nn.Linear(dim, dim * 3 * self.mod_param_sets, bias=bias)
self.act_fn = nn.SiLU()
- def forward(self, temb: torch.Tensor) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]:
+ def forward(self, temb: torch.Tensor) -> torch.Tensor:
mod = self.act_fn(temb)
mod = self.linear(mod)
+ return mod
+ @staticmethod
+ # split inside the transformer blocks, to avoid passing tuples into checkpoints https://github.com/huggingface/diffusers/issues/12776
+ def split(mod: torch.Tensor, mod_param_sets: int) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]:
if mod.ndim == 2:
mod = mod.unsqueeze(1)
- mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1)
+ mod_params = torch.chunk(mod, 3 * mod_param_sets, dim=-1)
# Return tuple of 3-tuples of modulation params shift/scale/gate
- return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets))
+ return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(mod_param_sets))
class Flux2Transformer2DModel(
@@ -674,7 +680,7 @@ class Flux2Transformer2DModel(
The number of dimensions to use for the pooled projection.
guidance_embeds (`bool`, defaults to `True`):
Whether to use guidance embeddings for guidance-distilled variant of the model.
- axes_dims_rope (`Tuple[int]`, defaults to `(32, 32, 32, 32)`):
+ axes_dims_rope (`tuple[int]`, defaults to `(32, 32, 32, 32)`):
The dimensions to use for the rotary positional embeddings.
"""
@@ -697,7 +703,7 @@ def __init__(
self,
patch_size: int = 1,
in_channels: int = 128,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 8,
num_single_layers: int = 48,
attention_head_dim: int = 128,
@@ -705,7 +711,7 @@ def __init__(
joint_attention_dim: int = 15360,
timestep_guidance_channels: int = 256,
mlp_ratio: float = 3.0,
- axes_dims_rope: Tuple[int, ...] = (32, 32, 32, 32),
+ axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
rope_theta: int = 2000,
eps: float = 1e-6,
guidance_embeds: bool = True,
@@ -774,6 +780,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -782,9 +789,9 @@ def forward(
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
guidance: torch.Tensor = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ ) -> torch.Tensor | Transformer2DModelOutput:
"""
The [`FluxTransformer2DModel`] forward method.
@@ -810,20 +817,6 @@ def forward(
`tuple` where the first element is the sample tensor.
"""
# 0. Handle input arguments
- if joint_attention_kwargs is not None:
- joint_attention_kwargs = joint_attention_kwargs.copy()
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
num_txt_tokens = encoder_hidden_states.shape[1]
@@ -837,7 +830,7 @@ def forward(
double_stream_mod_img = self.double_stream_modulation_img(temb)
double_stream_mod_txt = self.double_stream_modulation_txt(temb)
- single_stream_mod = self.single_stream_modulation(temb)[0]
+ single_stream_mod = self.single_stream_modulation(temb)
# 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
hidden_states = self.x_embedder(hidden_states)
@@ -874,8 +867,8 @@ def forward(
encoder_hidden_states, hidden_states = block(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
- temb_mod_params_img=double_stream_mod_img,
- temb_mod_params_txt=double_stream_mod_txt,
+ temb_mod_img=double_stream_mod_img,
+ temb_mod_txt=double_stream_mod_txt,
image_rotary_emb=concat_rotary_emb,
joint_attention_kwargs=joint_attention_kwargs,
)
@@ -897,7 +890,7 @@ def forward(
hidden_states = block(
hidden_states=hidden_states,
encoder_hidden_states=None,
- temb_mod_params=single_stream_mod,
+ temb_mod=single_stream_mod,
image_rotary_emb=concat_rotary_emb,
joint_attention_kwargs=joint_attention_kwargs,
)
@@ -908,10 +901,6 @@ def forward(
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_glm_image.py b/src/diffusers/models/transformers/transformer_glm_image.py
index 6f7ed2fca1c9..8413b24fef45 100644
--- a/src/diffusers/models/transformers/transformer_glm_image.py
+++ b/src/diffusers/models/transformers/transformer_glm_image.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -104,7 +104,7 @@ def __init__(self, embedding_dim: int, dim: int) -> None:
def forward(
self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
dtype = hidden_states.dtype
norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype)
@@ -148,9 +148,9 @@ class GlmImageLayerKVCache:
"""
def __init__(self):
- self.k_caches: List[Optional[torch.Tensor]] = []
- self.v_caches: List[Optional[torch.Tensor]] = []
- self.mode: Optional[str] = None # "write", "read", "skip"
+ self.k_caches: list[torch.Tensor | None] = []
+ self.v_caches: list[torch.Tensor | None] = []
+ self.mode: str | None = None # "write", "read", "skip"
self.current_sample_idx: int = 0 # Current sample index for writing
def store(self, k: torch.Tensor, v: torch.Tensor):
@@ -231,7 +231,7 @@ def __init__(self, num_layers: int):
def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
return self.caches[layer_idx]
- def set_mode(self, mode: Optional[str]):
+ def set_mode(self, mode: str):
if mode is not None and mode not in ["write", "read", "skip"]:
raise ValueError(f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'")
for cache in self.caches:
@@ -269,10 +269,10 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- kv_cache: Optional[GlmImageLayerKVCache] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ kv_cache: GlmImageLayerKVCache | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
dtype = encoder_hidden_states.dtype
batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
@@ -381,14 +381,12 @@ def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[
- Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]]
- ] = None,
- attention_mask: Optional[Dict[str, torch.Tensor]] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- kv_cache: Optional[GlmImageLayerKVCache] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ temb: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None,
+ attention_mask: dict[str, torch.Tensor] | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ kv_cache: GlmImageLayerKVCache | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# 1. Timestep conditioning
(
norm_hidden_states,
@@ -439,7 +437,7 @@ def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
self.patch_size = patch_size
self.theta = theta
- def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
batch_size, num_channels, height, width = hidden_states.shape
height, width = height // self.patch_size, width // self.patch_size
@@ -604,14 +602,12 @@ def forward(
timestep: torch.LongTensor,
target_size: torch.Tensor,
crop_coords: torch.Tensor,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- attention_mask: Optional[torch.Tensor] = None,
- kv_caches: Optional[GlmImageKVCache] = None,
- image_rotary_emb: Optional[
- Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]]
- ] = None,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
+ attention_mask: torch.Tensor | None = None,
+ kv_caches: GlmImageKVCache | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None,
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_channels, height, width = hidden_states.shape
# 1. RoPE
diff --git a/src/diffusers/models/transformers/transformer_hidream_image.py b/src/diffusers/models/transformers/transformer_hidream_image.py
index 4a5aee29abc4..57c0a16eee6f 100644
--- a/src/diffusers/models/transformers/transformer_hidream_image.py
+++ b/src/diffusers/models/transformers/transformer_hidream_image.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -8,7 +8,7 @@
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
from ...models.modeling_outputs import Transformer2DModelOutput
from ...models.modeling_utils import ModelMixin
-from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, deprecate, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import Attention
from ..embeddings import TimestepEmbedding, Timesteps
@@ -23,7 +23,7 @@ def __init__(
dim: int,
hidden_dim: int,
multiple_of: int = 256,
- ffn_dim_multiplier: Optional[float] = None,
+ ffn_dim_multiplier: float | None = None,
):
super().__init__()
hidden_dim = int(2 * hidden_dim / 3)
@@ -55,7 +55,7 @@ def __init__(self, hidden_size, frequency_embedding_size=256):
self.time_proj = Timesteps(num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0)
self.timestep_embedder = TimestepEmbedding(in_channels=frequency_embedding_size, time_embed_dim=hidden_size)
- def forward(self, timesteps: torch.Tensor, wdtype: Optional[torch.dtype] = None) -> torch.Tensor:
+ def forward(self, timesteps: torch.Tensor, wdtype: torch.dtype | None = None) -> torch.Tensor:
t_emb = self.time_proj(timesteps).to(dtype=wdtype)
t_emb = self.timestep_embedder(t_emb)
return t_emb
@@ -114,7 +114,7 @@ def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
class HiDreamImageEmbedND(nn.Module):
- def __init__(self, theta: int, axes_dim: List[int]):
+ def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -128,7 +128,7 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor:
return emb.unsqueeze(2)
-def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
@@ -205,8 +205,8 @@ def __call__(
self,
attn: HiDreamAttention,
hidden_states: torch.Tensor,
- hidden_states_masks: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
+ hidden_states_masks: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
image_rotary_emb: torch.Tensor = None,
*args,
**kwargs,
@@ -458,9 +458,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- hidden_states_masks: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ hidden_states_masks: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
image_rotary_emb: torch.Tensor = None,
) -> torch.Tensor:
wtype = hidden_states.dtype
@@ -530,11 +530,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- hidden_states_masks: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ hidden_states_masks: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
image_rotary_emb: torch.Tensor = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
wtype = hidden_states.dtype
(
shift_msa_i,
@@ -581,18 +581,18 @@ def forward(
class HiDreamBlock(nn.Module):
- def __init__(self, block: Union[HiDreamImageTransformerBlock, HiDreamImageSingleTransformerBlock]):
+ def __init__(self, block: HiDreamImageTransformerBlock | HiDreamImageSingleTransformerBlock):
super().__init__()
self.block = block
def forward(
self,
hidden_states: torch.Tensor,
- hidden_states_masks: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- temb: Optional[torch.Tensor] = None,
+ hidden_states_masks: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ temb: torch.Tensor | None = None,
image_rotary_emb: torch.Tensor = None,
- ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
return self.block(
hidden_states=hidden_states,
hidden_states_masks=hidden_states_masks,
@@ -609,20 +609,20 @@ class HiDreamImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin,
@register_to_config
def __init__(
self,
- patch_size: Optional[int] = None,
+ patch_size: int | None = None,
in_channels: int = 64,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 16,
num_single_layers: int = 32,
attention_head_dim: int = 128,
num_attention_heads: int = 20,
- caption_channels: List[int] = None,
+ caption_channels: list[int] = None,
text_emb_dim: int = 2048,
num_routed_experts: int = 4,
num_activated_experts: int = 2,
- axes_dims_rope: Tuple[int, int] = (32, 32),
- max_resolution: Tuple[int, int] = (128, 128),
- llama_layers: List[int] = None,
+ axes_dims_rope: tuple[int, int] = (32, 32),
+ max_resolution: tuple[int, int] = (128, 128),
+ llama_layers: list[int] = None,
force_inference_output: bool = False,
):
super().__init__()
@@ -681,7 +681,7 @@ def __init__(
self.gradient_checkpointing = False
- def unpatchify(self, x: torch.Tensor, img_sizes: List[Tuple[int, int]], is_training: bool) -> List[torch.Tensor]:
+ def unpatchify(self, x: torch.Tensor, img_sizes: list[tuple[int, int]], is_training: bool) -> list[torch.Tensor]:
if is_training and not self.config.force_inference_output:
B, S, F = x.shape
C = F // (self.config.patch_size * self.config.patch_size)
@@ -773,6 +773,7 @@ def patchify(self, hidden_states):
return hidden_states, hidden_states_masks, img_sizes, img_ids
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -780,13 +781,13 @@ def forward(
encoder_hidden_states_t5: torch.Tensor = None,
encoder_hidden_states_llama3: torch.Tensor = None,
pooled_embeds: torch.Tensor = None,
- img_ids: Optional[torch.Tensor] = None,
- img_sizes: Optional[List[Tuple[int, int]]] = None,
- hidden_states_masks: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ img_ids: torch.Tensor | None = None,
+ img_sizes: list[tuple[int, int]] | None = None,
+ hidden_states_masks: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
**kwargs,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
encoder_hidden_states = kwargs.get("encoder_hidden_states", None)
if encoder_hidden_states is not None:
@@ -808,21 +809,6 @@ def forward(
"if `hidden_states_masks` is passed, `hidden_states` must be a 3D tensors with shape (batch_size, patch_height * patch_width, patch_size * patch_size * channels)"
)
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
# spatial forward
batch_size = hidden_states.shape[0]
hidden_states_type = hidden_states.dtype
@@ -933,10 +919,6 @@ def forward(
if hidden_states_masks is not None:
hidden_states_masks = hidden_states_masks[:, :image_tokens_seq_len]
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video.py b/src/diffusers/models/transformers/transformer_hunyuan_video.py
index 4f0775ac9fa0..aab593c93643 100644
--- a/src/diffusers/models/transformers/transformer_hunyuan_video.py
+++ b/src/diffusers/models/transformers/transformer_hunyuan_video.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -22,7 +22,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ..attention import AttentionMixin, FeedForward
from ..attention_dispatch import dispatch_attention_fn
from ..attention_processor import Attention
@@ -56,9 +56,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
if attn.add_q_proj is None and encoder_hidden_states is not None:
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
@@ -162,7 +162,7 @@ def __call__(
class HunyuanVideoPatchEmbed(nn.Module):
def __init__(
self,
- patch_size: Union[int, Tuple[int, int, int]] = 16,
+ patch_size: int | tuple[int, int, int] = 16,
in_chans: int = 3,
embed_dim: int = 768,
) -> None:
@@ -178,7 +178,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
class HunyuanVideoAdaNorm(nn.Module):
- def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
+ def __init__(self, in_features: int, out_features: int | None = None) -> None:
super().__init__()
out_features = out_features or 2 * in_features
@@ -187,7 +187,7 @@ def __init__(self, in_features: int, out_features: Optional[int] = None) -> None
def forward(
self, temb: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
temb = self.linear(self.nonlinearity(temb))
gate_msa, gate_mlp = temb.chunk(2, dim=1)
gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
@@ -216,7 +216,7 @@ def forward(
emb: torch.Tensor,
token_replace_emb: torch.Tensor,
first_frame_num_tokens: int,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
emb = self.linear(self.silu(emb))
token_replace_emb = self.linear(self.silu(token_replace_emb))
@@ -267,7 +267,7 @@ def forward(
emb: torch.Tensor,
token_replace_emb: torch.Tensor,
first_frame_num_tokens: int,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
emb = self.linear(self.silu(emb))
token_replace_emb = self.linear(self.silu(token_replace_emb))
@@ -292,7 +292,7 @@ def __init__(
embedding_dim: int,
pooled_projection_dim: int,
guidance_embeds: bool,
- image_condition_type: Optional[str] = None,
+ image_condition_type: str | None = None,
):
super().__init__()
@@ -307,8 +307,8 @@ def __init__(
self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
def forward(
- self, timestep: torch.Tensor, pooled_projection: torch.Tensor, guidance: Optional[torch.Tensor] = None
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ self, timestep: torch.Tensor, pooled_projection: torch.Tensor, guidance: torch.Tensor | None = None
+ ) -> tuple[torch.Tensor, torch.Tensor]:
timesteps_proj = self.time_proj(timestep)
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) # (N, D)
pooled_projections = self.text_embedder(pooled_projection)
@@ -360,7 +360,7 @@ def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
norm_hidden_states = self.norm1(hidden_states)
@@ -408,7 +408,7 @@ def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
) -> None:
self_attn_mask = None
if attention_mask is not None:
@@ -458,7 +458,7 @@ def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
- attention_mask: Optional[torch.LongTensor] = None,
+ attention_mask: torch.LongTensor | None = None,
) -> torch.Tensor:
if attention_mask is None:
pooled_projections = hidden_states.mean(dim=1)
@@ -476,7 +476,7 @@ def forward(
class HunyuanVideoRotaryPosEmbed(nn.Module):
- def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None:
+ def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], theta: float = 256.0) -> None:
super().__init__()
self.patch_size = patch_size
@@ -544,11 +544,11 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
*args,
**kwargs,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_length = encoder_hidden_states.shape[1]
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
@@ -624,11 +624,11 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
*args,
**kwargs,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# 1. Input normalization
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -699,11 +699,11 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
token_replace_emb: torch.Tensor = None,
num_tokens: int = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_length = encoder_hidden_states.shape[1]
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
@@ -783,11 +783,11 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
token_replace_emb: torch.Tensor = None,
num_tokens: int = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# 1. Input normalization
(
norm_hidden_states,
@@ -875,7 +875,7 @@ class HunyuanVideoTransformer3DModel(
The dimension of the pooled projection of the text embeddings.
rope_theta (`float`, defaults to `256.0`):
The value of theta to use in the RoPE layer.
- rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ rope_axes_dim (`tuple[int]`, defaults to `(16, 56, 56)`):
The dimensions of the axes to use in the RoPE layer.
image_condition_type (`str`, *optional*, defaults to `None`):
The type of image conditioning to use. If `None`, no image conditioning is used. If `latent_concat`, the
@@ -916,8 +916,8 @@ def __init__(
text_embed_dim: int = 4096,
pooled_projection_dim: int = 768,
rope_theta: float = 256.0,
- rope_axes_dim: Tuple[int, ...] = (16, 56, 56),
- image_condition_type: Optional[str] = None,
+ rope_axes_dim: tuple[int, ...] = (16, 56, 56),
+ image_condition_type: str | None = None,
) -> None:
super().__init__()
@@ -989,6 +989,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -997,24 +998,9 @@ def forward(
encoder_attention_mask: torch.Tensor,
pooled_projections: torch.Tensor,
guidance: torch.Tensor = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p, p_t = self.config.patch_size, self.config.patch_size_t
post_patch_num_frames = num_frames // p_t
@@ -1104,10 +1090,6 @@ def forward(
hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (hidden_states,)
diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video15.py b/src/diffusers/models/transformers/transformer_hunyuan_video15.py
index 293ba996ea98..222b0791d650 100644
--- a/src/diffusers/models/transformers/transformer_hunyuan_video15.py
+++ b/src/diffusers/models/transformers/transformer_hunyuan_video15.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -22,7 +22,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ..attention import AttentionMixin, FeedForward
from ..attention_dispatch import dispatch_attention_fn
from ..attention_processor import Attention
@@ -55,9 +55,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
# 1. QKV projections
query = attn.to_q(hidden_states)
@@ -140,7 +140,7 @@ def __call__(
class HunyuanVideo15PatchEmbed(nn.Module):
def __init__(
self,
- patch_size: Union[int, Tuple[int, int, int]] = 16,
+ patch_size: int | tuple[int, int, int] = 16,
in_chans: int = 3,
embed_dim: int = 768,
) -> None:
@@ -156,7 +156,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
class HunyuanVideo15AdaNorm(nn.Module):
- def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
+ def __init__(self, in_features: int, out_features: int | None = None) -> None:
super().__init__()
out_features = out_features or 2 * in_features
@@ -165,7 +165,7 @@ def __init__(self, in_features: int, out_features: Optional[int] = None) -> None
def forward(
self, temb: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
temb = self.linear(self.nonlinearity(temb))
gate_msa, gate_mlp = temb.chunk(2, dim=1)
gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
@@ -200,7 +200,7 @@ def __init__(self, embedding_dim: int, use_meanflow: bool = False):
def forward(
self,
timestep: torch.Tensor,
- timestep_r: Optional[torch.Tensor] = None,
+ timestep_r: torch.Tensor | None = None,
) -> torch.Tensor:
timesteps_proj = self.time_proj(timestep)
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=timestep.dtype))
@@ -244,7 +244,7 @@ def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
norm_hidden_states = self.norm1(hidden_states)
@@ -292,7 +292,7 @@ def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
) -> None:
self_attn_mask = None
if attention_mask is not None:
@@ -341,7 +341,7 @@ def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
- attention_mask: Optional[torch.LongTensor] = None,
+ attention_mask: torch.LongTensor | None = None,
) -> torch.Tensor:
if attention_mask is None:
pooled_projections = hidden_states.mean(dim=1)
@@ -359,7 +359,7 @@ def forward(
class HunyuanVideo15RotaryPosEmbed(nn.Module):
- def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None:
+ def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], theta: float = 256.0) -> None:
super().__init__()
self.patch_size = patch_size
@@ -468,11 +468,11 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
*args,
**kwargs,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# 1. Input normalization
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -542,7 +542,7 @@ class HunyuanVideo15Transformer3DModel(
The dimension of the pooled projection of the text embeddings.
rope_theta (`float`, defaults to `256.0`):
The value of theta to use in the RoPE layer.
- rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ rope_axes_dim (`tuple[int]`, defaults to `(16, 56, 56)`):
The dimensions of the axes to use in the RoPE layer.
"""
@@ -576,7 +576,7 @@ def __init__(
text_embed_2_dim: int = 1472,
image_embed_dim: int = 1152,
rope_theta: float = 256.0,
- rope_axes_dim: Tuple[int, ...] = (16, 56, 56),
+ rope_axes_dim: tuple[int, ...] = (16, 56, 56),
# YiYi Notes: config based on target_size_config https://github.com/yiyixuxu/hy15/blob/main/hyvideo/pipelines/hunyuan_video_pipeline.py#L205
target_size: int = 640, # did not name sample_size since it is in pixel spaces
task_type: str = "i2v",
@@ -620,34 +620,20 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
encoder_attention_mask: torch.Tensor,
- timestep_r: Optional[torch.LongTensor] = None,
- encoder_hidden_states_2: Optional[torch.Tensor] = None,
- encoder_attention_mask_2: Optional[torch.Tensor] = None,
- image_embeds: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ timestep_r: torch.LongTensor | None = None,
+ encoder_hidden_states_2: torch.Tensor | None = None,
+ encoder_attention_mask_2: torch.Tensor | None = None,
+ image_embeds: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
post_patch_num_frames = num_frames // p_t
@@ -783,10 +769,6 @@ def forward(
hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (hidden_states,)
diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py b/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py
index 601ba0f0b472..f005c4d4cd51 100644
--- a/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py
+++ b/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -20,7 +20,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, get_logger, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, get_logger
from ..cache_utils import CacheMixin
from ..embeddings import get_1d_rotary_pos_embed
from ..modeling_outputs import Transformer2DModelOutput
@@ -39,7 +39,7 @@
class HunyuanVideoFramepackRotaryPosEmbed(nn.Module):
- def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None:
+ def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], theta: float = 256.0) -> None:
super().__init__()
self.patch_size = patch_size
@@ -91,9 +91,9 @@ def __init__(self, in_channels: int, inner_dim: int):
def forward(
self,
- latents_clean: Optional[torch.Tensor] = None,
- latents_clean_2x: Optional[torch.Tensor] = None,
- latents_clean_4x: Optional[torch.Tensor] = None,
+ latents_clean: torch.Tensor | None = None,
+ latents_clean_2x: torch.Tensor | None = None,
+ latents_clean_4x: torch.Tensor | None = None,
):
if latents_clean is not None:
latents_clean = self.proj(latents_clean)
@@ -139,8 +139,8 @@ def __init__(
text_embed_dim: int = 4096,
pooled_projection_dim: int = 768,
rope_theta: float = 256.0,
- rope_axes_dim: Tuple[int, ...] = (16, 56, 56),
- image_condition_type: Optional[str] = None,
+ rope_axes_dim: tuple[int, ...] = (16, 56, 56),
+ image_condition_type: str | None = None,
has_image_proj: int = False,
image_proj_dim: int = 1152,
has_clean_x_embedder: int = False,
@@ -198,6 +198,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -207,31 +208,16 @@ def forward(
pooled_projections: torch.Tensor,
image_embeds: torch.Tensor,
indices_latents: torch.Tensor,
- guidance: Optional[torch.Tensor] = None,
- latents_clean: Optional[torch.Tensor] = None,
- indices_latents_clean: Optional[torch.Tensor] = None,
- latents_history_2x: Optional[torch.Tensor] = None,
- indices_latents_history_2x: Optional[torch.Tensor] = None,
- latents_history_4x: Optional[torch.Tensor] = None,
- indices_latents_history_4x: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ guidance: torch.Tensor | None = None,
+ latents_clean: torch.Tensor | None = None,
+ indices_latents_clean: torch.Tensor | None = None,
+ latents_history_2x: torch.Tensor | None = None,
+ indices_latents_history_2x: torch.Tensor | None = None,
+ latents_history_4x: torch.Tensor | None = None,
+ indices_latents_history_4x: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p, p_t = self.config.patch_size, self.config.patch_size_t
post_patch_num_frames = num_frames // p_t
@@ -337,10 +323,6 @@ def forward(
hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (hidden_states,)
return Transformer2DModelOutput(sample=hidden_states)
@@ -348,13 +330,13 @@ def forward(
def _pack_history_states(
self,
hidden_states: torch.Tensor,
- latents_clean: Optional[torch.Tensor] = None,
- latents_history_2x: Optional[torch.Tensor] = None,
- latents_history_4x: Optional[torch.Tensor] = None,
- image_rotary_emb: Tuple[torch.Tensor, torch.Tensor] = None,
- image_rotary_emb_clean: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- image_rotary_emb_history_2x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- image_rotary_emb_history_4x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ latents_clean: torch.Tensor | None = None,
+ latents_history_2x: torch.Tensor | None = None,
+ latents_history_4x: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] = None,
+ image_rotary_emb_clean: tuple[torch.Tensor, torch.Tensor] | None = None,
+ image_rotary_emb_history_2x: tuple[torch.Tensor, torch.Tensor] | None = None,
+ image_rotary_emb_history_4x: tuple[torch.Tensor, torch.Tensor] | None = None,
height: int = None,
width: int = None,
):
@@ -381,10 +363,10 @@ def _pack_history_states(
def _pad_rotary_emb(
self,
- image_rotary_emb: Tuple[torch.Tensor],
+ image_rotary_emb: tuple[torch.Tensor],
height: int,
width: int,
- kernel_size: Tuple[int, int, int],
+ kernel_size: tuple[int, int, int],
):
# freqs_cos, freqs_sin have shape [W * H * T, D / 2], where D is attention head dim
freqs_cos, freqs_sin = image_rotary_emb
diff --git a/src/diffusers/models/transformers/transformer_hunyuanimage.py b/src/diffusers/models/transformers/transformer_hunyuanimage.py
index d626e322ad6f..a2d3d9229963 100644
--- a/src/diffusers/models/transformers/transformer_hunyuanimage.py
+++ b/src/diffusers/models/transformers/transformer_hunyuanimage.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -23,7 +23,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import AttentionMixin, FeedForward
from ..attention_dispatch import dispatch_attention_fn
@@ -57,9 +57,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
if attn.add_q_proj is None and encoder_hidden_states is not None:
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
@@ -157,7 +157,7 @@ def __call__(
class HunyuanImagePatchEmbed(nn.Module):
def __init__(
self,
- patch_size: Union[Tuple[int, int], Tuple[int, int, int]] = (16, 16),
+ patch_size: tuple[int, int, tuple[int, int, int]] = (16, 16),
in_chans: int = 3,
embed_dim: int = 768,
) -> None:
@@ -198,7 +198,7 @@ def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
class HunyuanImageAdaNorm(nn.Module):
- def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
+ def __init__(self, in_features: int, out_features: int | None = None) -> None:
super().__init__()
out_features = out_features or 2 * in_features
@@ -207,7 +207,7 @@ def __init__(self, in_features: int, out_features: Optional[int] = None) -> None
def forward(
self, temb: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
temb = self.linear(self.nonlinearity(temb))
gate_msa, gate_mlp = temb.chunk(2, dim=1)
gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
@@ -241,9 +241,9 @@ def __init__(
def forward(
self,
timestep: torch.Tensor,
- timestep_r: Optional[torch.Tensor] = None,
- guidance: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ timestep_r: torch.Tensor | None = None,
+ guidance: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
timesteps_proj = self.time_proj(timestep)
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=timestep.dtype))
@@ -295,7 +295,7 @@ def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
norm_hidden_states = self.norm1(hidden_states)
@@ -343,7 +343,7 @@ def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
) -> None:
self_attn_mask = None
if attention_mask is not None:
@@ -394,7 +394,7 @@ def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
- attention_mask: Optional[torch.LongTensor] = None,
+ attention_mask: torch.LongTensor | None = None,
) -> torch.Tensor:
if attention_mask is None:
pooled_hidden_states = hidden_states.mean(dim=1)
@@ -412,9 +412,7 @@ def forward(
class HunyuanImageRotaryPosEmbed(nn.Module):
- def __init__(
- self, patch_size: Union[Tuple, List[int]], rope_dim: Union[Tuple, List[int]], theta: float = 256.0
- ) -> None:
+ def __init__(self, patch_size: tuple | list[int], rope_dim: tuple | list[int], theta: float = 256.0) -> None:
super().__init__()
if not isinstance(patch_size, (tuple, list)) or len(patch_size) not in [2, 3]:
@@ -496,8 +494,8 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -577,11 +575,11 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
*args,
**kwargs,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# 1. Input normalization
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -653,7 +651,7 @@ class HunyuanImageTransformer2DModel(
The dimension of the pooled projection of the text embeddings.
rope_theta (`float`, defaults to `256.0`):
The value of theta to use in the RoPE layer.
- rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ rope_axes_dim (`tuple[int]`, defaults to `(16, 56, 56)`):
The dimensions of the axes to use in the RoPE layer.
image_condition_type (`str`, *optional*, defaults to `None`):
The type of image conditioning to use. If `None`, no image conditioning is used. If `latent_concat`, the
@@ -682,13 +680,13 @@ def __init__(
num_single_layers: int = 40,
num_refiner_layers: int = 2,
mlp_ratio: float = 4.0,
- patch_size: Tuple[int, int] = (1, 1),
+ patch_size: tuple[int, int] = (1, 1),
qk_norm: str = "rms_norm",
guidance_embeds: bool = False,
text_embed_dim: int = 3584,
- text_embed_2_dim: Optional[int] = None,
+ text_embed_2_dim: int | None = None,
rope_theta: float = 256.0,
- rope_axes_dim: Tuple[int, ...] = (64, 64),
+ rope_axes_dim: tuple[int, ...] = (64, 64),
use_meanflow: bool = False,
) -> None:
super().__init__()
@@ -742,34 +740,20 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
encoder_attention_mask: torch.Tensor,
- timestep_r: Optional[torch.LongTensor] = None,
- encoder_hidden_states_2: Optional[torch.Tensor] = None,
- encoder_attention_mask_2: Optional[torch.Tensor] = None,
- guidance: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ timestep_r: torch.LongTensor | None = None,
+ encoder_hidden_states_2: torch.Tensor | None = None,
+ encoder_attention_mask_2: torch.Tensor | None = None,
+ guidance: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> torch.Tensor | dict[str, torch.Tensor]:
if hidden_states.ndim == 4:
batch_size, channels, height, width = hidden_states.shape
sizes = (height, width)
@@ -900,10 +884,6 @@ def forward(
]
hidden_states = hidden_states.reshape(*final_dims)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (hidden_states,)
diff --git a/src/diffusers/models/transformers/transformer_kandinsky.py b/src/diffusers/models/transformers/transformer_kandinsky.py
index 57b28991d255..88ef70d546c8 100644
--- a/src/diffusers/models/transformers/transformer_kandinsky.py
+++ b/src/diffusers/models/transformers/transformer_kandinsky.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -366,9 +366,9 @@ def __init__(self, num_channels, head_dim, processor=None):
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- sparse_params: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ sparse_params: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -596,12 +596,12 @@ def forward(
encoder_hidden_states: torch.Tensor, # text_embed
timestep: torch.Tensor, # time
pooled_projections: torch.Tensor, # pooled_text_embed
- visual_rope_pos: Tuple[int, int, int],
+ visual_rope_pos: tuple[int, int, int],
text_rope_pos: torch.LongTensor,
- scale_factor: Tuple[float, float, float] = (1.0, 1.0, 1.0),
- sparse_params: Optional[Dict[str, Any]] = None,
+ scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0),
+ sparse_params: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[Transformer2DModelOutput, torch.FloatTensor]:
+ ) -> Transformer2DModelOutput | torch.FloatTensor:
"""
Forward pass of the Kandinsky5 3D Transformer.
@@ -610,10 +610,10 @@ def forward(
encoder_hidden_states (`torch.FloatTensor`): Text embeddings
timestep (`torch.Tensor` or `float` or `int`): Current timestep
pooled_projections (`torch.FloatTensor`): Pooled text embeddings
- visual_rope_pos (`Tuple[int, int, int]`): Position for visual RoPE
+ visual_rope_pos (`tuple[int, int, int]`): Position for visual RoPE
text_rope_pos (`torch.LongTensor`): Position for text RoPE
- scale_factor (`Tuple[float, float, float]`, optional): Scale factor for RoPE
- sparse_params (`Dict[str, Any]`, optional): Parameters for sparse attention
+ scale_factor (`tuple[float, float, float]`, optional): Scale factor for RoPE
+ sparse_params (`dict[str, Any]`, optional): Parameters for sparse attention
return_dict (`bool`, optional): Whether to return a dictionary
Returns:
diff --git a/src/diffusers/models/transformers/transformer_longcat_image.py b/src/diffusers/models/transformers/transformer_longcat_image.py
index 3d38da1dfcf5..7a000fa2b2ce 100644
--- a/src/diffusers/models/transformers/transformer_longcat_image.py
+++ b/src/diffusers/models/transformers/transformer_longcat_image.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -78,8 +78,8 @@ def __call__(
attn: "LongCatImageAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
attn, hidden_states, encoder_hidden_states
@@ -145,12 +145,12 @@ def __init__(
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
- added_kv_proj_dim: Optional[int] = None,
- added_proj_bias: Optional[bool] = True,
+ added_kv_proj_dim: int | None = None,
+ added_proj_bias: bool | None = True,
out_bias: bool = True,
eps: float = 1e-5,
out_dim: int = None,
- context_pre_only: Optional[bool] = None,
+ context_pre_only: bool | None = None,
pre_only: bool = False,
elementwise_affine: bool = True,
processor=None,
@@ -195,9 +195,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -238,9 +238,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_len = encoder_hidden_states.shape[1]
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
@@ -298,9 +298,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -351,7 +351,7 @@ def forward(
class LongCatImagePosEmbed(nn.Module):
- def __init__(self, theta: int, axes_dim: List[int]):
+ def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -420,7 +420,7 @@ def __init__(
num_attention_heads: int = 24,
joint_attention_dim: int = 3584,
pooled_projection_dim: int = 3584,
- axes_dims_rope: List[int] = [16, 56, 56],
+ axes_dims_rope: list[int] = [16, 56, 56],
):
super().__init__()
self.out_channels = in_channels
@@ -472,7 +472,7 @@ def forward(
txt_ids: torch.Tensor = None,
guidance: torch.Tensor = None,
return_dict: bool = True,
- ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+ ) -> torch.FloatTensor | Transformer2DModelOutput:
"""
The forward method.
diff --git a/src/diffusers/models/transformers/transformer_ltx.py b/src/diffusers/models/transformers/transformer_ltx.py
index 685c73c07c75..0034d636761b 100644
--- a/src/diffusers/models/transformers/transformer_ltx.py
+++ b/src/diffusers/models/transformers/transformer_ltx.py
@@ -15,14 +15,14 @@
import inspect
import math
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, deprecate, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, deprecate, is_torch_version, logging
from ...utils.torch_utils import maybe_allow_in_graph
from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
@@ -64,9 +64,9 @@ def __call__(
self,
attn: "LTXAttention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
@@ -124,7 +124,7 @@ def __init__(
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = True,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
out_bias: bool = True,
qk_norm: str = "rms_norm_across_heads",
processor=None,
@@ -161,9 +161,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -203,7 +203,7 @@ def _prepare_video_coords(
num_frames: int,
height: int,
width: int,
- rope_interpolation_scale: Tuple[torch.Tensor, float, float],
+ rope_interpolation_scale: tuple[torch.Tensor, float, float],
device: torch.device,
) -> torch.Tensor:
# Always compute rope in fp32
@@ -226,12 +226,12 @@ def _prepare_video_coords(
def forward(
self,
hidden_states: torch.Tensor,
- num_frames: Optional[int] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- rope_interpolation_scale: Optional[Tuple[torch.Tensor, float, float]] = None,
- video_coords: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ num_frames: int | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ rope_interpolation_scale: tuple[torch.Tensor, float, float] | None = None,
+ video_coords: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
batch_size = hidden_states.size(0)
if video_coords is None:
@@ -346,8 +346,8 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size = hidden_states.size(0)
norm_hidden_states = self.norm1(hidden_states)
@@ -491,35 +491,21 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_attention_mask: torch.Tensor,
- num_frames: Optional[int] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- rope_interpolation_scale: Optional[Union[Tuple[float, float, float], torch.Tensor]] = None,
- video_coords: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ num_frames: int | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ rope_interpolation_scale: tuple[float, float, float] | torch.Tensor | None = None,
+ video_coords: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
) -> torch.Tensor:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
image_rotary_emb = self.rope(hidden_states, num_frames, height, width, rope_interpolation_scale, video_coords)
# convert encoder_attention_mask to a bias the same way we do for attention_mask
@@ -568,10 +554,6 @@ def forward(
hidden_states = hidden_states * (1 + scale) + shift
output = self.proj_out(hidden_states)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index b88f096e8033..db949ca34a1f 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -15,21 +15,14 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import (
- USE_PEFT_BACKEND,
- BaseOutput,
- is_torch_version,
- logging,
- scale_lora_layers,
- unscale_lora_layers,
-)
+from ...utils import BaseOutput, apply_lora_scale, is_torch_version, logging
from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
from ..attention_dispatch import dispatch_attention_fn
@@ -42,7 +35,7 @@
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
cos, sin = freqs
x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) # [B, S, C // 2]
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
@@ -50,16 +43,14 @@ def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: Tuple[torch.Tensor, tor
return out
-def apply_split_rotary_emb(x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
cos, sin = freqs
x_dtype = x.dtype
needs_reshape = False
if x.ndim != 4 and cos.ndim == 4:
- # cos is (#b, h, t, r) -> reshape x to (b, h, t, dim_per_head)
- # The cos/sin batch dim may only be broadcastable, so take batch size from x
- b = x.shape[0]
- _, h, t, _ = cos.shape
+ # cos is (b, h, t, r) -> reshape x to (b, h, t, dim_per_head)
+ b, h, t, _ = cos.shape
x = x.reshape(b, t, h, -1).swapaxes(1, 2)
needs_reshape = True
@@ -141,10 +132,10 @@ def __init__(self, embedding_dim: int, num_mod_params: int = 6, use_additional_c
def forward(
self,
timestep: torch.Tensor,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- batch_size: Optional[int] = None,
- hidden_dtype: Optional[torch.dtype] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ batch_size: int | None = None,
+ hidden_dtype: torch.dtype | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
# No modulation happening here.
added_cond_kwargs = added_cond_kwargs or {"resolution": None, "aspect_ratio": None}
embedded_timestep = self.emb(timestep, **added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_dtype)
@@ -171,10 +162,10 @@ def __call__(
self,
attn: "LTX2Attention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- query_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- key_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
@@ -243,7 +234,7 @@ def __init__(
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = True,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
out_bias: bool = True,
qk_norm: str = "rms_norm_across_heads",
norm_eps: float = 1e-6,
@@ -282,10 +273,10 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- query_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- key_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -451,14 +442,14 @@ def forward(
temb_ca_audio_scale_shift: torch.Tensor,
temb_ca_gate: torch.Tensor,
temb_ca_audio_gate: torch.Tensor,
- video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- ca_video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- ca_audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- audio_encoder_attention_mask: Optional[torch.Tensor] = None,
- a2v_cross_attention_mask: Optional[torch.Tensor] = None,
- v2a_cross_attention_mask: Optional[torch.Tensor] = None,
+ video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ ca_video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ ca_audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ audio_encoder_attention_mask: torch.Tensor | None = None,
+ a2v_cross_attention_mask: torch.Tensor | None = None,
+ v2a_cross_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size = hidden_states.size(0)
@@ -621,7 +612,7 @@ def __init__(
base_width: int = 2048,
sampling_rate: int = 16000,
hop_length: int = 160,
- scale_factors: Tuple[int, ...] = (8, 32, 32),
+ scale_factors: tuple[int, ...] = (8, 32, 32),
theta: float = 10000.0,
causal_offset: int = 1,
modality: str = "video",
@@ -791,8 +782,8 @@ def prepare_coords(self, *args, **kwargs):
return self.prepare_audio_coords(*args, **kwargs)
def forward(
- self, coords: torch.Tensor, device: Optional[Union[str, torch.device]] = None
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ self, coords: torch.Tensor, device: str | torch.device | None = None
+ ) -> tuple[torch.Tensor, torch.Tensor]:
device = device or coords.device
# Number of spatiotemporal dimensions (3 for video, 1 (temporal) for audio and cross attn)
@@ -917,18 +908,18 @@ class LTX2VideoTransformer3DModel(
def __init__(
self,
in_channels: int = 128, # Video Arguments
- out_channels: Optional[int] = 128,
+ out_channels: int | None = 128,
patch_size: int = 1,
patch_size_t: int = 1,
num_attention_heads: int = 32,
attention_head_dim: int = 128,
cross_attention_dim: int = 4096,
- vae_scale_factors: Tuple[int, int, int] = (8, 32, 32),
+ vae_scale_factors: tuple[int, int, int] = (8, 32, 32),
pos_embed_max_pos: int = 20,
base_height: int = 2048,
base_width: int = 2048,
audio_in_channels: int = 128, # Audio Arguments
- audio_out_channels: Optional[int] = 128,
+ audio_out_channels: int | None = 128,
audio_patch_size: int = 1,
audio_patch_size_t: int = 1,
audio_num_attention_heads: int = 32,
@@ -1101,6 +1092,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -1108,17 +1100,17 @@ def forward(
encoder_hidden_states: torch.Tensor,
audio_encoder_hidden_states: torch.Tensor,
timestep: torch.LongTensor,
- audio_timestep: Optional[torch.LongTensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- audio_encoder_attention_mask: Optional[torch.Tensor] = None,
- num_frames: Optional[int] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ audio_timestep: torch.LongTensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ audio_encoder_attention_mask: torch.Tensor | None = None,
+ num_frames: int | None = None,
+ height: int | None = None,
+ width: int | None = None,
fps: float = 24.0,
- audio_num_frames: Optional[int] = None,
- video_coords: Optional[torch.Tensor] = None,
- audio_coords: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ audio_num_frames: int | None = None,
+ video_coords: torch.Tensor | None = None,
+ audio_coords: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
) -> torch.Tensor:
"""
@@ -1160,7 +1152,7 @@ def forward(
audio_coords (`torch.Tensor`, *optional*):
The audio coordinates to be used when calculating the rotary positional embeddings (RoPE) of shape
`(batch_size, 1, num_audio_tokens, 2)`. If not supplied, this will be calculated inside `forward`.
- attention_kwargs (`Dict[str, Any]`, *optional*):
+ attention_kwargs (`dict[str, Any]`, *optional*):
Optional dict of keyword args to be passed to the attention processor.
return_dict (`bool`, *optional*, defaults to `True`):
Whether to return a dict-like structured output of type `AudioVisualModelOutput` or a tuple.
@@ -1171,21 +1163,6 @@ def forward(
`tuple` is returned where the first element is the denoised video latent patch sequence and the second
element is the denoised audio latent patch sequence.
"""
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
# Determine timestep for audio.
audio_timestep = audio_timestep if audio_timestep is not None else timestep
@@ -1341,10 +1318,6 @@ def forward(
audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift
audio_output = self.audio_proj_out(audio_hidden_states)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output, audio_output)
return AudioVisualModelOutput(sample=output, audio_sample=audio_output)
diff --git a/src/diffusers/models/transformers/transformer_lumina2.py b/src/diffusers/models/transformers/transformer_lumina2.py
index 77121edb9fc9..03e2841f8bcb 100644
--- a/src/diffusers/models/transformers/transformer_lumina2.py
+++ b/src/diffusers/models/transformers/transformer_lumina2.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -22,7 +22,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
from ...loaders.single_file_model import FromOriginalModelMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ..attention import LuminaFeedForward
from ..attention_processor import Attention
from ..embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
@@ -58,7 +58,7 @@ def __init__(
def forward(
self, hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
timestep_proj = self.time_proj(timestep).type_as(hidden_states)
time_embed = self.timestep_embedder(timestep_proj)
caption_embed = self.caption_embedder(encoder_hidden_states)
@@ -80,9 +80,9 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
- base_sequence_length: Optional[int] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
+ base_sequence_length: int | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = hidden_states.shape
@@ -202,7 +202,7 @@ def forward(
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
image_rotary_emb: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
) -> torch.Tensor:
if self.modulation:
norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
@@ -231,7 +231,7 @@ def forward(
class Lumina2RotaryPosEmbed(nn.Module):
- def __init__(self, theta: int, axes_dim: List[int], axes_lens: List[int] = (300, 512, 512), patch_size: int = 2):
+ def __init__(self, theta: int, axes_dim: list[int], axes_lens: list[int] = (300, 512, 512), patch_size: int = 2):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -240,7 +240,7 @@ def __init__(self, theta: int, axes_dim: List[int], axes_lens: List[int] = (300,
self.freqs_cis = self._precompute_freqs_cis(axes_dim, axes_lens, theta)
- def _precompute_freqs_cis(self, axes_dim: List[int], axes_lens: List[int], theta: int) -> List[torch.Tensor]:
+ def _precompute_freqs_cis(self, axes_dim: list[int], axes_lens: list[int], theta: int) -> list[torch.Tensor]:
freqs_cis = []
freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
@@ -368,18 +368,18 @@ def __init__(
sample_size: int = 128,
patch_size: int = 2,
in_channels: int = 16,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
hidden_size: int = 2304,
num_layers: int = 26,
num_refiner_layers: int = 2,
num_attention_heads: int = 24,
num_kv_heads: int = 8,
multiple_of: int = 256,
- ffn_dim_multiplier: Optional[float] = None,
+ ffn_dim_multiplier: float | None = None,
norm_eps: float = 1e-5,
scaling_factor: float = 1.0,
- axes_dim_rope: Tuple[int, int, int] = (32, 32, 32),
- axes_lens: Tuple[int, int, int] = (300, 512, 512),
+ axes_dim_rope: tuple[int, int, int] = (32, 32, 32),
+ axes_lens: tuple[int, int, int] = (300, 512, 512),
cap_feat_dim: int = 1024,
) -> None:
super().__init__()
@@ -455,30 +455,16 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
encoder_attention_mask: torch.Tensor,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[torch.Tensor, Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> torch.Tensor | Transformer2DModelOutput:
# 1. Condition, positional & patch embedding
batch_size, _, height, width = hidden_states.shape
@@ -539,10 +525,6 @@ def forward(
)
output = torch.stack(output, dim=0)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
diff --git a/src/diffusers/models/transformers/transformer_mochi.py b/src/diffusers/models/transformers/transformer_mochi.py
index 63911fe7c10d..31106e3a0476 100644
--- a/src/diffusers/models/transformers/transformer_mochi.py
+++ b/src/diffusers/models/transformers/transformer_mochi.py
@@ -13,7 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, Optional, Tuple
+from typing import Any
import torch
import torch.nn as nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
from ...loaders.single_file_model import FromOriginalModelMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import FeedForward
from ..attention_processor import MochiAttention, MochiAttnProcessor2_0
@@ -104,7 +104,7 @@ def __init__(
def forward(
self, hidden_states: torch.Tensor, emb: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
hidden_states_dtype = hidden_states.dtype
emb = self.linear(self.silu(emb))
@@ -205,8 +205,8 @@ def forward(
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
encoder_attention_mask: torch.Tensor,
- image_rotary_emb: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
if not self.context_pre_only:
@@ -268,8 +268,8 @@ def _get_positions(
num_frames: int,
height: int,
width: int,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
) -> torch.Tensor:
scale = (self.target_area / (height * width)) ** 0.5
@@ -297,9 +297,9 @@ def forward(
num_frames: int,
height: int,
width: int,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
pos = self._get_positions(num_frames, height, width, device, dtype)
rope_cos, rope_sin = self._create_rope(pos_frequencies, pos)
return rope_cos, rope_sin
@@ -348,7 +348,7 @@ def __init__(
num_layers: int = 48,
pooled_projection_dim: int = 1536,
in_channels: int = 12,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
qk_norm: str = "rms_norm",
text_embed_dim: int = 4096,
time_embed_dim: int = 256,
@@ -404,30 +404,16 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_attention_mask: torch.Tensor,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
) -> torch.Tensor:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p = self.config.patch_size
@@ -479,10 +465,6 @@ def forward(
hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5)
output = hidden_states.reshape(batch_size, -1, num_frames, height, width)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
diff --git a/src/diffusers/models/transformers/transformer_omnigen.py b/src/diffusers/models/transformers/transformer_omnigen.py
index 6939cac0a3a7..bd8bb107e25c 100644
--- a/src/diffusers/models/transformers/transformer_omnigen.py
+++ b/src/diffusers/models/transformers/transformer_omnigen.py
@@ -13,7 +13,6 @@
# limitations under the License.
import math
-from typing import Dict, List, Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -200,8 +199,8 @@ def __call__(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = hidden_states.shape
@@ -308,7 +307,7 @@ class OmniGenTransformer2DModel(ModelMixin, ConfigMixin):
The size of the vocabulary of the embedding vocabulary.
rope_base (`int`, default to `10000`):
The default theta value to use when creating RoPE.
- rope_scaling (`Dict`, optional):
+ rope_scaling (`dict`, optional):
The scaling factors for the RoPE. Must contain `short_factor` and `long_factor`.
pos_embed_max_size (`int`, default to `192`):
The maximum size of the positional embeddings.
@@ -342,7 +341,7 @@ def __init__(
max_position_embeddings: int = 131072,
original_max_position_embeddings: int = 4096,
rope_base: int = 10000,
- rope_scaling: Dict = None,
+ rope_scaling: dict = None,
pos_embed_max_size: int = 192,
time_step_dim: int = 256,
flip_sin_to_cos: bool = True,
@@ -387,8 +386,8 @@ def __init__(
self.gradient_checkpointing = False
def _get_multimodal_embeddings(
- self, input_ids: torch.Tensor, input_img_latents: List[torch.Tensor], input_image_sizes: Dict
- ) -> Optional[torch.Tensor]:
+ self, input_ids: torch.Tensor, input_img_latents: list[torch.Tensor], input_image_sizes: dict
+ ) -> torch.Tensor | None:
if input_ids is None:
return None
@@ -408,14 +407,14 @@ def _get_multimodal_embeddings(
def forward(
self,
hidden_states: torch.Tensor,
- timestep: Union[int, float, torch.FloatTensor],
+ timestep: int | float | torch.FloatTensor,
input_ids: torch.Tensor,
- input_img_latents: List[torch.Tensor],
- input_image_sizes: Dict[int, List[int]],
+ input_img_latents: list[torch.Tensor],
+ input_image_sizes: dict[int, list[int]],
attention_mask: torch.Tensor,
position_ids: torch.Tensor,
return_dict: bool = True,
- ) -> Union[Transformer2DModelOutput, Tuple[torch.Tensor]]:
+ ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
batch_size, num_channels, height, width = hidden_states.shape
p = self.config.patch_size
post_patch_height, post_patch_width = height // p, width // p
diff --git a/src/diffusers/models/transformers/transformer_ovis_image.py b/src/diffusers/models/transformers/transformer_ovis_image.py
index 139ceaefa4e9..8280a4871f10 100644
--- a/src/diffusers/models/transformers/transformer_ovis_image.py
+++ b/src/diffusers/models/transformers/transformer_ovis_image.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -78,8 +78,8 @@ def __call__(
attn: "OvisImageAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
attn, hidden_states, encoder_hidden_states
@@ -145,12 +145,12 @@ def __init__(
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
- added_kv_proj_dim: Optional[int] = None,
- added_proj_bias: Optional[bool] = True,
+ added_kv_proj_dim: int | None = None,
+ added_proj_bias: bool | None = True,
out_bias: bool = True,
eps: float = 1e-5,
out_dim: int = None,
- context_pre_only: Optional[bool] = None,
+ context_pre_only: bool | None = None,
pre_only: bool = False,
elementwise_affine: bool = True,
processor=None,
@@ -195,9 +195,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
@@ -238,9 +238,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
text_seq_len = encoder_hidden_states.shape[1]
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
@@ -301,9 +301,9 @@ def forward(
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -354,7 +354,7 @@ def forward(
class OvisImagePosEmbed(nn.Module):
- def __init__(self, theta: int, axes_dim: List[int]):
+ def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -413,7 +413,7 @@ class OvisImageTransformer2DModel(
joint_attention_dim (`int`, defaults to `2048`):
The number of dimensions to use for the joint attention (embedding/channel dimension of
`encoder_hidden_states`).
- axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`):
The dimensions to use for the rotary positional embeddings.
"""
@@ -427,13 +427,13 @@ def __init__(
self,
patch_size: int = 1,
in_channels: int = 64,
- out_channels: Optional[int] = 64,
+ out_channels: int | None = 64,
num_layers: int = 6,
num_single_layers: int = 27,
attention_head_dim: int = 128,
num_attention_heads: int = 24,
joint_attention_dim: int = 2048,
- axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+ axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
):
super().__init__()
self.out_channels = out_channels or in_channels
@@ -483,7 +483,7 @@ def forward(
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
return_dict: bool = True,
- ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ ) -> torch.Tensor | Transformer2DModelOutput:
"""
The [`OvisImageTransformer2DModel`] forward method.
diff --git a/src/diffusers/models/transformers/transformer_prx.py b/src/diffusers/models/transformers/transformer_prx.py
index a87c120fdcd7..3c8e8ae4e2c9 100644
--- a/src/diffusers/models/transformers/transformer_prx.py
+++ b/src/diffusers/models/transformers/transformer_prx.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
from torch import nn
@@ -96,9 +96,9 @@ def __call__(
self,
attn: "PRXAttention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -234,9 +234,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
return self.processor(
@@ -263,10 +263,10 @@ class PRXEmbedND(nn.Module):
theta (int):
Scaling factor that controls the frequency spectrum of the rotary embeddings.
axes_dim (list[int]):
- List of embedding dimensions for each axis (each must be even).
+ list of embedding dimensions for each axis (each must be even).
"""
- def __init__(self, dim: int, theta: int, axes_dim: List[int]):
+ def __init__(self, dim: int, theta: int, axes_dim: list[int]):
super().__init__()
self.dim = dim
self.theta = theta
@@ -346,7 +346,7 @@ def __init__(self, dim: int):
def forward(
self, vec: torch.Tensor
- ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+ ) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(6, dim=-1)
return tuple(out[:3]), tuple(out[3:])
@@ -389,7 +389,7 @@ def __init__(
hidden_size: int,
num_heads: int,
mlp_ratio: float = 4.0,
- qk_scale: Optional[float] = None,
+ qk_scale: float | None = None,
):
super().__init__()
@@ -430,8 +430,8 @@ def forward(
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
image_rotary_emb: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- **kwargs: Dict[str, Any],
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: dict[str, Any],
) -> torch.Tensor:
r"""
Runs modulation-gated cross-attention and MLP, with residual connections.
@@ -607,7 +607,7 @@ class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
depth (`int`, *optional*, defaults to 16):
Number of transformer blocks.
axes_dim (`list[int]`, *optional*):
- List of dimensions for each positional embedding axis. Defaults to `[32, 32]`.
+ list of dimensions for each positional embedding axis. Defaults to `[32, 32]`.
theta (`int`, *optional*, defaults to 10000):
Frequency scaling factor for rotary embeddings.
time_factor (`float`, *optional*, defaults to 1000.0):
@@ -728,10 +728,10 @@ def forward(
hidden_states: torch.Tensor,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_mask: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]:
+ ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput:
r"""
Forward pass of the PRXTransformer2DModel.
diff --git a/src/diffusers/models/transformers/transformer_qwenimage.py b/src/diffusers/models/transformers/transformer_qwenimage.py
index 74461a07667e..a54cb3b8e092 100644
--- a/src/diffusers/models/transformers/transformer_qwenimage.py
+++ b/src/diffusers/models/transformers/transformer_qwenimage.py
@@ -15,7 +15,7 @@
import functools
import math
from math import prod
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -24,7 +24,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, deprecate, logging
from ...utils.torch_utils import maybe_allow_in_graph
from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
from ..attention import AttentionMixin, FeedForward
@@ -96,10 +96,10 @@ def get_timestep_embedding(
def apply_rotary_emb_qwen(
x: torch.Tensor,
- freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+ freqs_cis: torch.Tensor | tuple[torch.Tensor],
use_real: bool = True,
use_real_unbind_dim: int = -1,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
"""
Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
@@ -109,10 +109,10 @@ def apply_rotary_emb_qwen(
Args:
x (`torch.Tensor`):
Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply
- freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+ freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
Returns:
- Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+ tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings.
"""
if use_real:
cos, sin = freqs_cis # [S, D]
@@ -143,8 +143,8 @@ def apply_rotary_emb_qwen(
def compute_text_seq_len_from_mask(
- encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: Optional[torch.Tensor]
-) -> Tuple[int, Optional[torch.Tensor], Optional[torch.Tensor]]:
+ encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: torch.Tensor | None
+) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:
"""
Compute text sequence length without assuming contiguous masks. Returns length for RoPE and a normalized bool mask.
"""
@@ -164,7 +164,11 @@ def compute_text_seq_len_from_mask(
position_ids = torch.arange(text_seq_len, device=encoder_hidden_states.device, dtype=torch.long)
active_positions = torch.where(encoder_hidden_states_mask, position_ids, position_ids.new_zeros(()))
has_active = encoder_hidden_states_mask.any(dim=1)
- per_sample_len = torch.where(has_active, active_positions.max(dim=1).values + 1, torch.as_tensor(text_seq_len))
+ per_sample_len = torch.where(
+ has_active,
+ active_positions.max(dim=1).values + 1,
+ torch.as_tensor(text_seq_len, device=encoder_hidden_states.device),
+ )
return text_seq_len, per_sample_len, encoder_hidden_states_mask
@@ -194,7 +198,7 @@ def forward(self, timestep, hidden_states, addition_t_cond=None):
class QwenEmbedRope(nn.Module):
- def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
+ def __init__(self, theta: int, axes_dim: list[int], scale_rope=False):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -232,16 +236,16 @@ def rope_params(self, index, dim, theta=10000):
def forward(
self,
- video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]],
- txt_seq_lens: Optional[List[int]] = None,
+ video_fhw: tuple[int, int, int, list[tuple[int, int, int]]],
+ txt_seq_lens: list[int] | None = None,
device: torch.device = None,
- max_txt_seq_len: Optional[Union[int, torch.Tensor]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ max_txt_seq_len: int | torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
"""
Args:
- video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`):
+ video_fhw (`tuple[int, int, int]` or `list[tuple[int, int, int]]`):
A list of 3 integers [frame, height, width] representing the shape of the video.
- txt_seq_lens (`List[int]`, *optional*, **Deprecated**):
+ txt_seq_lens (`list[int]`, *optional*, **Deprecated**):
Deprecated parameter. Use `max_txt_seq_len` instead. If provided, the maximum value will be used.
device: (`torch.device`, *optional*):
The device on which to perform the RoPE computation.
@@ -329,7 +333,7 @@ def _compute_video_freqs(
class QwenEmbedLayer3DRope(nn.Module):
- def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
+ def __init__(self, theta: int, axes_dim: list[int], scale_rope=False):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
@@ -366,13 +370,13 @@ def rope_params(self, index, dim, theta=10000):
def forward(
self,
- video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]],
- max_txt_seq_len: Union[int, torch.Tensor],
+ video_fhw: tuple[int, int, int, list[tuple[int, int, int]]],
+ max_txt_seq_len: int | torch.Tensor,
device: torch.device = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
"""
Args:
- video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`):
+ video_fhw (`tuple[int, int, int]` or `list[tuple[int, int, int]]`):
A list of 3 integers [frame, height, width] representing the shape of the video, or a list of layer
structures.
max_txt_seq_len (`int` or `torch.Tensor`):
@@ -490,8 +494,8 @@ def __call__(
hidden_states: torch.FloatTensor, # Image stream
encoder_hidden_states: torch.FloatTensor = None, # Text stream
encoder_hidden_states_mask: torch.FloatTensor = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- image_rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.FloatTensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
) -> torch.FloatTensor:
if encoder_hidden_states is None:
raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")
@@ -663,10 +667,10 @@ def forward(
encoder_hidden_states: torch.Tensor,
encoder_hidden_states_mask: torch.Tensor,
temb: torch.Tensor,
- image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- modulate_index: Optional[List[int]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ modulate_index: list[int] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# Get modulation parameters for both streams
img_mod_params = self.img_mod(temb) # [B, 6*dim]
@@ -753,7 +757,7 @@ class QwenImageTransformer2DModel(
`encoder_hidden_states`).
guidance_embeds (`bool`, defaults to `False`):
Whether to use guidance embeddings for guidance-distilled variant of the model.
- axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+ axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`):
The dimensions to use for the rotary positional embeddings.
"""
@@ -782,13 +786,13 @@ def __init__(
self,
patch_size: int = 2,
in_channels: int = 64,
- out_channels: Optional[int] = 16,
+ out_channels: int | None = 16,
num_layers: int = 60,
attention_head_dim: int = 128,
num_attention_heads: int = 24,
joint_attention_dim: int = 3584,
guidance_embeds: bool = False, # TODO: this should probably be removed
- axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+ axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
zero_cond_t: bool = False,
use_additional_t_cond: bool = False,
use_layer3d_rope: bool = False,
@@ -829,20 +833,21 @@ def __init__(
self.gradient_checkpointing = False
self.zero_cond_t = zero_cond_t
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
encoder_hidden_states_mask: torch.Tensor = None,
timestep: torch.LongTensor = None,
- img_shapes: Optional[List[Tuple[int, int, int]]] = None,
- txt_seq_lens: Optional[List[int]] = None,
+ img_shapes: list[tuple[int, int, int]] | None = None,
+ txt_seq_lens: list[int] | None = None,
guidance: torch.Tensor = None, # TODO: this should probably be removed
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
controlnet_block_samples=None,
additional_t_cond=None,
return_dict: bool = True,
- ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ ) -> torch.Tensor | Transformer2DModelOutput:
"""
The [`QwenTransformer2DModel`] forward method.
@@ -857,9 +862,9 @@ def forward(
(not just contiguous valid tokens followed by padding) since it's applied element-wise in attention.
timestep ( `torch.LongTensor`):
Used to indicate denoising step.
- img_shapes (`List[Tuple[int, int, int]]`, *optional*):
+ img_shapes (`list[tuple[int, int, int]]`, *optional*):
Image shapes for RoPE computation.
- txt_seq_lens (`List[int]`, *optional*, **Deprecated**):
+ txt_seq_lens (`list[int]`, *optional*, **Deprecated**):
Deprecated parameter. Use `encoder_hidden_states_mask` instead. If provided, the maximum value will be
used to compute RoPE sequence length.
guidance (`torch.Tensor`, *optional*):
@@ -887,20 +892,6 @@ def forward(
"The mask-based approach is more flexible and supports variable-length sequences.",
standard_warn=False,
)
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
hidden_states = self.img_in(hidden_states)
@@ -981,10 +972,6 @@ def forward(
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_sana_video.py b/src/diffusers/models/transformers/transformer_sana_video.py
index a4f90342631a..f833c0e842c3 100644
--- a/src/diffusers/models/transformers/transformer_sana_video.py
+++ b/src/diffusers/models/transformers/transformer_sana_video.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn.functional as F
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ..attention import AttentionMixin
from ..attention_dispatch import dispatch_attention_fn
from ..attention_processor import Attention
@@ -40,7 +40,7 @@ def __init__(
in_channels: int,
out_channels: int,
expand_ratio: float = 4,
- norm_type: Optional[str] = None,
+ norm_type: str | None = None,
residual_connection: bool = True,
) -> None:
super().__init__()
@@ -103,9 +103,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
original_dtype = hidden_states.dtype
@@ -176,7 +176,7 @@ class WanRotaryPosEmbed(nn.Module):
def __init__(
self,
attention_head_dim: int,
- patch_size: Tuple[int, int, int],
+ patch_size: tuple[int, int, int],
max_seq_len: int,
theta: float = 10000.0,
):
@@ -290,8 +290,8 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
@@ -358,15 +358,15 @@ def __init__(
num_attention_heads: int = 20,
attention_head_dim: int = 112,
dropout: float = 0.0,
- num_cross_attention_heads: Optional[int] = 20,
- cross_attention_head_dim: Optional[int] = 112,
- cross_attention_dim: Optional[int] = 2240,
+ num_cross_attention_heads: int | None = 20,
+ cross_attention_head_dim: int | None = 112,
+ cross_attention_dim: int | None = 2240,
attention_bias: bool = True,
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-6,
attention_out_bias: bool = True,
mlp_ratio: float = 3.0,
- qk_norm: Optional[str] = "rms_norm_across_heads",
+ qk_norm: str | None = "rms_norm_across_heads",
rope_max_seq_len: int = 1024,
) -> None:
super().__init__()
@@ -409,14 +409,14 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- timestep: Optional[torch.LongTensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ timestep: torch.LongTensor | None = None,
frames: int = None,
height: int = None,
width: int = None,
- rotary_emb: Optional[torch.Tensor] = None,
+ rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size = hidden_states.shape[0]
@@ -503,25 +503,25 @@ class SanaVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Fro
def __init__(
self,
in_channels: int = 16,
- out_channels: Optional[int] = 16,
+ out_channels: int | None = 16,
num_attention_heads: int = 20,
attention_head_dim: int = 112,
num_layers: int = 20,
- num_cross_attention_heads: Optional[int] = 20,
- cross_attention_head_dim: Optional[int] = 112,
- cross_attention_dim: Optional[int] = 2240,
+ num_cross_attention_heads: int | None = 20,
+ cross_attention_head_dim: int | None = 112,
+ cross_attention_dim: int | None = 2240,
caption_channels: int = 2304,
mlp_ratio: float = 2.5,
dropout: float = 0.0,
attention_bias: bool = False,
sample_size: int = 30,
- patch_size: Tuple[int, int, int] = (1, 2, 2),
+ patch_size: tuple[int, int, int] = (1, 2, 2),
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-6,
- interpolation_scale: Optional[int] = None,
+ interpolation_scale: int | None = None,
guidance_embeds: bool = False,
guidance_embeds_scale: float = 0.1,
- qk_norm: Optional[str] = "rms_norm_across_heads",
+ qk_norm: str | None = "rms_norm_across_heads",
rope_max_seq_len: int = 1024,
) -> None:
super().__init__()
@@ -570,33 +570,19 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep: torch.Tensor,
- guidance: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_block_samples: Optional[Tuple[torch.Tensor]] = None,
+ guidance: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ controlnet_block_samples: tuple[torch.Tensor] | None = None,
return_dict: bool = True,
- ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput:
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
@@ -695,10 +681,6 @@ def forward(
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_sd3.py b/src/diffusers/models/transformers/transformer_sd3.py
index 05391e047b7a..ead657d0cfd2 100644
--- a/src/diffusers/models/transformers/transformer_sd3.py
+++ b/src/diffusers/models/transformers/transformer_sd3.py
@@ -11,14 +11,14 @@
# 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.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin, SD3Transformer2DLoadersMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import AttentionMixin, FeedForward, JointTransformerBlock
from ..attention_processor import (
@@ -106,7 +106,7 @@ class SD3Transformer2DModel(
The number of latent channels in the output.
pos_embed_max_size (`int`, defaults to `96`):
The maximum latent height/width of positional embeddings.
- dual_attention_layers (`Tuple[int, ...]`, defaults to `()`):
+ dual_attention_layers (`tuple[int, ...]`, defaults to `()`):
The number of dual-stream transformer blocks to use.
qk_norm (`str`, *optional*, defaults to `None`):
The normalization to use for query and key in the attention layer. If `None`, no normalization is used.
@@ -130,10 +130,10 @@ def __init__(
pooled_projection_dim: int = 2048,
out_channels: int = 16,
pos_embed_max_size: int = 96,
- dual_attention_layers: Tuple[
+ dual_attention_layers: tuple[
int, ...
] = (), # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5
- qk_norm: Optional[str] = None,
+ qk_norm: str | None = None,
):
super().__init__()
self.out_channels = out_channels if out_channels is not None else in_channels
@@ -172,7 +172,7 @@ def __init__(
self.gradient_checkpointing = False
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
- def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+ def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -245,17 +245,18 @@ def unfuse_qkv_projections(self):
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
+ @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
pooled_projections: torch.Tensor = None,
timestep: torch.LongTensor = None,
- block_controlnet_hidden_states: List = None,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ block_controlnet_hidden_states: list = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- skip_layers: Optional[List[int]] = None,
- ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ skip_layers: list[int] | None = None,
+ ) -> torch.Tensor | Transformer2DModelOutput:
"""
The [`SD3Transformer2DModel`] forward method.
@@ -284,20 +285,6 @@ def forward(
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
- if joint_attention_kwargs is not None:
- joint_attention_kwargs = joint_attention_kwargs.copy()
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
- )
height, width = hidden_states.shape[-2:]
@@ -352,10 +339,6 @@ def forward(
shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_skyreels_v2.py b/src/diffusers/models/transformers/transformer_skyreels_v2.py
index 2b9fc5b8d9fb..9067e32ea5c3 100644
--- a/src/diffusers/models/transformers/transformer_skyreels_v2.py
+++ b/src/diffusers/models/transformers/transformer_skyreels_v2.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, deprecate, logging
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
from ..attention_dispatch import dispatch_attention_fn
@@ -85,9 +85,9 @@ def __call__(
self,
attn: "SkyReelsV2Attention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
encoder_hidden_states_img = None
if attn.add_k_proj is not None:
@@ -188,8 +188,8 @@ def __init__(
dim_head: int = 64,
eps: float = 1e-5,
dropout: float = 0.0,
- added_kv_proj_dim: Optional[int] = None,
- cross_attention_dim_head: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
+ cross_attention_dim_head: int | None = None,
processor=None,
is_cross_attention=None,
):
@@ -275,9 +275,9 @@ def unfuse_projections(self):
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs)
@@ -335,8 +335,8 @@ def __init__(
time_freq_dim: int,
time_proj_dim: int,
text_embed_dim: int,
- image_embed_dim: Optional[int] = None,
- pos_embed_seq_len: Optional[int] = None,
+ image_embed_dim: int | None = None,
+ pos_embed_seq_len: int | None = None,
):
super().__init__()
@@ -354,7 +354,7 @@ def forward(
self,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
):
timestep = self.timesteps_proj(timestep)
@@ -375,7 +375,7 @@ class SkyReelsV2RotaryPosEmbed(nn.Module):
def __init__(
self,
attention_head_dim: int,
- patch_size: Tuple[int, int, int],
+ patch_size: tuple[int, int, int],
max_seq_len: int,
theta: float = 10000.0,
):
@@ -445,7 +445,7 @@ def __init__(
qk_norm: str = "rms_norm_across_heads",
cross_attn_norm: bool = False,
eps: float = 1e-6,
- added_kv_proj_dim: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
):
super().__init__()
@@ -522,7 +522,7 @@ class SkyReelsV2Transformer3DModel(
A Transformer model for video-like data used in the Wan-based SkyReels-V2 model.
Args:
- patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`):
+ patch_size (`tuple[int]`, defaults to `(1, 2, 2)`):
3D patch dimensions for video embedding (t_patch, h_patch, w_patch).
num_attention_heads (`int`, defaults to `16`):
Fixed length for text embeddings.
@@ -540,7 +540,7 @@ class SkyReelsV2Transformer3DModel(
Intermediate dimension in feed-forward network.
num_layers (`int`, defaults to `32`):
The number of layers of transformer blocks to use.
- window_size (`Tuple[int]`, defaults to `(-1, -1)`):
+ window_size (`tuple[int]`, defaults to `(-1, -1)`):
Window size for local attention (-1 indicates global attention).
cross_attn_norm (`bool`, defaults to `True`):
Enable cross-attention normalization.
@@ -570,7 +570,7 @@ class SkyReelsV2Transformer3DModel(
@register_to_config
def __init__(
self,
- patch_size: Tuple[int, ...] = (1, 2, 2),
+ patch_size: tuple[int] = (1, 2, 2),
num_attention_heads: int = 16,
attention_head_dim: int = 128,
in_channels: int = 16,
@@ -580,12 +580,12 @@ def __init__(
ffn_dim: int = 8192,
num_layers: int = 32,
cross_attn_norm: bool = True,
- qk_norm: Optional[str] = "rms_norm_across_heads",
+ qk_norm: str | None = "rms_norm_across_heads",
eps: float = 1e-6,
- image_dim: Optional[int] = None,
- added_kv_proj_dim: Optional[int] = None,
+ image_dim: int | None = None,
+ added_kv_proj_dim: int | None = None,
rope_max_seq_len: int = 1024,
- pos_embed_seq_len: Optional[int] = None,
+ pos_embed_seq_len: int | None = None,
inject_sample_info: bool = False,
num_frame_per_block: int = 1,
) -> None:
@@ -630,32 +630,18 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
enable_diffusion_forcing: bool = False,
- fps: Optional[torch.Tensor] = None,
+ fps: torch.Tensor | None = None,
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> torch.Tensor | dict[str, torch.Tensor]:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.config.patch_size
post_patch_num_frames = num_frames // p_t
@@ -771,10 +757,6 @@ def forward(
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_temporal.py b/src/diffusers/models/transformers/transformer_temporal.py
index ffaf31d04570..b6fedcb26cc8 100644
--- a/src/diffusers/models/transformers/transformer_temporal.py
+++ b/src/diffusers/models/transformers/transformer_temporal.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, Optional
+from typing import Any
import torch
from torch import nn
@@ -74,19 +74,19 @@ def __init__(
self,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
- out_channels: Optional[int] = None,
+ in_channels: int | None = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
norm_num_groups: int = 32,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
- sample_size: Optional[int] = None,
+ sample_size: int | None = None,
activation_fn: str = "geglu",
norm_elementwise_affine: bool = True,
double_self_attention: bool = True,
- positional_embeddings: Optional[str] = None,
- num_positional_embeddings: Optional[int] = None,
+ positional_embeddings: str | None = None,
+ num_positional_embeddings: int | None = None,
):
super().__init__()
self.num_attention_heads = num_attention_heads
@@ -123,11 +123,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.LongTensor] = None,
- timestep: Optional[torch.LongTensor] = None,
+ encoder_hidden_states: torch.LongTensor | None = None,
+ timestep: torch.LongTensor | None = None,
class_labels: torch.LongTensor = None,
num_frames: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
) -> TransformerTemporalModelOutput:
"""
@@ -222,9 +222,9 @@ def __init__(
num_attention_heads: int = 16,
attention_head_dim: int = 88,
in_channels: int = 320,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 1,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
):
super().__init__()
self.num_attention_heads = num_attention_heads
@@ -280,8 +280,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- image_only_indicator: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ image_only_indicator: torch.Tensor | None = None,
return_dict: bool = True,
):
"""
diff --git a/src/diffusers/models/transformers/transformer_wan.py b/src/diffusers/models/transformers/transformer_wan.py
index 132f615f2199..5926bbb8e713 100644
--- a/src/diffusers/models/transformers/transformer_wan.py
+++ b/src/diffusers/models/transformers/transformer_wan.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, deprecate, logging
from ...utils.torch_utils import maybe_allow_in_graph
from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
@@ -42,7 +42,7 @@ def _get_qkv_projections(attn: "WanAttention", hidden_states: torch.Tensor, enco
encoder_hidden_states = hidden_states
if attn.fused_projections:
- if attn.cross_attention_dim_head is None:
+ if not attn.is_cross_attention:
# In self-attention layers, we can fuse the entire QKV projection into a single linear
query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
else:
@@ -79,9 +79,9 @@ def __call__(
self,
attn: "WanAttention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
encoder_hidden_states_img = None
if attn.add_k_proj is not None:
@@ -183,8 +183,8 @@ def __init__(
dim_head: int = 64,
eps: float = 1e-5,
dropout: float = 0.0,
- added_kv_proj_dim: Optional[int] = None,
- cross_attention_dim_head: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
+ cross_attention_dim_head: int | None = None,
processor=None,
is_cross_attention=None,
):
@@ -214,7 +214,10 @@ def __init__(
self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True)
self.norm_added_k = torch.nn.RMSNorm(dim_head * heads, eps=eps)
- self.is_cross_attention = cross_attention_dim_head is not None
+ if is_cross_attention is not None:
+ self.is_cross_attention = is_cross_attention
+ else:
+ self.is_cross_attention = cross_attention_dim_head is not None
self.set_processor(processor)
@@ -222,7 +225,7 @@ def fuse_projections(self):
if getattr(self, "fused_projections", False):
return
- if self.cross_attention_dim_head is None:
+ if not self.is_cross_attention:
concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])
concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])
out_features, in_features = concatenated_weights.shape
@@ -270,9 +273,9 @@ def unfuse_projections(self):
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs)
@@ -309,8 +312,8 @@ def __init__(
time_freq_dim: int,
time_proj_dim: int,
text_embed_dim: int,
- image_embed_dim: Optional[int] = None,
- pos_embed_seq_len: Optional[int] = None,
+ image_embed_dim: int | None = None,
+ pos_embed_seq_len: int | None = None,
):
super().__init__()
@@ -328,8 +331,8 @@ def forward(
self,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
- timestep_seq_len: Optional[int] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
+ timestep_seq_len: int | None = None,
):
timestep = self.timesteps_proj(timestep)
if timestep_seq_len is not None:
@@ -352,7 +355,7 @@ class WanRotaryPosEmbed(nn.Module):
def __init__(
self,
attention_head_dim: int,
- patch_size: Tuple[int, int, int],
+ patch_size: tuple[int, int, int],
max_seq_len: int,
theta: float = 10000.0,
):
@@ -423,7 +426,7 @@ def __init__(
qk_norm: str = "rms_norm_across_heads",
cross_attn_norm: bool = False,
eps: float = 1e-6,
- added_kv_proj_dim: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
):
super().__init__()
@@ -508,7 +511,7 @@ class WanTransformer3DModel(
A Transformer model for video-like data used in the Wan model.
Args:
- patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`):
+ patch_size (`tuple[int]`, defaults to `(1, 2, 2)`):
3D patch dimensions for video embedding (t_patch, h_patch, w_patch).
num_attention_heads (`int`, defaults to `40`):
Fixed length for text embeddings.
@@ -526,7 +529,7 @@ class WanTransformer3DModel(
Intermediate dimension in feed-forward network.
num_layers (`int`, defaults to `40`):
The number of layers of transformer blocks to use.
- window_size (`Tuple[int]`, defaults to `(-1, -1)`):
+ window_size (`tuple[int]`, defaults to `(-1, -1)`):
Window size for local attention (-1 indicates global attention).
cross_attn_norm (`bool`, defaults to `True`):
Enable cross-attention normalization.
@@ -568,7 +571,7 @@ class WanTransformer3DModel(
@register_to_config
def __init__(
self,
- patch_size: Tuple[int, ...] = (1, 2, 2),
+ patch_size: tuple[int, ...] = (1, 2, 2),
num_attention_heads: int = 40,
attention_head_dim: int = 128,
in_channels: int = 16,
@@ -578,12 +581,12 @@ def __init__(
ffn_dim: int = 13824,
num_layers: int = 40,
cross_attn_norm: bool = True,
- qk_norm: Optional[str] = "rms_norm_across_heads",
+ qk_norm: str | None = "rms_norm_across_heads",
eps: float = 1e-6,
- image_dim: Optional[int] = None,
- added_kv_proj_dim: Optional[int] = None,
+ image_dim: int | None = None,
+ added_kv_proj_dim: int | None = None,
rope_max_seq_len: int = 1024,
- pos_embed_seq_len: Optional[int] = None,
+ pos_embed_seq_len: int | None = None,
) -> None:
super().__init__()
@@ -622,30 +625,16 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> torch.Tensor | dict[str, torch.Tensor]:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.config.patch_size
post_patch_num_frames = num_frames // p_t
@@ -713,10 +702,6 @@ def forward(
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_wan_animate.py b/src/diffusers/models/transformers/transformer_wan_animate.py
index ae1c3095a17f..b5b15fe06099 100644
--- a/src/diffusers/models/transformers/transformer_wan_animate.py
+++ b/src/diffusers/models/transformers/transformer_wan_animate.py
@@ -13,7 +13,7 @@
# limitations under the License.
import math
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
from ..attention_dispatch import dispatch_attention_fn
from ..cache_utils import CacheMixin
@@ -54,7 +54,7 @@ def _get_qkv_projections(attn: "WanAttention", hidden_states: torch.Tensor, enco
encoder_hidden_states = hidden_states
if attn.fused_projections:
- if attn.cross_attention_dim_head is None:
+ if not attn.is_cross_attention:
# In self-attention layers, we can fuse the entire QKV projection into a single linear
query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
else:
@@ -83,7 +83,7 @@ class FusedLeakyReLU(nn.Module):
Fused LeakyRelu with scale factor and channel-wise bias.
"""
- def __init__(self, negative_slope: float = 0.2, scale: float = 2**0.5, bias_channels: Optional[int] = None):
+ def __init__(self, negative_slope: float = 0.2, scale: float = 2**0.5, bias_channels: int | None = None):
super().__init__()
self.negative_slope = negative_slope
self.scale = scale
@@ -117,7 +117,7 @@ def __init__(
stride: int = 1,
padding: int = 0,
bias: bool = True,
- blur_kernel: Optional[Tuple[int, ...]] = None,
+ blur_kernel: tuple[int, ...] | None = None,
blur_upsample_factor: int = 1,
use_activation: bool = True,
):
@@ -231,7 +231,7 @@ def __init__(
out_channels: int,
kernel_size: int = 3,
kernel_size_skip: int = 1,
- blur_kernel: Tuple[int, ...] = (1, 3, 3, 1),
+ blur_kernel: tuple[int, ...] = (1, 3, 3, 1),
downsample_factor: int = 2,
):
super().__init__()
@@ -288,7 +288,7 @@ def __init__(
motion_dim: int = 20,
out_dim: int = 512,
motion_blocks: int = 5,
- channels: Optional[Dict[str, int]] = None,
+ channels: dict[str, int] | None = None,
):
super().__init__()
self.size = size
@@ -435,8 +435,8 @@ def __call__(
self,
attn: "WanAnimateFaceBlockCrossAttention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
# encoder_hidden_states corresponds to the motion vec
# attention_mask corresponds to the motion mask (if any)
@@ -501,14 +501,17 @@ def __init__(
heads: int = 8,
dim_head: int = 64,
eps: float = 1e-6,
- cross_attention_dim_head: Optional[int] = None,
+ cross_attention_dim_head: int | None = None,
+ bias: bool = True,
processor=None,
):
super().__init__()
self.inner_dim = dim_head * heads
self.heads = heads
- self.cross_attention_head_dim = cross_attention_dim_head
+ self.cross_attention_dim_head = cross_attention_dim_head
self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads
+ self.use_bias = bias
+ self.is_cross_attention = cross_attention_dim_head is not None
# 1. Pre-Attention Norms for the hidden_states (video latents) and encoder_hidden_states (motion vector).
# NOTE: this is not used in "vanilla" WanAttention
@@ -516,10 +519,10 @@ def __init__(
self.pre_norm_kv = nn.LayerNorm(dim, eps, elementwise_affine=False)
# 2. QKV and Output Projections
- self.to_q = torch.nn.Linear(dim, self.inner_dim, bias=True)
- self.to_k = torch.nn.Linear(dim, self.kv_inner_dim, bias=True)
- self.to_v = torch.nn.Linear(dim, self.kv_inner_dim, bias=True)
- self.to_out = torch.nn.Linear(self.inner_dim, dim, bias=True)
+ self.to_q = torch.nn.Linear(dim, self.inner_dim, bias=bias)
+ self.to_k = torch.nn.Linear(dim, self.kv_inner_dim, bias=bias)
+ self.to_v = torch.nn.Linear(dim, self.kv_inner_dim, bias=bias)
+ self.to_out = torch.nn.Linear(self.inner_dim, dim, bias=bias)
# 3. QK Norm
# NOTE: this is applied after the reshape, so only over dim_head rather than dim_head * heads
@@ -534,8 +537,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
return self.processor(self, hidden_states, encoder_hidden_states, attention_mask)
@@ -556,9 +559,9 @@ def __call__(
self,
attn: "WanAttention",
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
encoder_hidden_states_img = None
if attn.add_k_proj is not None:
@@ -651,8 +654,8 @@ def __init__(
dim_head: int = 64,
eps: float = 1e-5,
dropout: float = 0.0,
- added_kv_proj_dim: Optional[int] = None,
- cross_attention_dim_head: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
+ cross_attention_dim_head: int | None = None,
processor=None,
is_cross_attention=None,
):
@@ -682,7 +685,10 @@ def __init__(
self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True)
self.norm_added_k = torch.nn.RMSNorm(dim_head * heads, eps=eps)
- self.is_cross_attention = cross_attention_dim_head is not None
+ if is_cross_attention is not None:
+ self.is_cross_attention = is_cross_attention
+ else:
+ self.is_cross_attention = cross_attention_dim_head is not None
self.set_processor(processor)
@@ -690,7 +696,7 @@ def fuse_projections(self):
if getattr(self, "fused_projections", False):
return
- if self.cross_attention_dim_head is None:
+ if not self.is_cross_attention:
concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])
concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])
out_features, in_features = concatenated_weights.shape
@@ -738,9 +744,9 @@ def unfuse_projections(self):
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs)
@@ -779,8 +785,8 @@ def __init__(
time_freq_dim: int,
time_proj_dim: int,
text_embed_dim: int,
- image_embed_dim: Optional[int] = None,
- pos_embed_seq_len: Optional[int] = None,
+ image_embed_dim: int | None = None,
+ pos_embed_seq_len: int | None = None,
):
super().__init__()
@@ -798,8 +804,8 @@ def forward(
self,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
- timestep_seq_len: Optional[int] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
+ timestep_seq_len: int | None = None,
):
timestep = self.timesteps_proj(timestep)
if timestep_seq_len is not None:
@@ -825,7 +831,7 @@ class WanRotaryPosEmbed(nn.Module):
def __init__(
self,
attention_head_dim: int,
- patch_size: Tuple[int, int, int],
+ patch_size: tuple[int, int, int],
max_seq_len: int,
theta: float = 10000.0,
):
@@ -896,7 +902,7 @@ def __init__(
qk_norm: str = "rms_norm_across_heads",
cross_attn_norm: bool = False,
eps: float = 1e-6,
- added_kv_proj_dim: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
):
super().__init__()
@@ -981,7 +987,7 @@ class WanAnimateTransformer3DModel(
A Transformer model for video-like data used in the WanAnimate model.
Args:
- patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`):
+ patch_size (`tuple[int]`, defaults to `(1, 2, 2)`):
3D patch dimensions for video embedding (t_patch, h_patch, w_patch).
num_attention_heads (`int`, defaults to `40`):
Fixed length for text embeddings.
@@ -999,7 +1005,7 @@ class WanAnimateTransformer3DModel(
Intermediate dimension in feed-forward network.
num_layers (`int`, defaults to `40`):
The number of layers of transformer blocks to use.
- window_size (`Tuple[int]`, defaults to `(-1, -1)`):
+ window_size (`tuple[int]`, defaults to `(-1, -1)`):
Window size for local attention (-1 indicates global attention).
cross_attn_norm (`bool`, defaults to `True`):
Enable cross-attention normalization.
@@ -1030,24 +1036,24 @@ class WanAnimateTransformer3DModel(
@register_to_config
def __init__(
self,
- patch_size: Tuple[int] = (1, 2, 2),
+ patch_size: tuple[int] = (1, 2, 2),
num_attention_heads: int = 40,
attention_head_dim: int = 128,
- in_channels: Optional[int] = 36,
- latent_channels: Optional[int] = 16,
- out_channels: Optional[int] = 16,
+ in_channels: int | None = 36,
+ latent_channels: int | None = 16,
+ out_channels: int | None = 16,
text_dim: int = 4096,
freq_dim: int = 256,
ffn_dim: int = 13824,
num_layers: int = 40,
cross_attn_norm: bool = True,
- qk_norm: Optional[str] = "rms_norm_across_heads",
+ qk_norm: str | None = "rms_norm_across_heads",
eps: float = 1e-6,
- image_dim: Optional[int] = 1280,
- added_kv_proj_dim: Optional[int] = None,
+ image_dim: int | None = 1280,
+ added_kv_proj_dim: int | None = None,
rope_max_seq_len: int = 1024,
- pos_embed_seq_len: Optional[int] = None,
- motion_encoder_channel_sizes: Optional[Dict[str, int]] = None, # Start of Wan Animate-specific args
+ pos_embed_seq_len: int | None = None,
+ motion_encoder_channel_sizes: dict[str, int] | None = None, # Start of Wan Animate-specific args
motion_encoder_size: int = 512,
motion_style_dim: int = 512,
motion_dim: int = 20,
@@ -1141,18 +1147,19 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
- pose_hidden_states: Optional[torch.Tensor] = None,
- face_pixel_values: Optional[torch.Tensor] = None,
- motion_encode_batch_size: Optional[int] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
+ pose_hidden_states: torch.Tensor | None = None,
+ face_pixel_values: torch.Tensor | None = None,
+ motion_encode_batch_size: int | None = None,
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> torch.Tensor | dict[str, torch.Tensor]:
"""
Forward pass of Wan2.2-Animate transformer model.
@@ -1179,21 +1186,6 @@ def forward(
Whether to return the output as a dict or tuple.
"""
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
# Check that shapes match up
if pose_hidden_states is not None and pose_hidden_states.shape[2] + 1 != hidden_states.shape[2]:
raise ValueError(
@@ -1294,10 +1286,6 @@ def forward(
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_wan_vace.py b/src/diffusers/models/transformers/transformer_wan_vace.py
index 1be4f73e33e2..7c2e205ee3ed 100644
--- a/src/diffusers/models/transformers/transformer_wan_vace.py
+++ b/src/diffusers/models/transformers/transformer_wan_vace.py
@@ -13,14 +13,14 @@
# limitations under the License.
import math
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import apply_lora_scale, logging
from ..attention import AttentionMixin, FeedForward
from ..cache_utils import CacheMixin
from ..modeling_outputs import Transformer2DModelOutput
@@ -47,7 +47,7 @@ def __init__(
qk_norm: str = "rms_norm_across_heads",
cross_attn_norm: bool = False,
eps: float = 1e-6,
- added_kv_proj_dim: Optional[int] = None,
+ added_kv_proj_dim: int | None = None,
apply_input_projection: bool = False,
apply_output_projection: bool = False,
):
@@ -76,6 +76,7 @@ def __init__(
eps=eps,
added_kv_proj_dim=added_kv_proj_dim,
processor=WanAttnProcessor(),
+ is_cross_attention=True,
)
self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
@@ -141,7 +142,7 @@ class WanVACETransformer3DModel(
A Transformer model for video-like data used in the Wan model.
Args:
- patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`):
+ patch_size (`tuple[int]`, defaults to `(1, 2, 2)`):
3D patch dimensions for video embedding (t_patch, h_patch, w_patch).
num_attention_heads (`int`, defaults to `40`):
Fixed length for text embeddings.
@@ -159,7 +160,7 @@ class WanVACETransformer3DModel(
Intermediate dimension in feed-forward network.
num_layers (`int`, defaults to `40`):
The number of layers of transformer blocks to use.
- window_size (`Tuple[int]`, defaults to `(-1, -1)`):
+ window_size (`tuple[int]`, defaults to `(-1, -1)`):
Window size for local attention (-1 indicates global attention).
cross_attn_norm (`bool`, defaults to `True`):
Enable cross-attention normalization.
@@ -178,11 +179,12 @@ class WanVACETransformer3DModel(
_no_split_modules = ["WanTransformerBlock", "WanVACETransformerBlock"]
_keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"]
_keys_to_ignore_on_load_unexpected = ["norm_added_q"]
+ _repeated_blocks = ["WanTransformerBlock", "WanVACETransformerBlock"]
@register_to_config
def __init__(
self,
- patch_size: Tuple[int, ...] = (1, 2, 2),
+ patch_size: tuple[int, ...] = (1, 2, 2),
num_attention_heads: int = 40,
attention_head_dim: int = 128,
in_channels: int = 16,
@@ -192,13 +194,13 @@ def __init__(
ffn_dim: int = 13824,
num_layers: int = 40,
cross_attn_norm: bool = True,
- qk_norm: Optional[str] = "rms_norm_across_heads",
+ qk_norm: str | None = "rms_norm_across_heads",
eps: float = 1e-6,
- image_dim: Optional[int] = None,
- added_kv_proj_dim: Optional[int] = None,
+ image_dim: int | None = None,
+ added_kv_proj_dim: int | None = None,
rope_max_seq_len: int = 1024,
- pos_embed_seq_len: Optional[int] = None,
- vace_layers: List[int] = [0, 5, 10, 15, 20, 25, 30, 35],
+ pos_embed_seq_len: int | None = None,
+ vace_layers: list[int] = [0, 5, 10, 15, 20, 25, 30, 35],
vace_in_channels: int = 96,
) -> None:
super().__init__()
@@ -261,32 +263,18 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
- encoder_hidden_states_image: Optional[torch.Tensor] = None,
+ encoder_hidden_states_image: torch.Tensor | None = None,
control_hidden_states: torch.Tensor = None,
control_hidden_states_scale: torch.Tensor = None,
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
- else:
- if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
- logger.warning(
- "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
- )
-
+ attention_kwargs: dict[str, Any] | None = None,
+ ) -> torch.Tensor | dict[str, torch.Tensor]:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.config.patch_size
post_patch_num_frames = num_frames // p_t
@@ -379,10 +367,6 @@ def forward(
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (output,)
diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py
index 5983c34ab640..3bbf78bc5e01 100644
--- a/src/diffusers/models/transformers/transformer_z_image.py
+++ b/src/diffusers/models/transformers/transformer_z_image.py
@@ -13,7 +13,6 @@
# limitations under the License.
import math
-from typing import Dict, List, Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -92,9 +91,9 @@ def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- freqs_cis: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ freqs_cis: torch.Tensor | None = None,
) -> torch.Tensor:
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
@@ -229,10 +228,10 @@ def forward(
x: torch.Tensor,
attn_mask: torch.Tensor,
freqs_cis: torch.Tensor,
- adaln_input: Optional[torch.Tensor] = None,
- noise_mask: Optional[torch.Tensor] = None,
- adaln_noisy: Optional[torch.Tensor] = None,
- adaln_clean: Optional[torch.Tensor] = None,
+ adaln_input: torch.Tensor | None = None,
+ noise_mask: torch.Tensor | None = None,
+ adaln_noisy: torch.Tensor | None = None,
+ adaln_clean: torch.Tensor | None = None,
):
if self.modulation:
seq_len = x.shape[1]
@@ -315,8 +314,8 @@ class RopeEmbedder:
def __init__(
self,
theta: float = 256.0,
- axes_dims: List[int] = (16, 56, 56),
- axes_lens: List[int] = (64, 128, 128),
+ axes_dims: list[int] = (16, 56, 56),
+ axes_lens: list[int] = (64, 128, 128),
):
self.theta = theta
self.axes_dims = axes_dims
@@ -325,7 +324,7 @@ def __init__(
self.freqs_cis = None
@staticmethod
- def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0):
+ def precompute_freqs_cis(dim: list[int], end: list[int], theta: float = 256.0):
with torch.device("cpu"):
freqs_cis = []
for i, (d, e) in enumerate(zip(dim, end)):
@@ -482,12 +481,12 @@ def __init__(
def unpatchify(
self,
- x: List[torch.Tensor],
- size: List[Tuple],
+ x: list[torch.Tensor],
+ size: list[tuple],
patch_size,
f_patch_size,
- x_pos_offsets: Optional[List[Tuple[int, int]]] = None,
- ) -> List[torch.Tensor]:
+ x_pos_offsets: list[tuple[int, int]] | None = None,
+ ) -> list[torch.Tensor]:
pH = pW = patch_size
pF = f_patch_size
bsz = len(x)
@@ -552,10 +551,10 @@ def _patchify_image(self, image: torch.Tensor, patch_size: int, f_patch_size: in
def _pad_with_ids(
self,
feat: torch.Tensor,
- pos_grid_size: Tuple,
- pos_start: Tuple,
+ pos_grid_size: tuple,
+ pos_start: tuple,
device: torch.device,
- noise_mask_val: Optional[int] = None,
+ noise_mask_val: int | None = None,
):
"""Pad feature to SEQ_MULTI_OF, create position IDs and pad mask."""
ori_len = len(feat)
@@ -587,7 +586,7 @@ def _pad_with_ids(
return padded_feat, pos_ids, pad_mask, total_len, noise_mask
def patchify_and_embed(
- self, all_image: List[torch.Tensor], all_cap_feats: List[torch.Tensor], patch_size: int, f_patch_size: int
+ self, all_image: list[torch.Tensor], all_cap_feats: list[torch.Tensor], patch_size: int, f_patch_size: int
):
"""Patchify for basic mode: single image per batch item."""
device = all_image[0].device
@@ -625,12 +624,12 @@ def patchify_and_embed(
def patchify_and_embed_omni(
self,
- all_x: List[List[torch.Tensor]],
- all_cap_feats: List[List[torch.Tensor]],
- all_siglip_feats: List[List[torch.Tensor]],
+ all_x: list[list[torch.Tensor]],
+ all_cap_feats: list[list[torch.Tensor]],
+ all_siglip_feats: list[list[torch.Tensor]],
patch_size: int,
f_patch_size: int,
- images_noise_mask: List[List[int]],
+ images_noise_mask: list[list[int]],
):
"""Patchify for omni mode: multiple images per batch item with noise masks."""
bsz = len(all_x)
@@ -764,11 +763,11 @@ def patchify_and_embed_omni(
def _prepare_sequence(
self,
- feats: List[torch.Tensor],
- pos_ids: List[torch.Tensor],
- inner_pad_mask: List[torch.Tensor],
+ feats: list[torch.Tensor],
+ pos_ids: list[torch.Tensor],
+ inner_pad_mask: list[torch.Tensor],
pad_token: torch.nn.Parameter,
- noise_mask: Optional[List[List[int]]] = None,
+ noise_mask: list[list[int]] | None = None,
device: torch.device = None,
):
"""Prepare sequence: apply pad token, RoPE embed, pad to batch, create attention mask."""
@@ -808,16 +807,16 @@ def _build_unified_sequence(
self,
x: torch.Tensor,
x_freqs: torch.Tensor,
- x_seqlens: List[int],
- x_noise_mask: Optional[List[List[int]]],
+ x_seqlens: list[int],
+ x_noise_mask: list[list[int]] | None,
cap: torch.Tensor,
cap_freqs: torch.Tensor,
- cap_seqlens: List[int],
- cap_noise_mask: Optional[List[List[int]]],
- siglip: Optional[torch.Tensor],
- siglip_freqs: Optional[torch.Tensor],
- siglip_seqlens: Optional[List[int]],
- siglip_noise_mask: Optional[List[List[int]]],
+ cap_seqlens: list[int],
+ cap_noise_mask: list[list[int]] | None,
+ siglip: torch.Tensor | None,
+ siglip_freqs: torch.Tensor | None,
+ siglip_seqlens: list[int] | None,
+ siglip_noise_mask: list[list[int]] | None,
omni_mode: bool,
device: torch.device,
):
@@ -887,13 +886,13 @@ def _build_unified_sequence(
def forward(
self,
- x: Union[List[torch.Tensor], List[List[torch.Tensor]]],
+ x: list[torch.Tensor, list[list[torch.Tensor]]],
t,
- cap_feats: Union[List[torch.Tensor], List[List[torch.Tensor]]],
+ cap_feats: list[torch.Tensor, list[list[torch.Tensor]]],
return_dict: bool = True,
- controlnet_block_samples: Optional[Dict[int, torch.Tensor]] = None,
- siglip_feats: Optional[List[List[torch.Tensor]]] = None,
- image_noise_mask: Optional[List[List[int]]] = None,
+ controlnet_block_samples: dict[int, torch.Tensor] | None = None,
+ siglip_feats: list[list[torch.Tensor]] | None = None,
+ image_noise_mask: list[list[int]] | None = None,
patch_size: int = 2,
f_patch_size: int = 1,
):
diff --git a/src/diffusers/models/unets/unet_1d.py b/src/diffusers/models/unets/unet_1d.py
index a027c553ed06..83ffe1f6f8cb 100644
--- a/src/diffusers/models/unets/unet_1d.py
+++ b/src/diffusers/models/unets/unet_1d.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -56,12 +55,12 @@ class UNet1DModel(ModelMixin, ConfigMixin):
freq_shift (`float`, *optional*, defaults to 0.0): Frequency shift for Fourier time embedding.
flip_sin_to_cos (`bool`, *optional*, defaults to `False`):
Whether to flip sin to cos for Fourier time embedding.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")`):
- Tuple of downsample block types.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")`):
- Tuple of upsample block types.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(32, 32, 64)`):
- Tuple of block output channels.
+ down_block_types (`tuple[str]`, *optional*, defaults to `("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")`):
+ tuple of downsample block types.
+ up_block_types (`tuple[str]`, *optional*, defaults to `("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")`):
+ tuple of upsample block types.
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(32, 32, 64)`):
+ tuple of block output channels.
mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock1D"`): Block type for middle of UNet.
out_block_type (`str`, *optional*, defaults to `None`): Optional output processing block of UNet.
act_fn (`str`, *optional*, defaults to `None`): Optional activation function in UNet blocks.
@@ -77,20 +76,20 @@ class UNet1DModel(ModelMixin, ConfigMixin):
def __init__(
self,
sample_size: int = 65536,
- sample_rate: Optional[int] = None,
+ sample_rate: int | None = None,
in_channels: int = 2,
out_channels: int = 2,
extra_in_channels: int = 0,
time_embedding_type: str = "fourier",
- time_embedding_dim: Optional[int] = None,
+ time_embedding_dim: int | None = None,
flip_sin_to_cos: bool = True,
use_timestep_embedding: bool = False,
freq_shift: float = 0.0,
- down_block_types: Tuple[str, ...] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D"),
- up_block_types: Tuple[str, ...] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip"),
+ down_block_types: tuple[str, ...] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D"),
+ up_block_types: tuple[str, ...] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip"),
mid_block_type: str = "UNetMidBlock1D",
out_block_type: str = None,
- block_out_channels: Tuple[int, ...] = (32, 32, 64),
+ block_out_channels: tuple[int, ...] = (32, 32, 64),
act_fn: str = None,
norm_num_groups: int = 8,
layers_per_block: int = 1,
@@ -206,9 +205,9 @@ def __init__(
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
return_dict: bool = True,
- ) -> Union[UNet1DOutput, Tuple]:
+ ) -> UNet1DOutput | tuple:
r"""
The [`UNet1DModel`] forward method.
diff --git a/src/diffusers/models/unets/unet_1d_blocks.py b/src/diffusers/models/unets/unet_1d_blocks.py
index 58cbdfd005b6..a21521eaed4c 100644
--- a/src/diffusers/models/unets/unet_1d_blocks.py
+++ b/src/diffusers/models/unets/unet_1d_blocks.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import math
-from typing import Optional, Tuple, Union
import torch
import torch.nn.functional as F
@@ -26,13 +25,13 @@ class DownResnetBlock1D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 1,
conv_shortcut: bool = False,
temb_channels: int = 32,
groups: int = 32,
- groups_out: Optional[int] = None,
- non_linearity: Optional[str] = None,
+ groups_out: int | None = None,
+ non_linearity: str | None = None,
time_embedding_norm: str = "default",
output_scale_factor: float = 1.0,
add_downsample: bool = True,
@@ -66,7 +65,7 @@ def __init__(
if add_downsample:
self.downsample = Downsample1D(out_channels, use_conv=True, padding=1)
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
output_states = ()
hidden_states = self.resnets[0](hidden_states, temb)
@@ -88,12 +87,12 @@ class UpResnetBlock1D(nn.Module):
def __init__(
self,
in_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
num_layers: int = 1,
temb_channels: int = 32,
groups: int = 32,
- groups_out: Optional[int] = None,
- non_linearity: Optional[str] = None,
+ groups_out: int | None = None,
+ non_linearity: str | None = None,
time_embedding_norm: str = "default",
output_scale_factor: float = 1.0,
add_upsample: bool = True,
@@ -129,8 +128,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Optional[Tuple[torch.Tensor, ...]] = None,
- temb: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...] | None = None,
+ temb: torch.Tensor | None = None,
) -> torch.Tensor:
if res_hidden_states_tuple is not None:
res_hidden_states = res_hidden_states_tuple[-1]
@@ -161,7 +160,7 @@ def __init__(self, in_channels: int, out_channels: int, embed_dim: int):
self.res2 = ResidualTemporalBlock1D(in_channels // 2, in_channels // 4, embed_dim=embed_dim)
self.down2 = Downsample1D(out_channels // 4, use_conv=True)
- def forward(self, x: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, x: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
x = self.res1(x, temb)
x = self.down1(x)
x = self.res2(x, temb)
@@ -178,7 +177,7 @@ def __init__(
num_layers: int = 1,
add_downsample: bool = False,
add_upsample: bool = False,
- non_linearity: Optional[str] = None,
+ non_linearity: str | None = None,
):
super().__init__()
self.in_channels = in_channels
@@ -230,7 +229,7 @@ def __init__(self, num_groups_out: int, out_channels: int, embed_dim: int, act_f
self.final_conv1d_act = get_activation(act_fn)
self.final_conv1d_2 = nn.Conv1d(embed_dim, out_channels, 1)
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = self.final_conv1d_1(hidden_states)
hidden_states = rearrange_dims(hidden_states)
hidden_states = self.final_conv1d_gn(hidden_states)
@@ -305,7 +304,7 @@ def __init__(self, kernel: str = "linear", pad_mode: str = "reflect"):
self.pad = kernel_1d.shape[0] // 2 - 1
self.register_buffer("kernel", kernel_1d)
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = F.pad(hidden_states, ((self.pad + 1) // 2,) * 2, self.pad_mode)
weight = hidden_states.new_zeros([hidden_states.shape[1], hidden_states.shape[1], self.kernel.shape[0]])
indices = torch.arange(hidden_states.shape[1], device=hidden_states.device)
@@ -407,7 +406,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
class UNetMidBlock1D(nn.Module):
- def __init__(self, mid_channels: int, in_channels: int, out_channels: Optional[int] = None):
+ def __init__(self, mid_channels: int, in_channels: int, out_channels: int | None = None):
super().__init__()
out_channels = in_channels if out_channels is None else out_channels
@@ -435,7 +434,7 @@ def __init__(self, mid_channels: int, in_channels: int, out_channels: Optional[i
self.attentions = nn.ModuleList(attentions)
self.resnets = nn.ModuleList(resnets)
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = self.down(hidden_states)
for attn, resnet in zip(self.attentions, self.resnets):
hidden_states = resnet(hidden_states)
@@ -447,7 +446,7 @@ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = No
class AttnDownBlock1D(nn.Module):
- def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None):
+ def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None = None):
super().__init__()
mid_channels = out_channels if mid_channels is None else mid_channels
@@ -466,7 +465,7 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[i
self.attentions = nn.ModuleList(attentions)
self.resnets = nn.ModuleList(resnets)
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = self.down(hidden_states)
for resnet, attn in zip(self.resnets, self.attentions):
@@ -477,7 +476,7 @@ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = No
class DownBlock1D(nn.Module):
- def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None):
+ def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None = None):
super().__init__()
mid_channels = out_channels if mid_channels is None else mid_channels
@@ -490,7 +489,7 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[i
self.resnets = nn.ModuleList(resnets)
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = self.down(hidden_states)
for resnet in self.resnets:
@@ -500,7 +499,7 @@ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = No
class DownBlock1DNoSkip(nn.Module):
- def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None):
+ def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None = None):
super().__init__()
mid_channels = out_channels if mid_channels is None else mid_channels
@@ -512,7 +511,7 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[i
self.resnets = nn.ModuleList(resnets)
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = torch.cat([hidden_states, temb], dim=1)
for resnet in self.resnets:
hidden_states = resnet(hidden_states)
@@ -521,7 +520,7 @@ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = No
class AttnUpBlock1D(nn.Module):
- def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None):
+ def __init__(self, in_channels: int, out_channels: int, mid_channels: int | None = None):
super().__init__()
mid_channels = out_channels if mid_channels is None else mid_channels
@@ -543,8 +542,8 @@ def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[i
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
) -> torch.Tensor:
res_hidden_states = res_hidden_states_tuple[-1]
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
@@ -559,7 +558,7 @@ def forward(
class UpBlock1D(nn.Module):
- def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None):
+ def __init__(self, in_channels: int, out_channels: int, mid_channels: int | None = None):
super().__init__()
mid_channels = in_channels if mid_channels is None else mid_channels
@@ -575,8 +574,8 @@ def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[i
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
) -> torch.Tensor:
res_hidden_states = res_hidden_states_tuple[-1]
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
@@ -590,7 +589,7 @@ def forward(
class UpBlock1DNoSkip(nn.Module):
- def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None):
+ def __init__(self, in_channels: int, out_channels: int, mid_channels: int | None = None):
super().__init__()
mid_channels = in_channels if mid_channels is None else mid_channels
@@ -605,8 +604,8 @@ def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[i
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
) -> torch.Tensor:
res_hidden_states = res_hidden_states_tuple[-1]
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
@@ -617,10 +616,10 @@ def forward(
return hidden_states
-DownBlockType = Union[DownResnetBlock1D, DownBlock1D, AttnDownBlock1D, DownBlock1DNoSkip]
-MidBlockType = Union[MidResTemporalBlock1D, ValueFunctionMidBlock1D, UNetMidBlock1D]
-OutBlockType = Union[OutConv1DBlock, OutValueFunctionBlock]
-UpBlockType = Union[UpResnetBlock1D, UpBlock1D, AttnUpBlock1D, UpBlock1DNoSkip]
+DownBlockType = DownResnetBlock1D | DownBlock1D | AttnDownBlock1D | DownBlock1DNoSkip
+MidBlockType = MidResTemporalBlock1D | ValueFunctionMidBlock1D | UNetMidBlock1D
+OutBlockType = OutConv1DBlock | OutValueFunctionBlock
+UpBlockType = UpResnetBlock1D | UpBlock1D | AttnUpBlock1D | UpBlock1DNoSkip
def get_down_block(
@@ -694,7 +693,7 @@ def get_mid_block(
def get_out_block(
*, out_block_type: str, num_groups_out: int, embed_dim: int, out_channels: int, act_fn: str, fc_dim: int
-) -> Optional[OutBlockType]:
+) -> OutBlockType | None:
if out_block_type == "OutConv1DBlock":
return OutConv1DBlock(num_groups_out, out_channels, embed_dim, act_fn)
elif out_block_type == "ValueFunction":
diff --git a/src/diffusers/models/unets/unet_2d.py b/src/diffusers/models/unets/unet_2d.py
index 2588a9c518bd..4e54f757d120 100644
--- a/src/diffusers/models/unets/unet_2d.py
+++ b/src/diffusers/models/unets/unet_2d.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -45,7 +44,7 @@ class UNet2DModel(ModelMixin, ConfigMixin):
for all models (such as downloading or saving).
Parameters:
- sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+ sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`):
Height and width of input/output sample. Dimensions must be a multiple of `2 ** (len(block_out_channels) -
1)`.
in_channels (`int`, *optional*, defaults to 3): Number of channels in the input sample.
@@ -55,14 +54,14 @@ class UNet2DModel(ModelMixin, ConfigMixin):
freq_shift (`int`, *optional*, defaults to 0): Frequency shift for Fourier time embedding.
flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
Whether to flip sin to cos for Fourier time embedding.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`):
- Tuple of downsample block types.
+ down_block_types (`tuple[str]`, *optional*, defaults to `("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`):
+ tuple of downsample block types.
mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2D"`):
Block type for middle of UNet, it can be either `UNetMidBlock2D` or `None`.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`):
- Tuple of upsample block types.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(224, 448, 672, 896)`):
- Tuple of block output channels.
+ up_block_types (`tuple[str]`, *optional*, defaults to `("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`):
+ tuple of upsample block types.
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(224, 448, 672, 896)`):
+ tuple of block output channels.
layers_per_block (`int`, *optional*, defaults to `2`): The number of layers per block.
mid_block_scale_factor (`float`, *optional*, defaults to `1`): The scale factor for the mid block.
downsample_padding (`int`, *optional*, defaults to `1`): The padding for the downsample convolution.
@@ -95,18 +94,18 @@ class UNet2DModel(ModelMixin, ConfigMixin):
@register_to_config
def __init__(
self,
- sample_size: Optional[Union[int, Tuple[int, int]]] = None,
+ sample_size: int | tuple[int, int] | None = None,
in_channels: int = 3,
out_channels: int = 3,
center_input_sample: bool = False,
time_embedding_type: str = "positional",
- time_embedding_dim: Optional[int] = None,
+ time_embedding_dim: int | None = None,
freq_shift: int = 0,
flip_sin_to_cos: bool = True,
- down_block_types: Tuple[str, ...] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
- mid_block_type: Optional[str] = "UNetMidBlock2D",
- up_block_types: Tuple[str, ...] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
- block_out_channels: Tuple[int, ...] = (224, 448, 672, 896),
+ down_block_types: tuple[str, ...] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
+ mid_block_type: str | None = "UNetMidBlock2D",
+ up_block_types: tuple[str, ...] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
+ block_out_channels: tuple[int, ...] = (224, 448, 672, 896),
layers_per_block: int = 2,
mid_block_scale_factor: float = 1,
downsample_padding: int = 1,
@@ -114,15 +113,15 @@ def __init__(
upsample_type: str = "conv",
dropout: float = 0.0,
act_fn: str = "silu",
- attention_head_dim: Optional[int] = 8,
+ attention_head_dim: int | None = 8,
norm_num_groups: int = 32,
- attn_norm_num_groups: Optional[int] = None,
+ attn_norm_num_groups: int | None = None,
norm_eps: float = 1e-5,
resnet_time_scale_shift: str = "default",
add_attention: bool = True,
- class_embed_type: Optional[str] = None,
- num_class_embeds: Optional[int] = None,
- num_train_timesteps: Optional[int] = None,
+ class_embed_type: str | None = None,
+ num_class_embeds: int | None = None,
+ num_train_timesteps: int | None = None,
):
super().__init__()
@@ -250,10 +249,10 @@ def __init__(
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
- class_labels: Optional[torch.Tensor] = None,
+ timestep: torch.Tensor | float | int,
+ class_labels: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[UNet2DOutput, Tuple]:
+ ) -> UNet2DOutput | tuple:
r"""
The [`UNet2DModel`] forward method.
diff --git a/src/diffusers/models/unets/unet_2d_blocks.py b/src/diffusers/models/unets/unet_2d_blocks.py
index 94a9245e567c..18a0b4ec5659 100644
--- a/src/diffusers/models/unets/unet_2d_blocks.py
+++ b/src/diffusers/models/unets/unet_2d_blocks.py
@@ -11,7 +11,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.
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -50,10 +50,10 @@ def get_down_block(
resnet_eps: float,
resnet_act_fn: str,
transformer_layers_per_block: int = 1,
- num_attention_heads: Optional[int] = None,
- resnet_groups: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
- downsample_padding: Optional[int] = None,
+ num_attention_heads: int | None = None,
+ resnet_groups: int | None = None,
+ cross_attention_dim: int | None = None,
+ downsample_padding: int | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
only_cross_attention: bool = False,
@@ -62,9 +62,9 @@ def get_down_block(
attention_type: str = "default",
resnet_skip_time_act: bool = False,
resnet_out_scale_factor: float = 1.0,
- cross_attention_norm: Optional[str] = None,
- attention_head_dim: Optional[int] = None,
- downsample_type: Optional[str] = None,
+ cross_attention_norm: str | None = None,
+ attention_head_dim: int | None = None,
+ downsample_type: str | None = None,
dropout: float = 0.0,
):
# If attn head dim is not defined, we default it to the number of heads
@@ -258,8 +258,8 @@ def get_mid_block(
resnet_groups: int,
output_scale_factor: float = 1.0,
transformer_layers_per_block: int = 1,
- num_attention_heads: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
+ num_attention_heads: int | None = None,
+ cross_attention_dim: int | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
mid_block_only_cross_attention: bool = False,
@@ -267,8 +267,8 @@ def get_mid_block(
resnet_time_scale_shift: str = "default",
attention_type: str = "default",
resnet_skip_time_act: bool = False,
- cross_attention_norm: Optional[str] = None,
- attention_head_dim: Optional[int] = 1,
+ cross_attention_norm: str | None = None,
+ attention_head_dim: int | None = 1,
dropout: float = 0.0,
):
if mid_block_type == "UNetMidBlock2DCrossAttn":
@@ -334,11 +334,11 @@ def get_up_block(
add_upsample: bool,
resnet_eps: float,
resnet_act_fn: str,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
transformer_layers_per_block: int = 1,
- num_attention_heads: Optional[int] = None,
- resnet_groups: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
+ num_attention_heads: int | None = None,
+ resnet_groups: int | None = None,
+ cross_attention_dim: int | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
only_cross_attention: bool = False,
@@ -347,9 +347,9 @@ def get_up_block(
attention_type: str = "default",
resnet_skip_time_act: bool = False,
resnet_out_scale_factor: float = 1.0,
- cross_attention_norm: Optional[str] = None,
- attention_head_dim: Optional[int] = None,
- upsample_type: Optional[str] = None,
+ cross_attention_norm: str | None = None,
+ attention_head_dim: int | None = None,
+ upsample_type: str | None = None,
dropout: float = 0.0,
) -> nn.Module:
# If attn head dim is not defined, we default it to the number of heads
@@ -602,7 +602,7 @@ class UNetMidBlock2D(nn.Module):
resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.
resnet_groups (`int`, *optional*, defaults to 32):
The number of groups to use in the group normalization layers of the resnet blocks.
- attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.
+ attn_groups (`int | None`, *optional*, defaults to None): The number of groups for the attention blocks.
resnet_pre_norm (`bool`, *optional*, defaults to `True`):
Whether to use pre-normalization for the resnet blocks.
add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.
@@ -627,7 +627,7 @@ def __init__(
resnet_time_scale_shift: str = "default", # default, spatial
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
- attn_groups: Optional[int] = None,
+ attn_groups: int | None = None,
resnet_pre_norm: bool = True,
add_attention: bool = True,
attention_head_dim: int = 1,
@@ -733,7 +733,7 @@ def __init__(
self.gradient_checkpointing = False
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = self.resnets[0](hidden_states, temb)
for attn, resnet in zip(self.attentions, self.resnets[1:]):
if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -753,15 +753,15 @@ def __init__(
self,
in_channels: int,
temb_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
- resnet_groups_out: Optional[int] = None,
+ resnet_groups_out: int | None = None,
resnet_pre_norm: bool = True,
num_attention_heads: int = 1,
output_scale_factor: float = 1.0,
@@ -854,11 +854,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
@@ -907,7 +907,7 @@ def __init__(
cross_attention_dim: int = 1280,
skip_time_act: bool = False,
only_cross_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
):
super().__init__()
@@ -978,11 +978,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
if cross_attention_kwargs.get("scale", None) is not None:
@@ -1112,10 +1112,10 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ temb: torch.Tensor | None = None,
+ upsample_size: int | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
if cross_attention_kwargs.get("scale", None) is not None:
logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
@@ -1152,7 +1152,7 @@ def __init__(
temb_channels: int,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
@@ -1239,13 +1239,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- additional_residuals: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ additional_residuals: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
@@ -1344,8 +1344,8 @@ def __init__(
self.gradient_checkpointing = False
def forward(
- self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, *args, **kwargs
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -1635,11 +1635,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- skip_sample: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ skip_sample: torch.Tensor | None = None,
*args,
**kwargs,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...], torch.Tensor]:
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...], torch.Tensor]:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -1726,11 +1726,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- skip_sample: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ skip_sample: torch.Tensor | None = None,
*args,
**kwargs,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...], torch.Tensor]:
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...], torch.Tensor]:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -1818,8 +1818,8 @@ def __init__(
self.gradient_checkpointing = False
def forward(
- self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, *args, **kwargs
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -1862,7 +1862,7 @@ def __init__(
add_downsample: bool = True,
skip_time_act: bool = False,
only_cross_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
):
super().__init__()
@@ -1941,12 +1941,12 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
if cross_attention_kwargs.get("scale", None) is not None:
logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
@@ -2041,8 +2041,8 @@ def __init__(
self.gradient_checkpointing = False
def forward(
- self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, *args, **kwargs
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -2132,12 +2132,12 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
if cross_attention_kwargs.get("scale", None) is not None:
logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
@@ -2278,9 +2278,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ upsample_size: int | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -2318,10 +2318,10 @@ def __init__(
out_channels: int,
prev_output_channel: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
@@ -2405,13 +2405,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ upsample_size: int | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
@@ -2478,7 +2478,7 @@ def __init__(
prev_output_channel: int,
out_channels: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -2524,9 +2524,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ upsample_size: int | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -2577,7 +2577,7 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -2587,7 +2587,7 @@ def __init__(
resnet_pre_norm: bool = True,
output_scale_factor: float = 1.0,
add_upsample: bool = True,
- temb_channels: Optional[int] = None,
+ temb_channels: int | None = None,
):
super().__init__()
resnets = []
@@ -2634,7 +2634,7 @@ def __init__(
self.resolution_idx = resolution_idx
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
for resnet in self.resnets:
hidden_states = resnet(hidden_states, temb=temb)
@@ -2650,7 +2650,7 @@ def __init__(
self,
in_channels: int,
out_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -2661,7 +2661,7 @@ def __init__(
attention_head_dim: int = 1,
output_scale_factor: float = 1.0,
add_upsample: bool = True,
- temb_channels: Optional[int] = None,
+ temb_channels: int | None = None,
):
super().__init__()
resnets = []
@@ -2732,7 +2732,7 @@ def __init__(
self.resolution_idx = resolution_idx
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
for resnet, attn in zip(self.resnets, self.attentions):
hidden_states = resnet(hidden_states, temb=temb)
hidden_states = attn(hidden_states, temb=temb)
@@ -2751,7 +2751,7 @@ def __init__(
prev_output_channel: int,
out_channels: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -2841,12 +2841,12 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
skip_sample=None,
*args,
**kwargs,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -2885,7 +2885,7 @@ def __init__(
prev_output_channel: int,
out_channels: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -2953,12 +2953,12 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
skip_sample=None,
*args,
**kwargs,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -2995,7 +2995,7 @@ def __init__(
prev_output_channel: int,
out_channels: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -3060,9 +3060,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ upsample_size: int | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -3095,7 +3095,7 @@ def __init__(
out_channels: int,
prev_output_channel: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -3109,7 +3109,7 @@ def __init__(
add_upsample: bool = True,
skip_time_act: bool = False,
only_cross_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
):
super().__init__()
resnets = []
@@ -3190,13 +3190,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ upsample_size: int | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
if cross_attention_kwargs.get("scale", None) is not None:
@@ -3256,7 +3256,7 @@ def __init__(
num_layers: int = 5,
resnet_eps: float = 1e-5,
resnet_act_fn: str = "gelu",
- resnet_group_size: Optional[int] = 32,
+ resnet_group_size: int | None = 32,
add_upsample: bool = True,
):
super().__init__()
@@ -3298,9 +3298,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ upsample_size: int | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -3414,13 +3414,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ upsample_size: int | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
res_hidden_states_tuple = res_hidden_states_tuple[-1]
if res_hidden_states_tuple is not None:
@@ -3490,12 +3490,12 @@ def __init__(
num_attention_heads: int,
attention_head_dim: int,
dropout: float = 0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
upcast_attention: bool = False,
temb_channels: int = 768, # for ada_group_norm
add_self_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
group_size: int = 32,
):
super().__init__()
@@ -3536,13 +3536,13 @@ def _to_4d(self, hidden_states: torch.Tensor, height: int, weight: int) -> torch
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.Tensor] = None,
+ encoder_hidden_states: torch.Tensor | None = None,
# TODO: mark emb as non-optional (self.norm2 requires it).
# requires assessing impact of change to positional param interface.
- emb: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ emb: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
if cross_attention_kwargs.get("scale", None) is not None:
diff --git a/src/diffusers/models/unets/unet_2d_condition.py b/src/diffusers/models/unets/unet_2d_condition.py
index e669aa51a54e..deae25899475 100644
--- a/src/diffusers/models/unets/unet_2d_condition.py
+++ b/src/diffusers/models/unets/unet_2d_condition.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -20,7 +20,12 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin
from ...loaders.single_file_model import FromOriginalModelMixin
-from ...utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from ...utils import (
+ BaseOutput,
+ apply_lora_scale,
+ deprecate,
+ logging,
+)
from ..activations import get_activation
from ..attention import AttentionMixin
from ..attention_processor import (
@@ -78,7 +83,7 @@ class UNet2DConditionModel(
for all models (such as downloading or saving).
Parameters:
- sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+ sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`):
Height and width of input/output sample.
in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
@@ -86,17 +91,17 @@ class UNet2DConditionModel(
flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
Whether to flip the sin to cos in the time embedding.
freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
The tuple of downsample blocks to use.
mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or
`UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
The tuple of upsample blocks to use.
- only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
+ only_cross_attention(`bool` or `tuple[bool]`, *optional*, default to `False`):
Whether to include self-attention in the basic transformer blocks, see
[`~models.attention.BasicTransformerBlock`].
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
@@ -106,15 +111,15 @@ class UNet2DConditionModel(
norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
If `None`, normalization and activation layers is skipped in post-processing.
norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
- cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+ cross_attention_dim (`int` or `tuple[int]`, *optional*, defaults to 1280):
The dimension of the cross attention features.
- transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+ transformer_layers_per_block (`int`, `tuple[int]`, or `tuple[tuple]` , *optional*, defaults to 1):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
[`~models.unets.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unets.unet_2d_blocks.CrossAttnUpBlock2D`],
[`~models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
- reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None):
+ reverse_transformer_layers_per_block : (`tuple[tuple]`, *optional*, defaults to None):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling
- blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for
+ blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `tuple[tuple]` and for
[`~models.unets.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unets.unet_2d_blocks.CrossAttnUpBlock2D`],
[`~models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
encoder_hid_dim (`int`, *optional*, defaults to None):
@@ -171,63 +176,63 @@ class conditioning with `class_embed_type` equal to `None`.
@register_to_config
def __init__(
self,
- sample_size: Optional[Union[int, Tuple[int, int]]] = None,
+ sample_size: int | tuple[int, int] | None = None,
in_channels: int = 4,
out_channels: int = 4,
center_input_sample: bool = False,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
- up_block_types: Tuple[str, ...] = (
+ mid_block_type: str | None = "UNetMidBlock2DCrossAttn",
+ up_block_types: tuple[str, ...] = (
"UpBlock2D",
"CrossAttnUpBlock2D",
"CrossAttnUpBlock2D",
"CrossAttnUpBlock2D",
),
- only_cross_attention: Union[bool, Tuple[bool]] = False,
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
- layers_per_block: Union[int, Tuple[int]] = 2,
+ only_cross_attention: bool | tuple[bool] = False,
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
+ layers_per_block: int | tuple[int] = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
dropout: float = 0.0,
act_fn: str = "silu",
- norm_num_groups: Optional[int] = 32,
+ norm_num_groups: int | None = 32,
norm_eps: float = 1e-5,
- cross_attention_dim: Union[int, Tuple[int]] = 1280,
- transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
- reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
- encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
- attention_head_dim: Union[int, Tuple[int]] = 8,
- num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+ cross_attention_dim: int | tuple[int] = 1280,
+ transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1,
+ reverse_transformer_layers_per_block: tuple[tuple[int]] | None = None,
+ encoder_hid_dim: int | None = None,
+ encoder_hid_dim_type: str | None = None,
+ attention_head_dim: int | tuple[int] = 8,
+ num_attention_heads: int | tuple[int] | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
- addition_time_embed_dim: Optional[int] = None,
- num_class_embeds: Optional[int] = None,
+ class_embed_type: str | None = None,
+ addition_embed_type: str | None = None,
+ addition_time_embed_dim: int | None = None,
+ num_class_embeds: int | None = None,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
resnet_skip_time_act: bool = False,
resnet_out_scale_factor: float = 1.0,
time_embedding_type: str = "positional",
- time_embedding_dim: Optional[int] = None,
- time_embedding_act_fn: Optional[str] = None,
- timestep_post_act: Optional[str] = None,
- time_cond_proj_dim: Optional[int] = None,
+ time_embedding_dim: int | None = None,
+ time_embedding_act_fn: str | None = None,
+ timestep_post_act: str | None = None,
+ time_cond_proj_dim: int | None = None,
conv_in_kernel: int = 3,
conv_out_kernel: int = 3,
- projection_class_embeddings_input_dim: Optional[int] = None,
+ projection_class_embeddings_input_dim: int | None = None,
attention_type: str = "default",
class_embeddings_concat: bool = False,
- mid_block_only_cross_attention: Optional[bool] = None,
- cross_attention_norm: Optional[str] = None,
+ mid_block_only_cross_attention: bool | None = None,
+ cross_attention_norm: str | None = None,
addition_embed_type_num_heads: int = 64,
):
super().__init__()
@@ -491,16 +496,16 @@ def __init__(
def _check_config(
self,
- down_block_types: Tuple[str, ...],
- up_block_types: Tuple[str, ...],
- only_cross_attention: Union[bool, Tuple[bool]],
- block_out_channels: Tuple[int, ...],
- layers_per_block: Union[int, Tuple[int]],
- cross_attention_dim: Union[int, Tuple[int]],
- transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple[int]]],
+ down_block_types: tuple[str, ...],
+ up_block_types: tuple[str, ...],
+ only_cross_attention: bool | tuple[bool],
+ block_out_channels: tuple[int, ...],
+ layers_per_block: int | tuple[int],
+ cross_attention_dim: int | tuple[int],
+ transformer_layers_per_block: int | tuple[int, tuple[tuple[int]]],
reverse_transformer_layers_per_block: bool,
attention_head_dim: int,
- num_attention_heads: Optional[Union[int, Tuple[int]]],
+ num_attention_heads: int | tuple[int] | None,
):
if len(down_block_types) != len(up_block_types):
raise ValueError(
@@ -548,7 +553,7 @@ def _set_time_proj(
flip_sin_to_cos: bool,
freq_shift: float,
time_embedding_dim: int,
- ) -> Tuple[int, int]:
+ ) -> tuple[int, int]:
if time_embedding_type == "fourier":
time_embed_dim = time_embedding_dim or block_out_channels[0] * 2
if time_embed_dim % 2 != 0:
@@ -571,9 +576,9 @@ def _set_time_proj(
def _set_encoder_hid_proj(
self,
- encoder_hid_dim_type: Optional[str],
- cross_attention_dim: Union[int, Tuple[int]],
- encoder_hid_dim: Optional[int],
+ encoder_hid_dim_type: str | None,
+ cross_attention_dim: int | tuple[int],
+ encoder_hid_dim: int | None,
):
if encoder_hid_dim_type is None and encoder_hid_dim is not None:
encoder_hid_dim_type = "text_proj"
@@ -611,10 +616,10 @@ def _set_encoder_hid_proj(
def _set_class_embedding(
self,
- class_embed_type: Optional[str],
+ class_embed_type: str | None,
act_fn: str,
- num_class_embeds: Optional[int],
- projection_class_embeddings_input_dim: Optional[int],
+ num_class_embeds: int | None,
+ projection_class_embeddings_input_dim: int | None,
time_embed_dim: int,
timestep_input_dim: int,
):
@@ -650,12 +655,12 @@ def _set_add_embedding(
self,
addition_embed_type: str,
addition_embed_type_num_heads: int,
- addition_time_embed_dim: Optional[int],
+ addition_time_embed_dim: int | None,
flip_sin_to_cos: bool,
freq_shift: float,
- cross_attention_dim: Optional[int],
- encoder_hid_dim: Optional[int],
- projection_class_embeddings_input_dim: Optional[int],
+ cross_attention_dim: int | None,
+ encoder_hid_dim: int | None,
+ projection_class_embeddings_input_dim: int | None,
time_embed_dim: int,
):
if addition_embed_type == "text":
@@ -716,7 +721,7 @@ def set_default_attn_processor(self):
self.set_attn_processor(processor)
- def set_attention_slice(self, slice_size: Union[str, int, List[int]] = "auto"):
+ def set_attention_slice(self, slice_size: str | int | list[int] = "auto"):
r"""
Enable sliced attention computation.
@@ -770,7 +775,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
# Recursively walk through all the children.
# Any children which exposes the set_attention_slice method
# gets the message
- def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]):
if hasattr(module, "set_attention_slice"):
module.set_attention_slice(slice_size.pop())
@@ -843,9 +848,7 @@ def unfuse_qkv_projections(self):
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
- def get_time_embed(
- self, sample: torch.Tensor, timestep: Union[torch.Tensor, float, int]
- ) -> Optional[torch.Tensor]:
+ def get_time_embed(self, sample: torch.Tensor, timestep: torch.Tensor | float | int) -> torch.Tensor | None:
timesteps = timestep
if not torch.is_tensor(timesteps):
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
@@ -870,7 +873,7 @@ def get_time_embed(
t_emb = t_emb.to(dtype=sample.dtype)
return t_emb
- def get_class_embed(self, sample: torch.Tensor, class_labels: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+ def get_class_embed(self, sample: torch.Tensor, class_labels: torch.Tensor | None) -> torch.Tensor | None:
class_emb = None
if self.class_embedding is not None:
if class_labels is None:
@@ -887,8 +890,8 @@ def get_class_embed(self, sample: torch.Tensor, class_labels: Optional[torch.Ten
return class_emb
def get_aug_embed(
- self, emb: torch.Tensor, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any]
- ) -> Optional[torch.Tensor]:
+ self, emb: torch.Tensor, encoder_hidden_states: torch.Tensor, added_cond_kwargs: dict[str, Any]
+ ) -> torch.Tensor | None:
aug_emb = None
if self.config.addition_embed_type == "text":
aug_emb = self.add_embedding(encoder_hidden_states)
@@ -939,7 +942,7 @@ def get_aug_embed(
return aug_emb
def process_encoder_hidden_states(
- self, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any]
+ self, encoder_hidden_states: torch.Tensor, added_cond_kwargs: dict[str, Any]
) -> torch.Tensor:
if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
@@ -974,22 +977,23 @@ def process_encoder_hidden_states(
encoder_hidden_states = (encoder_hidden_states, image_embeds)
return encoder_hidden_states
+ @apply_lora_scale("cross_attention_kwargs")
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
- mid_block_additional_residual: Optional[torch.Tensor] = None,
- down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ down_block_additional_residuals: tuple[torch.Tensor] | None = None,
+ mid_block_additional_residual: torch.Tensor | None = None,
+ down_intrablock_additional_residuals: tuple[torch.Tensor] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[UNet2DConditionOutput, Tuple]:
+ ) -> UNet2DConditionOutput | tuple:
r"""
The [`UNet2DConditionModel`] forward method.
@@ -1112,18 +1116,6 @@ def forward(
cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
# 3. down
- # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated
- # to the internal blocks and will raise deprecation warnings. this will be confusing for our users.
- if cross_attention_kwargs is not None:
- cross_attention_kwargs = cross_attention_kwargs.copy()
- lora_scale = cross_attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
-
- if USE_PEFT_BACKEND:
- # weight the lora layers by setting `lora_scale` for each PEFT layer
- scale_lora_layers(self, lora_scale)
-
is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
# using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
is_adapter = down_intrablock_additional_residuals is not None
@@ -1239,10 +1231,6 @@ def forward(
sample = self.conv_act(sample)
sample = self.conv_out(sample)
- if USE_PEFT_BACKEND:
- # remove `lora_scale` from each PEFT layer
- unscale_lora_layers(self, lora_scale)
-
if not return_dict:
return (sample,)
diff --git a/src/diffusers/models/unets/unet_2d_condition_flax.py b/src/diffusers/models/unets/unet_2d_condition_flax.py
index 8d9a309afbcc..a361026fc0ea 100644
--- a/src/diffusers/models/unets/unet_2d_condition_flax.py
+++ b/src/diffusers/models/unets/unet_2d_condition_flax.py
@@ -11,7 +11,6 @@
# 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.
-from typing import Dict, Optional, Tuple, Union
import flax
import flax.linen as nn
@@ -74,20 +73,20 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):
The number of channels in the input sample.
out_channels (`int`, *optional*, defaults to 4):
The number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`):
The tuple of downsample blocks to use.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("FlaxUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D")`):
+ up_block_types (`tuple[str]`, *optional*, defaults to `("FlaxUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D")`):
The tuple of upsample blocks to use.
mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`. If `None`, the mid block layer
is skipped.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2):
The number of layers per block.
- attention_head_dim (`int` or `Tuple[int]`, *optional*, defaults to 8):
+ attention_head_dim (`int` or `tuple[int]`, *optional*, defaults to 8):
The dimension of the attention heads.
- num_attention_heads (`int` or `Tuple[int]`, *optional*):
+ num_attention_heads (`int` or `tuple[int]`, *optional*):
The number of attention heads.
cross_attention_dim (`int`, *optional*, defaults to 768):
The dimension of the cross attention features.
@@ -106,19 +105,19 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):
sample_size: int = 32
in_channels: int = 4
out_channels: int = 4
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
)
- up_block_types: Tuple[str, ...] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
- mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn"
- only_cross_attention: Union[bool, Tuple[bool]] = False
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280)
+ up_block_types: tuple[str, ...] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
+ mid_block_type: str | None = "UNetMidBlock2DCrossAttn"
+ only_cross_attention: bool | tuple[bool] = False
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280)
layers_per_block: int = 2
- attention_head_dim: Union[int, Tuple[int, ...]] = 8
- num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None
+ attention_head_dim: int | tuple[int, ...] = 8
+ num_attention_heads: int | tuple[int, ...] | None = None
cross_attention_dim: int = 1280
dropout: float = 0.0
use_linear_projection: bool = False
@@ -127,11 +126,11 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):
freq_shift: int = 0
use_memory_efficient_attention: bool = False
split_head_dim: bool = False
- transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1
- addition_embed_type: Optional[str] = None
- addition_time_embed_dim: Optional[int] = None
+ transformer_layers_per_block: int | tuple[int, ...] = 1
+ addition_embed_type: str | None = None
+ addition_time_embed_dim: int | None = None
addition_embed_type_num_heads: int = 64
- projection_class_embeddings_input_dim: Optional[int] = None
+ projection_class_embeddings_input_dim: int | None = None
def init_weights(self, rng: jax.Array) -> FrozenDict:
# init input tensors
@@ -338,14 +337,14 @@ def setup(self) -> None:
def __call__(
self,
sample: jnp.ndarray,
- timesteps: Union[jnp.ndarray, float, int],
+ timesteps: jnp.ndarray | float | int,
encoder_hidden_states: jnp.ndarray,
- added_cond_kwargs: Optional[Union[Dict, FrozenDict]] = None,
- down_block_additional_residuals: Optional[Tuple[jnp.ndarray, ...]] = None,
- mid_block_additional_residual: Optional[jnp.ndarray] = None,
+ added_cond_kwargs: dict | FrozenDict | None = None,
+ down_block_additional_residuals: tuple[jnp.ndarray, ...] | None = None,
+ mid_block_additional_residual: jnp.ndarray | None = None,
return_dict: bool = True,
train: bool = False,
- ) -> Union[FlaxUNet2DConditionOutput, Tuple[jnp.ndarray]]:
+ ) -> FlaxUNet2DConditionOutput | tuple[jnp.ndarray]:
r"""
Args:
sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor
diff --git a/src/diffusers/models/unets/unet_3d_blocks.py b/src/diffusers/models/unets/unet_3d_blocks.py
index 53c0f4bae38b..b5f5aae342b4 100644
--- a/src/diffusers/models/unets/unet_3d_blocks.py
+++ b/src/diffusers/models/unets/unet_3d_blocks.py
@@ -12,7 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, Optional, Tuple, Union
+from __future__ import annotations
+
+from typing import Any
import torch
from torch import nn
@@ -89,9 +91,9 @@ def get_down_block(
resnet_eps: float,
resnet_act_fn: str,
num_attention_heads: int,
- resnet_groups: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
- downsample_padding: Optional[int] = None,
+ resnet_groups: int | None = None,
+ cross_attention_dim: int | None = None,
+ downsample_padding: int | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = True,
only_cross_attention: bool = False,
@@ -99,15 +101,10 @@ def get_down_block(
resnet_time_scale_shift: str = "default",
temporal_num_attention_heads: int = 8,
temporal_max_seq_length: int = 32,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
- temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
+ temporal_transformer_layers_per_block: int | tuple[int] = 1,
dropout: float = 0.0,
-) -> Union[
- "DownBlock3D",
- "CrossAttnDownBlock3D",
- "DownBlockSpatioTemporal",
- "CrossAttnDownBlockSpatioTemporal",
-]:
+) -> "DownBlock3D" | "CrossAttnDownBlock3D" | "DownBlockSpatioTemporal" | "CrossAttnDownBlockSpatioTemporal":
if down_block_type == "DownBlock3D":
return DownBlock3D(
num_layers=num_layers,
@@ -182,26 +179,21 @@ def get_up_block(
resnet_eps: float,
resnet_act_fn: str,
num_attention_heads: int,
- resolution_idx: Optional[int] = None,
- resnet_groups: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
+ resolution_idx: int | None = None,
+ resnet_groups: int | None = None,
+ cross_attention_dim: int | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = True,
only_cross_attention: bool = False,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
temporal_num_attention_heads: int = 8,
- temporal_cross_attention_dim: Optional[int] = None,
+ temporal_cross_attention_dim: int | None = None,
temporal_max_seq_length: int = 32,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
- temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
+ temporal_transformer_layers_per_block: int | tuple[int] = 1,
dropout: float = 0.0,
-) -> Union[
- "UpBlock3D",
- "CrossAttnUpBlock3D",
- "UpBlockSpatioTemporal",
- "CrossAttnUpBlockSpatioTemporal",
-]:
+) -> "UpBlock3D" | "CrossAttnUpBlock3D" | "UpBlockSpatioTemporal" | "CrossAttnUpBlockSpatioTemporal":
if up_block_type == "UpBlock3D":
return UpBlock3D(
num_layers=num_layers,
@@ -376,11 +368,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
num_frames: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
) -> torch.Tensor:
hidden_states = self.resnets[0](hidden_states, temb)
hidden_states = self.temp_convs[0](hidden_states, num_frames=num_frames)
@@ -509,12 +501,12 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
num_frames: int = 1,
- cross_attention_kwargs: Dict[str, Any] = None,
- ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ cross_attention_kwargs: dict[str, Any] = None,
+ ) -> torch.Tensor | tuple[torch.Tensor, ...]:
# TODO(Patrick, William) - attention mask is not used
output_states = ()
@@ -616,9 +608,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
num_frames: int = 1,
- ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ ) -> torch.Tensor | tuple[torch.Tensor, ...]:
output_states = ()
for resnet, temp_conv in zip(self.resnets, self.temp_convs):
@@ -658,7 +650,7 @@ def __init__(
use_linear_projection: bool = False,
only_cross_attention: bool = False,
upcast_attention: bool = False,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
):
super().__init__()
resnets = []
@@ -734,13 +726,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ upsample_size: int | None = None,
+ attention_mask: torch.Tensor | None = None,
num_frames: int = 1,
- cross_attention_kwargs: Dict[str, Any] = None,
+ cross_attention_kwargs: dict[str, Any] = None,
) -> torch.Tensor:
is_freeu_enabled = (
getattr(self, "s1", None)
@@ -809,7 +801,7 @@ def __init__(
resnet_pre_norm: bool = True,
output_scale_factor: float = 1.0,
add_upsample: bool = True,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
):
super().__init__()
resnets = []
@@ -856,9 +848,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ upsample_size: int | None = None,
num_frames: int = 1,
) -> torch.Tensor:
is_freeu_enabled = (
@@ -1015,7 +1007,7 @@ def __init__(
in_channels: int,
temb_channels: int,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
num_attention_heads: int = 1,
cross_attention_dim: int = 1280,
):
@@ -1067,9 +1059,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- image_only_indicator: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ image_only_indicator: torch.Tensor | None = None,
) -> torch.Tensor:
hidden_states = self.resnets[0](
hidden_states,
@@ -1142,9 +1134,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- image_only_indicator: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ temb: torch.Tensor | None = None,
+ image_only_indicator: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
output_states = ()
for resnet in self.resnets:
if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -1170,7 +1162,7 @@ def __init__(
out_channels: int,
temb_channels: int,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
num_attention_heads: int = 1,
cross_attention_dim: int = 1280,
add_downsample: bool = True,
@@ -1227,10 +1219,10 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- image_only_indicator: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ image_only_indicator: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
output_states = ()
blocks = list(zip(self.resnets, self.attentions))
@@ -1271,7 +1263,7 @@ def __init__(
prev_output_channel: int,
out_channels: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
num_layers: int = 1,
resnet_eps: float = 1e-6,
add_upsample: bool = True,
@@ -1305,10 +1297,10 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- image_only_indicator: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ image_only_indicator: torch.Tensor | None = None,
+ upsample_size: int | None = None,
) -> torch.Tensor:
for resnet in self.resnets:
# pop res hidden states
@@ -1336,9 +1328,9 @@ def __init__(
out_channels: int,
prev_output_channel: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
num_attention_heads: int = 1,
cross_attention_dim: int = 1280,
@@ -1390,11 +1382,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- image_only_indicator: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ image_only_indicator: torch.Tensor | None = None,
+ upsample_size: int | None = None,
) -> torch.Tensor:
for resnet, attn in zip(self.resnets, self.attentions):
# pop res hidden states
diff --git a/src/diffusers/models/unets/unet_3d_condition.py b/src/diffusers/models/unets/unet_3d_condition.py
index 26dc50f84acd..5006e48feb46 100644
--- a/src/diffusers/models/unets/unet_3d_condition.py
+++ b/src/diffusers/models/unets/unet_3d_condition.py
@@ -14,7 +14,7 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -67,15 +67,15 @@ class UNet3DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DCondit
for all models (such as downloading or saving).
Parameters:
- sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+ sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`):
Height and width of input/output sample.
in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.
out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")`):
The tuple of downsample blocks to use.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")`):
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")`):
The tuple of upsample blocks to use.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
@@ -97,32 +97,32 @@ class UNet3DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DCondit
@register_to_config
def __init__(
self,
- sample_size: Optional[int] = None,
+ sample_size: int | None = None,
in_channels: int = 4,
out_channels: int = 4,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlock3D",
"CrossAttnDownBlock3D",
"CrossAttnDownBlock3D",
"DownBlock3D",
),
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"UpBlock3D",
"CrossAttnUpBlock3D",
"CrossAttnUpBlock3D",
"CrossAttnUpBlock3D",
),
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
layers_per_block: int = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
act_fn: str = "silu",
- norm_num_groups: Optional[int] = 32,
+ norm_num_groups: int | None = 32,
norm_eps: float = 1e-5,
cross_attention_dim: int = 1024,
- attention_head_dim: Union[int, Tuple[int]] = 64,
- num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
- time_cond_proj_dim: Optional[int] = None,
+ attention_head_dim: int | tuple[int] = 64,
+ num_attention_heads: int | tuple[int] | None = None,
+ time_cond_proj_dim: int | None = None,
):
super().__init__()
@@ -286,7 +286,7 @@ def __init__(
)
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice
- def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None:
+ def set_attention_slice(self, slice_size: str | int | list[int]) -> None:
r"""
Enable sliced attention computation.
@@ -340,7 +340,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
# Recursively walk through all the children.
# Any children which exposes the set_attention_slice method
# gets the message
- def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]):
if hasattr(module, "set_attention_slice"):
module.set_attention_slice(slice_size.pop())
@@ -351,7 +351,7 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i
for module in self.children():
fn_recursive_set_attention_slice(module, reversed_slice_size)
- def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+ def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -476,16 +476,16 @@ def unfuse_qkv_projections(self):
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
- mid_block_additional_residual: Optional[torch.Tensor] = None,
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ down_block_additional_residuals: tuple[torch.Tensor] | None = None,
+ mid_block_additional_residual: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[UNet3DConditionOutput, Tuple[torch.Tensor]]:
+ ) -> UNet3DConditionOutput | tuple[torch.Tensor]:
r"""
The [`UNet3DConditionModel`] forward method.
diff --git a/src/diffusers/models/unets/unet_i2vgen_xl.py b/src/diffusers/models/unets/unet_i2vgen_xl.py
index 0ada264417dd..5c3cfe91d5bd 100644
--- a/src/diffusers/models/unets/unet_i2vgen_xl.py
+++ b/src/diffusers/models/unets/unet_i2vgen_xl.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -51,7 +51,7 @@ def __init__(
attention_head_dim: int,
activation_fn: str = "geglu",
upcast_attention: bool = False,
- ff_inner_dim: Optional[int] = None,
+ ff_inner_dim: int | None = None,
dropout: int = 0.0,
):
super().__init__()
@@ -101,15 +101,15 @@ class I2VGenXLUNet(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoade
for all models (such as downloading or saving).
Parameters:
- sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+ sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`):
Height and width of input/output sample.
in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.
out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
The tuple of downsample blocks to use.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
The tuple of upsample blocks to use.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
@@ -124,27 +124,27 @@ class I2VGenXLUNet(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoade
@register_to_config
def __init__(
self,
- sample_size: Optional[int] = None,
+ sample_size: int | None = None,
in_channels: int = 4,
out_channels: int = 4,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlock3D",
"CrossAttnDownBlock3D",
"CrossAttnDownBlock3D",
"DownBlock3D",
),
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"UpBlock3D",
"CrossAttnUpBlock3D",
"CrossAttnUpBlock3D",
"CrossAttnUpBlock3D",
),
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
layers_per_block: int = 2,
- norm_num_groups: Optional[int] = 32,
+ norm_num_groups: int | None = 32,
cross_attention_dim: int = 1024,
- attention_head_dim: Union[int, Tuple[int]] = 64,
- num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+ attention_head_dim: int | tuple[int] = 64,
+ num_attention_heads: int | tuple[int] | None = None,
):
super().__init__()
@@ -313,7 +313,7 @@ def __init__(
self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1)
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
- def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+ def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -439,15 +439,15 @@ def unfuse_qkv_projections(self):
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
fps: torch.Tensor,
image_latents: torch.Tensor,
- image_embeddings: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ image_embeddings: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
- ) -> Union[UNet3DConditionOutput, Tuple[torch.Tensor]]:
+ ) -> UNet3DConditionOutput | tuple[torch.Tensor]:
r"""
The [`I2VGenXLUNet`] forward method.
diff --git a/src/diffusers/models/unets/unet_kandinsky3.py b/src/diffusers/models/unets/unet_kandinsky3.py
index 13f4641a4c50..6fa68b42ee30 100644
--- a/src/diffusers/models/unets/unet_kandinsky3.py
+++ b/src/diffusers/models/unets/unet_kandinsky3.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Tuple, Union
import torch
from torch import nn
@@ -54,9 +53,9 @@ def __init__(
time_embedding_dim: int = 1536,
groups: int = 32,
attention_head_dim: int = 64,
- layers_per_block: Union[int, Tuple[int]] = 3,
- block_out_channels: Tuple[int, ...] = (384, 768, 1536, 3072),
- cross_attention_dim: Union[int, Tuple[int]] = 4096,
+ layers_per_block: int | tuple[int] = 3,
+ block_out_channels: tuple[int, ...] = (384, 768, 1536, 3072),
+ cross_attention_dim: int | tuple[int] = 4096,
encoder_hid_dim: int = 4096,
):
super().__init__()
diff --git a/src/diffusers/models/unets/unet_motion_model.py b/src/diffusers/models/unets/unet_motion_model.py
index 5a93541501d3..97452eff05aa 100644
--- a/src/diffusers/models/unets/unet_motion_model.py
+++ b/src/diffusers/models/unets/unet_motion_model.py
@@ -13,7 +13,7 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -21,7 +21,7 @@
from ...configuration_utils import ConfigMixin, FrozenDict, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin, UNet2DConditionLoadersMixin
-from ...utils import BaseOutput, deprecate, logging
+from ...utils import BaseOutput, apply_lora_scale, deprecate, logging
from ...utils.torch_utils import apply_freeu
from ..attention import AttentionMixin, BasicTransformerBlock
from ..attention_processor import (
@@ -93,19 +93,19 @@ def __init__(
self,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
- out_channels: Optional[int] = None,
+ in_channels: int | None = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
norm_num_groups: int = 32,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
- sample_size: Optional[int] = None,
+ sample_size: int | None = None,
activation_fn: str = "geglu",
norm_elementwise_affine: bool = True,
double_self_attention: bool = True,
- positional_embeddings: Optional[str] = None,
- num_positional_embeddings: Optional[int] = None,
+ positional_embeddings: str | None = None,
+ num_positional_embeddings: int | None = None,
):
super().__init__()
self.num_attention_heads = num_attention_heads
@@ -142,11 +142,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- encoder_hidden_states: Optional[torch.LongTensor] = None,
- timestep: Optional[torch.LongTensor] = None,
- class_labels: Optional[torch.LongTensor] = None,
+ encoder_hidden_states: torch.LongTensor | None = None,
+ timestep: torch.LongTensor | None = None,
+ class_labels: torch.LongTensor | None = None,
num_frames: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
) -> torch.Tensor:
"""
The [`AnimateDiffTransformer3D`] forward method.
@@ -227,10 +227,10 @@ def __init__(
output_scale_factor: float = 1.0,
add_downsample: bool = True,
downsample_padding: int = 1,
- temporal_num_attention_heads: Union[int, Tuple[int]] = 1,
- temporal_cross_attention_dim: Optional[int] = None,
+ temporal_num_attention_heads: int | tuple[int] = 1,
+ temporal_cross_attention_dim: int | None = None,
temporal_max_seq_length: int = 32,
- temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ temporal_transformer_layers_per_block: int | tuple[int] = 1,
temporal_double_self_attention: bool = True,
):
super().__init__()
@@ -308,11 +308,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
num_frames: int = 1,
*args,
**kwargs,
- ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ ) -> torch.Tensor | tuple[torch.Tensor, ...]:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -347,7 +347,7 @@ def __init__(
temb_channels: int,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
@@ -363,10 +363,10 @@ def __init__(
only_cross_attention: bool = False,
upcast_attention: bool = False,
attention_type: str = "default",
- temporal_cross_attention_dim: Optional[int] = None,
+ temporal_cross_attention_dim: int | None = None,
temporal_num_attention_heads: int = 8,
temporal_max_seq_length: int = 32,
- temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ temporal_transformer_layers_per_block: int | tuple[int] = 1,
temporal_double_self_attention: bool = True,
):
super().__init__()
@@ -477,13 +477,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
num_frames: int = 1,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- additional_residuals: Optional[torch.Tensor] = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ additional_residuals: torch.Tensor | None = None,
):
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
@@ -531,10 +531,10 @@ def __init__(
out_channels: int,
prev_output_channel: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
@@ -549,10 +549,10 @@ def __init__(
only_cross_attention: bool = False,
upcast_attention: bool = False,
attention_type: str = "default",
- temporal_cross_attention_dim: Optional[int] = None,
+ temporal_cross_attention_dim: int | None = None,
temporal_num_attention_heads: int = 8,
temporal_max_seq_length: int = 32,
- temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ temporal_transformer_layers_per_block: int | tuple[int] = 1,
):
super().__init__()
resnets = []
@@ -653,13 +653,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ upsample_size: int | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
num_frames: int = 1,
) -> torch.Tensor:
if cross_attention_kwargs is not None:
@@ -723,7 +723,7 @@ def __init__(
prev_output_channel: int,
out_channels: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -733,10 +733,10 @@ def __init__(
resnet_pre_norm: bool = True,
output_scale_factor: float = 1.0,
add_upsample: bool = True,
- temporal_cross_attention_dim: Optional[int] = None,
+ temporal_cross_attention_dim: int | None = None,
temporal_num_attention_heads: int = 8,
temporal_max_seq_length: int = 32,
- temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ temporal_transformer_layers_per_block: int | tuple[int] = 1,
):
super().__init__()
resnets = []
@@ -798,8 +798,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
upsample_size=None,
num_frames: int = 1,
*args,
@@ -858,7 +858,7 @@ def __init__(
temb_channels: int,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
@@ -872,9 +872,9 @@ def __init__(
upcast_attention: bool = False,
attention_type: str = "default",
temporal_num_attention_heads: int = 1,
- temporal_cross_attention_dim: Optional[int] = None,
+ temporal_cross_attention_dim: int | None = None,
temporal_max_seq_length: int = 32,
- temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ temporal_transformer_layers_per_block: int | tuple[int] = 1,
):
super().__init__()
@@ -980,11 +980,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
num_frames: int = 1,
) -> torch.Tensor:
if cross_attention_kwargs is not None:
@@ -1021,10 +1021,10 @@ def __init__(
self,
in_channels: int,
layers_per_block: int = 2,
- transformer_layers_per_block: Union[int, Tuple[int]] = 8,
- num_attention_heads: Union[int, Tuple[int]] = 8,
+ transformer_layers_per_block: int | tuple[int] = 8,
+ num_attention_heads: int | tuple[int] = 8,
attention_bias: bool = False,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
activation_fn: str = "geglu",
norm_num_groups: int = 32,
max_seq_length: int = 32,
@@ -1061,31 +1061,31 @@ class MotionAdapter(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@register_to_config
def __init__(
self,
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
- motion_layers_per_block: Union[int, Tuple[int]] = 2,
- motion_transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple[int]]] = 1,
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
+ motion_layers_per_block: int | tuple[int] = 2,
+ motion_transformer_layers_per_block: int | tuple[int] | tuple[tuple[int]] = 1,
motion_mid_block_layers_per_block: int = 1,
- motion_transformer_layers_per_mid_block: Union[int, Tuple[int]] = 1,
- motion_num_attention_heads: Union[int, Tuple[int]] = 8,
+ motion_transformer_layers_per_mid_block: int | tuple[int] = 1,
+ motion_num_attention_heads: int | tuple[int] = 8,
motion_norm_num_groups: int = 32,
motion_max_seq_length: int = 32,
use_motion_mid_block: bool = True,
- conv_in_channels: Optional[int] = None,
+ conv_in_channels: int | None = None,
):
"""Container to store AnimateDiff Motion Modules
Args:
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each UNet block.
- motion_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 2):
+ motion_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 2):
The number of motion layers per UNet block.
- motion_transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple[int]]`, *optional*, defaults to 1):
+ motion_transformer_layers_per_block (`int`, `tuple[int]`, or `tuple[tuple[int]]`, *optional*, defaults to 1):
The number of transformer layers to use in each motion layer in each block.
motion_mid_block_layers_per_block (`int`, *optional*, defaults to 1):
The number of motion layers in the middle UNet block.
- motion_transformer_layers_per_mid_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
+ motion_transformer_layers_per_mid_block (`int` or `tuple[int]`, *optional*, defaults to 1):
The number of transformer layers to use in each motion layer in the middle block.
- motion_num_attention_heads (`int` or `Tuple[int]`, *optional*, defaults to 8):
+ motion_num_attention_heads (`int` or `tuple[int]`, *optional*, defaults to 8):
The number of heads to use in each attention layer of the motion module.
motion_norm_num_groups (`int`, *optional*, defaults to 32):
The number of groups to use in each group normalization layer of the motion module.
@@ -1209,48 +1209,48 @@ class UNetMotionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLo
@register_to_config
def __init__(
self,
- sample_size: Optional[int] = None,
+ sample_size: int | None = None,
in_channels: int = 4,
out_channels: int = 4,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlockMotion",
"CrossAttnDownBlockMotion",
"CrossAttnDownBlockMotion",
"DownBlockMotion",
),
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"UpBlockMotion",
"CrossAttnUpBlockMotion",
"CrossAttnUpBlockMotion",
"CrossAttnUpBlockMotion",
),
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
- layers_per_block: Union[int, Tuple[int]] = 2,
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
+ layers_per_block: int | tuple[int] = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
act_fn: str = "silu",
norm_num_groups: int = 32,
norm_eps: float = 1e-5,
cross_attention_dim: int = 1280,
- transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
- reverse_transformer_layers_per_block: Optional[Union[int, Tuple[int], Tuple[Tuple]]] = None,
- temporal_transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
- reverse_temporal_transformer_layers_per_block: Optional[Union[int, Tuple[int], Tuple[Tuple]]] = None,
- transformer_layers_per_mid_block: Optional[Union[int, Tuple[int]]] = None,
- temporal_transformer_layers_per_mid_block: Optional[Union[int, Tuple[int]]] = 1,
+ transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1,
+ reverse_transformer_layers_per_block: int | tuple[int] | tuple[tuple] | None = None,
+ temporal_transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1,
+ reverse_temporal_transformer_layers_per_block: int | tuple[int] | tuple[tuple] | None = None,
+ transformer_layers_per_mid_block: int | tuple[int] | None = None,
+ temporal_transformer_layers_per_mid_block: int | tuple[int] | None = 1,
use_linear_projection: bool = False,
- num_attention_heads: Union[int, Tuple[int, ...]] = 8,
+ num_attention_heads: int | tuple[int, ...] = 8,
motion_max_seq_length: int = 32,
- motion_num_attention_heads: Union[int, Tuple[int, ...]] = 8,
- reverse_motion_num_attention_heads: Optional[Union[int, Tuple[int, ...], Tuple[Tuple[int, ...], ...]]] = None,
+ motion_num_attention_heads: int | tuple[int, ...] = 8,
+ reverse_motion_num_attention_heads: int | tuple[int, ...] | tuple[tuple[int, ...], ...] | None = None,
use_motion_mid_block: bool = True,
mid_block_layers: int = 1,
- encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
- addition_time_embed_dim: Optional[int] = None,
- projection_class_embeddings_input_dim: Optional[int] = None,
- time_cond_proj_dim: Optional[int] = None,
+ encoder_hid_dim: int | None = None,
+ encoder_hid_dim_type: str | None = None,
+ addition_embed_type: str | None = None,
+ addition_time_embed_dim: int | None = None,
+ projection_class_embeddings_input_dim: int | None = None,
+ time_cond_proj_dim: int | None = None,
):
super().__init__()
@@ -1534,7 +1534,7 @@ def __init__(
def from_unet2d(
cls,
unet: UNet2DConditionModel,
- motion_adapter: Optional[MotionAdapter] = None,
+ motion_adapter: MotionAdapter | None = None,
load_weights: bool = True,
):
has_motion_adapter = motion_adapter is not None
@@ -1708,7 +1708,7 @@ def freeze_unet2d_params(self) -> None:
for param in motion_modules.parameters():
param.requires_grad = True
- def load_motion_modules(self, motion_adapter: Optional[MotionAdapter]) -> None:
+ def load_motion_modules(self, motion_adapter: MotionAdapter | None) -> None:
for i, down_block in enumerate(motion_adapter.down_blocks):
self.down_blocks[i].motion_modules.load_state_dict(down_block.motion_modules.state_dict())
for i, up_block in enumerate(motion_adapter.up_blocks):
@@ -1723,7 +1723,7 @@ def save_motion_modules(
save_directory: str,
is_main_process: bool = True,
safe_serialization: bool = True,
- variant: Optional[str] = None,
+ variant: str | None = None,
push_to_hub: bool = False,
**kwargs,
) -> None:
@@ -1753,7 +1753,7 @@ def save_motion_modules(
**kwargs,
)
- def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+ def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -1875,19 +1875,20 @@ def unfuse_qkv_projections(self):
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
+ @apply_lora_scale("cross_attention_kwargs")
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
- mid_block_additional_residual: Optional[torch.Tensor] = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ down_block_additional_residuals: tuple[torch.Tensor] | None = None,
+ mid_block_additional_residual: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[UNetMotionOutput, Tuple[torch.Tensor]]:
+ ) -> UNetMotionOutput | tuple[torch.Tensor]:
r"""
The [`UNetMotionModel`] forward method.
diff --git a/src/diffusers/models/unets/unet_spatio_temporal_condition.py b/src/diffusers/models/unets/unet_spatio_temporal_condition.py
index c0cd5fbdd489..eddeb9826b0c 100644
--- a/src/diffusers/models/unets/unet_spatio_temporal_condition.py
+++ b/src/diffusers/models/unets/unet_spatio_temporal_condition.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -39,29 +38,29 @@ class UNetSpatioTemporalConditionModel(ModelMixin, AttentionMixin, ConfigMixin,
for all models (such as downloading or saving).
Parameters:
- sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+ sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`):
Height and width of input/output sample.
in_channels (`int`, *optional*, defaults to 8): Number of channels in the input sample.
out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`):
The tuple of downsample blocks to use.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`):
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`):
The tuple of upsample blocks to use.
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
addition_time_embed_dim: (`int`, defaults to 256):
Dimension to to encode the additional time ids.
projection_class_embeddings_input_dim (`int`, defaults to 768):
The dimension of the projection of encoded `added_time_ids`.
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
- cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+ cross_attention_dim (`int` or `tuple[int]`, *optional*, defaults to 1280):
The dimension of the cross attention features.
- transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+ transformer_layers_per_block (`int`, `tuple[int]`, or `tuple[tuple]` , *optional*, defaults to 1):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
[`~models.unets.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`],
[`~models.unets.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`],
[`~models.unets.unet_3d_blocks.UNetMidBlockSpatioTemporal`].
- num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`):
+ num_attention_heads (`int`, `tuple[int]`, defaults to `(5, 10, 10, 20)`):
The number of attention heads.
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
"""
@@ -71,28 +70,28 @@ class UNetSpatioTemporalConditionModel(ModelMixin, AttentionMixin, ConfigMixin,
@register_to_config
def __init__(
self,
- sample_size: Optional[int] = None,
+ sample_size: int | None = None,
in_channels: int = 8,
out_channels: int = 4,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str, ...] = (
"CrossAttnDownBlockSpatioTemporal",
"CrossAttnDownBlockSpatioTemporal",
"CrossAttnDownBlockSpatioTemporal",
"DownBlockSpatioTemporal",
),
- up_block_types: Tuple[str, ...] = (
+ up_block_types: tuple[str, ...] = (
"UpBlockSpatioTemporal",
"CrossAttnUpBlockSpatioTemporal",
"CrossAttnUpBlockSpatioTemporal",
"CrossAttnUpBlockSpatioTemporal",
),
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+ block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280),
addition_time_embed_dim: int = 256,
projection_class_embeddings_input_dim: int = 768,
- layers_per_block: Union[int, Tuple[int]] = 2,
- cross_attention_dim: Union[int, Tuple[int]] = 1024,
- transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
- num_attention_heads: Union[int, Tuple[int, ...]] = (5, 10, 20, 20),
+ layers_per_block: int | tuple[int] = 2,
+ cross_attention_dim: int | tuple[int] = 1024,
+ transformer_layers_per_block: int | tuple[int, tuple[tuple]] = 1,
+ num_attention_heads: int | tuple[int, ...] = (5, 10, 20, 20),
num_frames: int = 25,
):
super().__init__()
@@ -260,7 +259,7 @@ def set_default_attn_processor(self):
self.set_attn_processor(processor)
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
- def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+ def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -292,11 +291,11 @@ def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
added_time_ids: torch.Tensor,
return_dict: bool = True,
- ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]:
+ ) -> UNetSpatioTemporalConditionOutput | tuple:
r"""
The [`UNetSpatioTemporalConditionModel`] forward method.
diff --git a/src/diffusers/models/unets/unet_stable_cascade.py b/src/diffusers/models/unets/unet_stable_cascade.py
index 23d358c1bf51..5a6f24ab794b 100644
--- a/src/diffusers/models/unets/unet_stable_cascade.py
+++ b/src/diffusers/models/unets/unet_stable_cascade.py
@@ -14,7 +14,6 @@
import math
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -145,30 +144,30 @@ def __init__(
timestep_ratio_embedding_dim: int = 64,
patch_size: int = 1,
conditioning_dim: int = 2048,
- block_out_channels: Tuple[int, ...] = (2048, 2048),
- num_attention_heads: Tuple[int, ...] = (32, 32),
- down_num_layers_per_block: Tuple[int, ...] = (8, 24),
- up_num_layers_per_block: Tuple[int, ...] = (24, 8),
- down_blocks_repeat_mappers: Optional[Tuple[int]] = (
+ block_out_channels: tuple[int, ...] = (2048, 2048),
+ num_attention_heads: tuple[int, ...] = (32, 32),
+ down_num_layers_per_block: tuple[int, ...] = (8, 24),
+ up_num_layers_per_block: tuple[int, ...] = (24, 8),
+ down_blocks_repeat_mappers: tuple[int] | None = (
1,
1,
),
- up_blocks_repeat_mappers: Optional[Tuple[int]] = (1, 1),
- block_types_per_layer: Tuple[Tuple[str]] = (
+ up_blocks_repeat_mappers: tuple[int] | None = (1, 1),
+ block_types_per_layer: tuple[tuple[str]] = (
("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"),
("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"),
),
- clip_text_in_channels: Optional[int] = None,
+ clip_text_in_channels: int | None = None,
clip_text_pooled_in_channels=1280,
- clip_image_in_channels: Optional[int] = None,
+ clip_image_in_channels: int | None = None,
clip_seq=4,
- effnet_in_channels: Optional[int] = None,
- pixel_mapper_in_channels: Optional[int] = None,
+ effnet_in_channels: int | None = None,
+ pixel_mapper_in_channels: int | None = None,
kernel_size=3,
- dropout: Union[float, Tuple[float]] = (0.1, 0.1),
- self_attn: Union[bool, Tuple[bool]] = True,
- timestep_conditioning_type: Tuple[str, ...] = ("sca", "crp"),
- switch_level: Optional[Tuple[bool]] = None,
+ dropout: float | tuple[float] = (0.1, 0.1),
+ self_attn: bool | tuple[bool] = True,
+ timestep_conditioning_type: tuple[str, ...] = ("sca", "crp"),
+ switch_level: tuple[bool] | None = None,
):
"""
@@ -183,20 +182,20 @@ def __init__(
Patch size to use for pixel unshuffling layer
conditioning_dim (`int`, defaults to 2048):
Dimension of the image and text conditional embedding.
- block_out_channels (Tuple[int], defaults to (2048, 2048)):
- Tuple of output channels for each block.
- num_attention_heads (Tuple[int], defaults to (32, 32)):
+ block_out_channels (tuple[int], defaults to (2048, 2048)):
+ tuple of output channels for each block.
+ num_attention_heads (tuple[int], defaults to (32, 32)):
Number of attention heads in each attention block. Set to -1 to if block types in a layer do not have
attention.
- down_num_layers_per_block (Tuple[int], defaults to [8, 24]):
+ down_num_layers_per_block (tuple[int], defaults to [8, 24]):
Number of layers in each down block.
- up_num_layers_per_block (Tuple[int], defaults to [24, 8]):
+ up_num_layers_per_block (tuple[int], defaults to [24, 8]):
Number of layers in each up block.
- down_blocks_repeat_mappers (Tuple[int], optional, defaults to [1, 1]):
+ down_blocks_repeat_mappers (tuple[int], optional, defaults to [1, 1]):
Number of 1x1 Convolutional layers to repeat in each down block.
- up_blocks_repeat_mappers (Tuple[int], optional, defaults to [1, 1]):
+ up_blocks_repeat_mappers (tuple[int], optional, defaults to [1, 1]):
Number of 1x1 Convolutional layers to repeat in each up block.
- block_types_per_layer (Tuple[Tuple[str]], optional,
+ block_types_per_layer (tuple[tuple[str]], optional,
defaults to (
("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"), ("SDCascadeResBlock",
"SDCascadeTimestepBlock", "SDCascadeAttnBlock")
@@ -214,14 +213,14 @@ def __init__(
Number of input channels for pixel mapper conditioning.
kernel_size (`int`, *optional*, defaults to 3):
Kernel size to use in the block convolutional layers.
- dropout (Tuple[float], *optional*, defaults to (0.1, 0.1)):
+ dropout (tuple[float], *optional*, defaults to (0.1, 0.1)):
Dropout to use per block.
- self_attn (Union[bool, Tuple[bool]]):
- Tuple of booleans that determine whether to use self attention in a block or not.
- timestep_conditioning_type (Tuple[str], defaults to ("sca", "crp")):
+ self_attn (bool | tuple[bool]):
+ tuple of booleans that determine whether to use self attention in a block or not.
+ timestep_conditioning_type (tuple[str], defaults to ("sca", "crp")):
Timestep conditioning type.
- switch_level (Optional[Tuple[bool]], *optional*, defaults to `None`):
- Tuple that indicates whether upsampling or downsampling should be applied in a block
+ switch_level (tuple[bool] | None, *optional*, defaults to `None`):
+ tuple that indicates whether upsampling or downsampling should be applied in a block
"""
super().__init__()
diff --git a/src/diffusers/models/unets/uvit_2d.py b/src/diffusers/models/unets/uvit_2d.py
index 4c99ef88ca19..836d41a7f946 100644
--- a/src/diffusers/models/unets/uvit_2d.py
+++ b/src/diffusers/models/unets/uvit_2d.py
@@ -21,6 +21,7 @@
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
+from ...utils import apply_lora_scale
from ..attention import AttentionMixin, BasicTransformerBlock, SkipFFTransformerBlock
from ..attention_processor import (
ADDED_KV_ATTENTION_PROCESSORS,
@@ -146,6 +147,7 @@ def __init__(
self.gradient_checkpointing = False
+ @apply_lora_scale("cross_attention_kwargs")
def forward(self, input_ids, encoder_hidden_states, pooled_text_emb, micro_conds, cross_attention_kwargs=None):
encoder_hidden_states = self.encoder_proj(encoder_hidden_states)
encoder_hidden_states = self.encoder_proj_layer_norm(encoder_hidden_states)
diff --git a/src/diffusers/models/upsampling.py b/src/diffusers/models/upsampling.py
index 8a47c69f1264..cd3986287303 100644
--- a/src/diffusers/models/upsampling.py
+++ b/src/diffusers/models/upsampling.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple
-
import torch
import torch.nn as nn
import torch.nn.functional as F
@@ -44,7 +42,7 @@ def __init__(
channels: int,
use_conv: bool = False,
use_conv_transpose: bool = False,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
name: str = "conv",
):
super().__init__()
@@ -94,9 +92,9 @@ def __init__(
channels: int,
use_conv: bool = False,
use_conv_transpose: bool = False,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
name: str = "conv",
- kernel_size: Optional[int] = None,
+ kernel_size: int | None = None,
padding=1,
norm_type=None,
eps=None,
@@ -139,7 +137,7 @@ def __init__(
else:
self.Conv2d_0 = conv
- def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None, *args, **kwargs) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, output_size: int | None = None, *args, **kwargs) -> torch.Tensor:
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
@@ -208,10 +206,10 @@ class FirUpsample2D(nn.Module):
def __init__(
self,
- channels: Optional[int] = None,
- out_channels: Optional[int] = None,
+ channels: int | None = None,
+ out_channels: int | None = None,
use_conv: bool = False,
- fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1),
+ fir_kernel: tuple[int, int, int, int] = (1, 3, 3, 1),
):
super().__init__()
out_channels = out_channels if out_channels else channels
@@ -224,8 +222,8 @@ def __init__(
def _upsample_2d(
self,
hidden_states: torch.Tensor,
- weight: Optional[torch.Tensor] = None,
- kernel: Optional[torch.Tensor] = None,
+ weight: torch.Tensor | None = None,
+ kernel: torch.Tensor | None = None,
factor: int = 2,
gain: float = 1,
) -> torch.Tensor:
@@ -425,7 +423,7 @@ def upfirdn2d_native(
kernel: torch.Tensor,
up: int = 1,
down: int = 1,
- pad: Tuple[int, int] = (0, 0),
+ pad: tuple[int, int] = (0, 0),
) -> torch.Tensor:
up_x = up_y = up
down_x = down_y = down
@@ -472,7 +470,7 @@ def upfirdn2d_native(
def upsample_2d(
hidden_states: torch.Tensor,
- kernel: Optional[torch.Tensor] = None,
+ kernel: torch.Tensor | None = None,
factor: int = 2,
gain: float = 1,
) -> torch.Tensor:
diff --git a/src/diffusers/models/vae_flax.py b/src/diffusers/models/vae_flax.py
index 5aad386a89e8..c7042840e4e0 100644
--- a/src/diffusers/models/vae_flax.py
+++ b/src/diffusers/models/vae_flax.py
@@ -16,7 +16,6 @@
import math
from functools import partial
-from typing import Tuple
import flax
import flax.linen as nn
@@ -514,10 +513,10 @@ class FlaxEncoder(nn.Module):
Input channels
out_channels (:obj:`int`, *optional*, defaults to 3):
Output channels
- down_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`):
+ down_block_types (:obj:`tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`):
DownEncoder block type
- block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`):
- Tuple containing the number of output channels for each block
+ block_out_channels (:obj:`tuple[str]`, *optional*, defaults to `(64,)`):
+ tuple[ containing the number of output channels for each block
layers_per_block (:obj:`int`, *optional*, defaults to `2`):
Number of Resnet layer for each block
norm_num_groups (:obj:`int`, *optional*, defaults to `32`):
@@ -532,8 +531,8 @@ class FlaxEncoder(nn.Module):
in_channels: int = 3
out_channels: int = 3
- down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",)
- block_out_channels: Tuple[int, ...] = (64,)
+ down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",)
+ block_out_channels: tuple[int, ...] = (64,)
layers_per_block: int = 2
norm_num_groups: int = 32
act_fn: str = "silu"
@@ -632,10 +631,10 @@ class FlaxDecoder(nn.Module):
Input channels
out_channels (:obj:`int`, *optional*, defaults to 3):
Output channels
- up_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`):
+ up_block_types (:obj:`tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`):
UpDecoder block type
- block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`):
- Tuple containing the number of output channels for each block
+ block_out_channels (:obj:`tuple[str]`, *optional*, defaults to `(64,)`):
+ tuple[ containing the number of output channels for each block
layers_per_block (:obj:`int`, *optional*, defaults to `2`):
Number of Resnet layer for each block
norm_num_groups (:obj:`int`, *optional*, defaults to `32`):
@@ -650,8 +649,8 @@ class FlaxDecoder(nn.Module):
in_channels: int = 3
out_channels: int = 3
- up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",)
- block_out_channels: Tuple[int, ...] = (64,)
+ up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",)
+ block_out_channels: tuple[int, ...] = (64,)
layers_per_block: int = 2
norm_num_groups: int = 32
act_fn: str = "silu"
@@ -794,12 +793,12 @@ class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin):
Number of channels in the input image.
out_channels (`int`, *optional*, defaults to 3):
Number of channels in the output.
- down_block_types (`Tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`):
- Tuple of downsample block types.
- up_block_types (`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`):
- Tuple of upsample block types.
- block_out_channels (`Tuple[str]`, *optional*, defaults to `(64,)`):
- Tuple of block output channels.
+ down_block_types (`tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`):
+ tuple[ of downsample block types.
+ up_block_types (`tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`):
+ tuple[ of upsample block types.
+ block_out_channels (`tuple[str]`, *optional*, defaults to `(64,)`):
+ tuple[ of block output channels.
layers_per_block (`int`, *optional*, defaults to `2`):
Number of ResNet layer for each block.
act_fn (`str`, *optional*, defaults to `silu`):
@@ -823,9 +822,9 @@ class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin):
in_channels: int = 3
out_channels: int = 3
- down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",)
- up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",)
- block_out_channels: Tuple[int, ...] = (64,)
+ down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",)
+ up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",)
+ block_out_channels: tuple[int, ...] = (64,)
layers_per_block: int = 1
act_fn: str = "silu"
latent_channels: int = 4
diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py
index 94b87c61c234..c9bebd8644f7 100644
--- a/src/diffusers/modular_pipelines/__init__.py
+++ b/src/diffusers/modular_pipelines/__init__.py
@@ -33,6 +33,7 @@
"ModularPipeline",
"AutoPipelineBlocks",
"SequentialPipelineBlocks",
+ "ConditionalPipelineBlocks",
"LoopSequentialPipelineBlocks",
"PipelineState",
"BlockState",
@@ -105,6 +106,7 @@
from .modular_pipeline import (
AutoPipelineBlocks,
BlockState,
+ ConditionalPipelineBlocks,
LoopSequentialPipelineBlocks,
ModularPipeline,
ModularPipelineBlocks,
diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py
index 4a7ea8502c86..e018381ba859 100644
--- a/src/diffusers/modular_pipelines/components_manager.py
+++ b/src/diffusers/modular_pipelines/components_manager.py
@@ -12,11 +12,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import copy
import time
from collections import OrderedDict
from itertools import combinations
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
import torch
@@ -53,9 +55,9 @@ class CustomOffloadHook(ModelHook):
def __init__(
self,
- execution_device: Optional[Union[str, int, torch.device]] = None,
- other_hooks: Optional[List["UserCustomOffloadHook"]] = None,
- offload_strategy: Optional["AutoOffloadStrategy"] = None,
+ execution_device: str | int | torch.device | None = None,
+ other_hooks: list["UserCustomOffloadHook"] | None = None,
+ offload_strategy: "AutoOffloadStrategy" | None = None,
):
self.execution_device = execution_device if execution_device is not None else PartialState().default_device
self.other_hooks = other_hooks
@@ -135,8 +137,8 @@ def add_other_hook(self, hook: "UserCustomOffloadHook"):
def custom_offload_with_hook(
model_id: str,
model: torch.nn.Module,
- execution_device: Union[str, int, torch.device] = None,
- offload_strategy: Optional["AutoOffloadStrategy"] = None,
+ execution_device: str | int | torch.device = None,
+ offload_strategy: "AutoOffloadStrategy" | None = None,
):
hook = CustomOffloadHook(execution_device=execution_device, offload_strategy=offload_strategy)
user_hook = UserCustomOffloadHook(model_id=model_id, model=model, hook=hook)
@@ -226,7 +228,7 @@ def search_best_candidate(module_sizes, min_memory_offload):
# utils for display component info in a readable format
# TODO: move to a different file
-def summarize_dict_by_value_and_parts(d: Dict[str, Any]) -> Dict[str, Any]:
+def summarize_dict_by_value_and_parts(d: dict[str, Any]) -> dict[str, Any]:
"""Summarizes a dictionary by finding common prefixes that share the same value.
For a dictionary with dot-separated keys like: {
@@ -247,7 +249,7 @@ def summarize_dict_by_value_and_parts(d: Dict[str, Any]) -> Dict[str, Any]:
value_to_keys[value_tuple] = []
value_to_keys[value_tuple].append(key)
- def find_common_prefix(keys: List[str]) -> str:
+ def find_common_prefix(keys: list[str]) -> str:
"""Find the shortest common prefix among a list of dot-separated keys."""
if not keys:
return ""
@@ -337,10 +339,10 @@ def __init__(self):
def _lookup_ids(
self,
- name: Optional[str] = None,
- collection: Optional[str] = None,
- load_id: Optional[str] = None,
- components: Optional[OrderedDict] = None,
+ name: str | None = None,
+ collection: str | None = None,
+ load_id: str | None = None,
+ components: OrderedDict | None = None,
):
"""
Lookup component_ids by name, collection, or load_id. Does not support pattern matching. Returns a set of
@@ -381,14 +383,14 @@ def _lookup_ids(
def _id_to_name(component_id: str):
return "_".join(component_id.split("_")[:-1])
- def add(self, name: str, component: Any, collection: Optional[str] = None):
+ def add(self, name: str, component: Any, collection: str | None = None):
"""
Add a component to the ComponentsManager.
Args:
name (str): The name of the component
component (Any): The component to add
- collection (Optional[str]): The collection to add the component to
+ collection (str | None): The collection to add the component to
Returns:
str: The unique component ID, which is generated as "{name}_{id(component)}" where
@@ -506,9 +508,9 @@ def remove(self, component_id: str = None):
# YiYi TODO: rename to search_components for now, may remove this method
def search_components(
self,
- names: Optional[str] = None,
- collection: Optional[str] = None,
- load_id: Optional[str] = None,
+ names: str | None = None,
+ collection: str | None = None,
+ load_id: str | None = None,
return_dict_with_names: bool = True,
):
"""
@@ -690,7 +692,7 @@ def matches_pattern(component_id, pattern, exact_match=False):
return get_return_dict(matches, return_dict_with_names)
- def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = None, memory_reserve_margin="3GB"):
+ def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memory_reserve_margin="3GB"):
"""
Enable automatic CPU offloading for all components.
@@ -702,7 +704,7 @@ def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = None,
5. Models stay on the execution device until another model needs memory and forces them off
Args:
- device (Union[str, int, torch.device]): The execution device where models are moved for forward passes
+ device (str | int | torch.device): The execution device where models are moved for forward passes
memory_reserve_margin (str): The memory reserve margin to use, default is 3GB. This is the amount of
memory to keep free on the device to avoid running out of memory during model
execution (e.g., for intermediate activations, gradients, etc.)
@@ -767,13 +769,13 @@ def disable_auto_cpu_offload(self):
def get_model_info(
self,
component_id: str,
- fields: Optional[Union[str, List[str]]] = None,
- ) -> Optional[Dict[str, Any]]:
+ fields: str | list[str] | None = None,
+ ) -> dict[str, Any] | None:
"""Get comprehensive information about a component.
Args:
component_id (str): Name of the component to get info for
- fields (Optional[Union[str, List[str]]]):
+ fields (str | list[str] | None):
Field(s) to return. Can be a string for single field or list of fields. If None, uses the
available_info_fields setting.
@@ -995,10 +997,10 @@ def format_device(component, info):
def get_one(
self,
- component_id: Optional[str] = None,
- name: Optional[str] = None,
- collection: Optional[str] = None,
- load_id: Optional[str] = None,
+ component_id: str | None = None,
+ name: str | None = None,
+ collection: str | None = None,
+ load_id: str | None = None,
) -> Any:
"""
Get a single component by either:
@@ -1007,10 +1009,10 @@ def get_one(
Raises an error if multiple components match or none are found.
Args:
- component_id (Optional[str]): Optional component ID to get
- name (Optional[str]): Component name or pattern
- collection (Optional[str]): Optional collection to filter by
- load_id (Optional[str]): Optional load_id to filter by
+ component_id (str | None): Optional component ID to get
+ name (str | None): Component name or pattern
+ collection (str | None): Optional collection to filter by
+ load_id (str | None): Optional load_id to filter by
Returns:
A single component
@@ -1038,16 +1040,16 @@ def get_one(
return next(iter(results.values()))
- def get_ids(self, names: Union[str, List[str]] = None, collection: Optional[str] = None):
+ def get_ids(self, names: str | list[str] = None, collection: str | None = None):
"""
Get component IDs by a list of names, optionally filtered by collection.
Args:
- names (Union[str, List[str]]): List of component names
- collection (Optional[str]): Optional collection to filter by
+ names (str | list[str]): list of component names
+ collection (str | None): Optional collection to filter by
Returns:
- List[str]: List of component IDs
+ list[str]: list of component IDs
"""
ids = set()
if not isinstance(names, list):
@@ -1056,18 +1058,18 @@ def get_ids(self, names: Union[str, List[str]] = None, collection: Optional[str]
ids.update(self._lookup_ids(name=name, collection=collection))
return list(ids)
- def get_components_by_ids(self, ids: List[str], return_dict_with_names: Optional[bool] = True):
+ def get_components_by_ids(self, ids: list[str], return_dict_with_names: bool | None = True):
"""
Get components by a list of IDs.
Args:
- ids (List[str]):
- List of component IDs
- return_dict_with_names (Optional[bool]):
+ ids (list[str]):
+ list of component IDs
+ return_dict_with_names (bool | None):
Whether to return a dictionary with component names as keys:
Returns:
- Dict[str, Any]: Dictionary of components.
+ dict[str, Any]: Dictionary of components.
- If return_dict_with_names=True, keys are component names.
- If return_dict_with_names=False, keys are component IDs.
@@ -1089,16 +1091,16 @@ def get_components_by_ids(self, ids: List[str], return_dict_with_names: Optional
else:
return components
- def get_components_by_names(self, names: List[str], collection: Optional[str] = None):
+ def get_components_by_names(self, names: list[str], collection: str | None = None):
"""
Get components by a list of names, optionally filtered by collection.
Args:
- names (List[str]): List of component names
- collection (Optional[str]): Optional collection to filter by
+ names (list[str]): list of component names
+ collection (str | None): Optional collection to filter by
Returns:
- Dict[str, Any]: Dictionary of components with component names as keys
+ dict[str, Any]: Dictionary of components with component names as keys
Raises:
ValueError: If duplicate component names are found in the search results
diff --git a/src/diffusers/modular_pipelines/flux/__init__.py b/src/diffusers/modular_pipelines/flux/__init__.py
index ec00986611c8..4754ed01ce6a 100644
--- a/src/diffusers/modular_pipelines/flux/__init__.py
+++ b/src/diffusers/modular_pipelines/flux/__init__.py
@@ -21,21 +21,8 @@
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
- _import_structure["encoders"] = ["FluxTextEncoderStep"]
- _import_structure["modular_blocks"] = [
- "ALL_BLOCKS",
- "AUTO_BLOCKS",
- "AUTO_BLOCKS_KONTEXT",
- "FLUX_KONTEXT_BLOCKS",
- "TEXT2IMAGE_BLOCKS",
- "FluxAutoBeforeDenoiseStep",
- "FluxAutoBlocks",
- "FluxAutoDecodeStep",
- "FluxAutoDenoiseStep",
- "FluxKontextAutoBlocks",
- "FluxKontextAutoDenoiseStep",
- "FluxKontextBeforeDenoiseStep",
- ]
+ _import_structure["modular_blocks_flux"] = ["FluxAutoBlocks"]
+ _import_structure["modular_blocks_flux_kontext"] = ["FluxKontextAutoBlocks"]
_import_structure["modular_pipeline"] = ["FluxKontextModularPipeline", "FluxModularPipeline"]
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
@@ -45,21 +32,8 @@
except OptionalDependencyNotAvailable:
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
else:
- from .encoders import FluxTextEncoderStep
- from .modular_blocks import (
- ALL_BLOCKS,
- AUTO_BLOCKS,
- AUTO_BLOCKS_KONTEXT,
- FLUX_KONTEXT_BLOCKS,
- TEXT2IMAGE_BLOCKS,
- FluxAutoBeforeDenoiseStep,
- FluxAutoBlocks,
- FluxAutoDecodeStep,
- FluxAutoDenoiseStep,
- FluxKontextAutoBlocks,
- FluxKontextAutoDenoiseStep,
- FluxKontextBeforeDenoiseStep,
- )
+ from .modular_blocks_flux import FluxAutoBlocks
+ from .modular_blocks_flux_kontext import FluxKontextAutoBlocks
from .modular_pipeline import FluxKontextModularPipeline, FluxModularPipeline
else:
import sys
diff --git a/src/diffusers/modular_pipelines/flux/before_denoise.py b/src/diffusers/modular_pipelines/flux/before_denoise.py
index daffec986535..c28154775f5a 100644
--- a/src/diffusers/modular_pipelines/flux/before_denoise.py
+++ b/src/diffusers/modular_pipelines/flux/before_denoise.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Union
import numpy as np
import torch
@@ -33,10 +32,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -51,15 +50,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -106,7 +105,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -156,7 +155,7 @@ class FluxSetTimestepsStep(ModularPipelineBlocks):
model_name = "flux"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
@property
@@ -164,7 +163,7 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for inference"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_inference_steps", default=50),
InputParam("timesteps"),
@@ -183,7 +182,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
OutputParam(
@@ -232,7 +231,7 @@ class FluxImg2ImgSetTimestepsStep(ModularPipelineBlocks):
model_name = "flux"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
@property
@@ -240,7 +239,7 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for inference"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_inference_steps", default=50),
InputParam("timesteps"),
@@ -259,7 +258,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
OutputParam(
@@ -322,7 +321,7 @@ class FluxPrepareLatentsStep(ModularPipelineBlocks):
model_name = "flux"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return []
@property
@@ -330,11 +329,11 @@ def description(self) -> str:
return "Prepare latents step that prepares the latents for the text-to-image generation process"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("height", type_hint=int),
InputParam("width", type_hint=int),
- InputParam("latents", type_hint=Optional[torch.Tensor]),
+ InputParam("latents", type_hint=torch.Tensor | None),
InputParam("num_images_per_prompt", type_hint=int, default=1),
InputParam("generator"),
InputParam(
@@ -347,7 +346,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
@@ -431,11 +430,11 @@ def description(self) -> str:
" `prepare_latents`. Both noise and image latents should already be patchified."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="latents",
@@ -458,7 +457,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="initial_noise",
@@ -507,7 +506,7 @@ def description(self) -> str:
return "Step that prepares the RoPE inputs for the denoising process. Should be placed after text encoder and latent preparation steps."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(name="height", required=True),
InputParam(name="width", required=True),
@@ -515,18 +514,18 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="txt_ids",
kwargs_type="denoiser_input_fields",
- type_hint=List[int],
+ type_hint=list[int],
description="The sequence lengths of the prompt embeds, used for RoPE calculation.",
),
OutputParam(
name="img_ids",
kwargs_type="denoiser_input_fields",
- type_hint=List[int],
+ type_hint=list[int],
description="The sequence lengths of the image latents, used for RoPE calculation.",
),
]
@@ -557,7 +556,7 @@ def description(self) -> str:
return "Step that prepares the RoPE inputs for the denoising process of Flux Kontext. Should be placed after text encoder and latent preparation steps."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(name="image_height"),
InputParam(name="image_width"),
@@ -567,18 +566,18 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="txt_ids",
kwargs_type="denoiser_input_fields",
- type_hint=List[int],
+ type_hint=list[int],
description="The sequence lengths of the prompt embeds, used for RoPE calculation.",
),
OutputParam(
name="img_ids",
kwargs_type="denoiser_input_fields",
- type_hint=List[int],
+ type_hint=list[int],
description="The sequence lengths of the image latents, used for RoPE calculation.",
),
]
diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py
index 846549b1a376..5da861e78fcb 100644
--- a/src/diffusers/modular_pipelines/flux/decoders.py
+++ b/src/diffusers/modular_pipelines/flux/decoders.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, List, Tuple, Union
+from typing import Any
import numpy as np
import PIL
@@ -49,7 +49,7 @@ class FluxDecodeStep(ModularPipelineBlocks):
model_name = "flux"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKL),
ComponentSpec(
@@ -65,7 +65,7 @@ def description(self) -> str:
return "Step that decodes the denoised latents into images"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("output_type", default="pil"),
InputParam("height", default=1024),
@@ -79,11 +79,11 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"images",
- type_hint=Union[List[PIL.Image.Image], torch.Tensor, np.ndarray],
+ type_hint=list[PIL.Image.Image] | torch.Tensor | np.ndarray,
description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array",
)
]
diff --git a/src/diffusers/modular_pipelines/flux/denoise.py b/src/diffusers/modular_pipelines/flux/denoise.py
index 5a769df1036d..babb4a867e59 100644
--- a/src/diffusers/modular_pipelines/flux/denoise.py
+++ b/src/diffusers/modular_pipelines/flux/denoise.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, List, Tuple
+from typing import Any
import torch
@@ -36,7 +36,7 @@ class FluxLoopDenoiser(ModularPipelineBlocks):
model_name = "flux"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", FluxTransformer2DModel)]
@property
@@ -48,7 +48,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("joint_attention_kwargs"),
InputParam(
@@ -113,7 +113,7 @@ class FluxKontextLoopDenoiser(ModularPipelineBlocks):
model_name = "flux-kontext"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", FluxTransformer2DModel)]
@property
@@ -125,7 +125,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("joint_attention_kwargs"),
InputParam(
@@ -203,7 +203,7 @@ class FluxLoopAfterDenoiser(ModularPipelineBlocks):
model_name = "flux"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
@property
@@ -215,15 +215,15 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return []
@property
- def intermediate_inputs(self) -> List[str]:
+ def intermediate_inputs(self) -> list[str]:
return [InputParam("generator")]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
@torch.no_grad()
@@ -254,14 +254,14 @@ def description(self) -> str:
)
@property
- def loop_expected_components(self) -> List[ComponentSpec]:
+ def loop_expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
ComponentSpec("transformer", FluxTransformer2DModel),
]
@property
- def loop_inputs(self) -> List[InputParam]:
+ def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
"timesteps",
diff --git a/src/diffusers/modular_pipelines/flux/encoders.py b/src/diffusers/modular_pipelines/flux/encoders.py
index f0314d4771b0..583c139ff22e 100644
--- a/src/diffusers/modular_pipelines/flux/encoders.py
+++ b/src/diffusers/modular_pipelines/flux/encoders.py
@@ -13,7 +13,6 @@
# limitations under the License.
import html
-from typing import List, Optional, Union
import regex as re
import torch
@@ -55,7 +54,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -90,7 +89,7 @@ def description(self) -> str:
return "Image Preprocess step."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
@@ -101,11 +100,11 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [InputParam("resized_image"), InputParam("image"), InputParam("height"), InputParam("width")]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam(name="processed_image")]
@staticmethod
@@ -151,7 +150,7 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
@@ -162,11 +161,11 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [InputParam("image"), InputParam("_auto_resize", type_hint=bool, default=True)]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam(name="processed_image")]
@torch.no_grad()
@@ -206,7 +205,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState):
return components, state
-class FluxVaeEncoderDynamicStep(ModularPipelineBlocks):
+class FluxVaeEncoderStep(ModularPipelineBlocks):
model_name = "flux"
def __init__(
@@ -241,17 +240,17 @@ def description(self) -> str:
return f"Dynamic VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}.\n"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
components = [ComponentSpec("vae", AutoencoderKL)]
return components
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [InputParam(self._image_input_name), InputParam("generator")]
return inputs
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
self._image_latents_output_name,
@@ -291,7 +290,7 @@ def description(self) -> str:
return "Text Encoder step that generate text_embeddings to guide the image generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", CLIPTextModel),
ComponentSpec("tokenizer", CLIPTokenizer),
@@ -300,7 +299,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("prompt"),
InputParam("prompt_2"),
@@ -309,7 +308,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
@@ -332,9 +331,7 @@ def check_inputs(block_state):
raise ValueError(f"`prompt` or `prompt_2` has to be of type `str` or `list` but is {type(prompt)}")
@staticmethod
- def _get_t5_prompt_embeds(
- components, prompt: Union[str, List[str]], max_sequence_length: int, device: torch.device
- ):
+ def _get_t5_prompt_embeds(components, prompt: str | list[str], max_sequence_length: int, device: torch.device):
dtype = components.text_encoder_2.dtype
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -365,7 +362,7 @@ def _get_t5_prompt_embeds(
return prompt_embeds
@staticmethod
- def _get_clip_prompt_embeds(components, prompt: Union[str, List[str]], device: torch.device):
+ def _get_clip_prompt_embeds(components, prompt: str | list[str], device: torch.device):
prompt = [prompt] if isinstance(prompt, str) else prompt
if isinstance(components, TextualInversionLoaderMixin):
@@ -401,13 +398,13 @@ def _get_clip_prompt_embeds(components, prompt: Union[str, List[str]], device: t
@staticmethod
def encode_prompt(
components,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]],
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ device: torch.device | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
device = device or components._execution_device
diff --git a/src/diffusers/modular_pipelines/flux/inputs.py b/src/diffusers/modular_pipelines/flux/inputs.py
index 45b1c6bc136f..9d2f69dbe26f 100644
--- a/src/diffusers/modular_pipelines/flux/inputs.py
+++ b/src/diffusers/modular_pipelines/flux/inputs.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List
import torch
@@ -42,7 +41,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_images_per_prompt", default=1),
InputParam(
@@ -62,7 +61,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"batch_size",
@@ -122,13 +121,13 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip
# Adapted from `QwenImageAdditionalInputsStep`
-class FluxInputsDynamicStep(ModularPipelineBlocks):
+class FluxAdditionalInputsStep(ModularPipelineBlocks):
model_name = "flux"
def __init__(
self,
- image_latent_inputs: List[str] = ["image_latents"],
- additional_batch_inputs: List[str] = [],
+ image_latent_inputs: list[str] = ["image_latents"],
+ additional_batch_inputs: list[str] = [],
):
if not isinstance(image_latent_inputs, list):
image_latent_inputs = [image_latent_inputs]
@@ -163,7 +162,7 @@ def description(self) -> str:
return summary_section + inputs_info + placement_section
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [
InputParam(name="num_images_per_prompt", default=1),
InputParam(name="batch_size", required=True),
@@ -182,7 +181,7 @@ def inputs(self) -> List[InputParam]:
return inputs
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(name="image_height", type_hint=int, description="The height of the image latents"),
OutputParam(name="image_width", type_hint=int, description="The width of the image latents"),
@@ -244,7 +243,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip
return components, state
-class FluxKontextInputsDynamicStep(FluxInputsDynamicStep):
+class FluxKontextAdditionalInputsStep(FluxAdditionalInputsStep):
model_name = "flux-kontext"
def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
@@ -257,7 +256,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip
continue
# 1. Calculate height/width from latents
- # Unlike the `FluxInputsDynamicStep`, we don't overwrite the `block.height` and `block.width`
+ # Unlike the `FluxAdditionalInputsStep`, we don't overwrite the `block.height` and `block.width`
height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor)
if not hasattr(block_state, "image_height"):
block_state.image_height = height
@@ -304,6 +303,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip
class FluxKontextSetResolutionStep(ModularPipelineBlocks):
model_name = "flux-kontext"
+ @property
def description(self):
return (
"Determines the height and width to be used during the subsequent computations.\n"
@@ -311,7 +311,7 @@ def description(self):
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [
InputParam(name="height"),
InputParam(name="width"),
@@ -320,7 +320,7 @@ def inputs(self) -> List[InputParam]:
return inputs
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(name="height", type_hint=int, description="The height of the initial noisy latents"),
OutputParam(name="width", type_hint=int, description="The width of the initial noisy latents"),
diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks.py b/src/diffusers/modular_pipelines/flux/modular_blocks.py
deleted file mode 100644
index bd9b2d1b40c9..000000000000
--- a/src/diffusers/modular_pipelines/flux/modular_blocks.py
+++ /dev/null
@@ -1,446 +0,0 @@
-# Copyright 2025 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from ...utils import logging
-from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
-from ..modular_pipeline_utils import InsertableDict
-from .before_denoise import (
- FluxImg2ImgPrepareLatentsStep,
- FluxImg2ImgSetTimestepsStep,
- FluxKontextRoPEInputsStep,
- FluxPrepareLatentsStep,
- FluxRoPEInputsStep,
- FluxSetTimestepsStep,
-)
-from .decoders import FluxDecodeStep
-from .denoise import FluxDenoiseStep, FluxKontextDenoiseStep
-from .encoders import (
- FluxKontextProcessImagesInputStep,
- FluxProcessImagesInputStep,
- FluxTextEncoderStep,
- FluxVaeEncoderDynamicStep,
-)
-from .inputs import (
- FluxInputsDynamicStep,
- FluxKontextInputsDynamicStep,
- FluxKontextSetResolutionStep,
- FluxTextInputStep,
-)
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-
-# vae encoder (run before before_denoise)
-FluxImg2ImgVaeEncoderBlocks = InsertableDict(
- [("preprocess", FluxProcessImagesInputStep()), ("encode", FluxVaeEncoderDynamicStep())]
-)
-
-
-class FluxImg2ImgVaeEncoderStep(SequentialPipelineBlocks):
- model_name = "flux"
-
- block_classes = FluxImg2ImgVaeEncoderBlocks.values()
- block_names = FluxImg2ImgVaeEncoderBlocks.keys()
-
- @property
- def description(self) -> str:
- return "Vae encoder step that preprocess andencode the image inputs into their latent representations."
-
-
-class FluxAutoVaeEncoderStep(AutoPipelineBlocks):
- block_classes = [FluxImg2ImgVaeEncoderStep]
- block_names = ["img2img"]
- block_trigger_inputs = ["image"]
-
- @property
- def description(self):
- return (
- "Vae encoder step that encode the image inputs into their latent representations.\n"
- + "This is an auto pipeline block that works for img2img tasks.\n"
- + " - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided."
- + " - if `image` is not provided, step will be skipped."
- )
-
-
-# Flux Kontext vae encoder (run before before_denoise)
-
-FluxKontextVaeEncoderBlocks = InsertableDict(
- [("preprocess", FluxKontextProcessImagesInputStep()), ("encode", FluxVaeEncoderDynamicStep(sample_mode="argmax"))]
-)
-
-
-class FluxKontextVaeEncoderStep(SequentialPipelineBlocks):
- model_name = "flux-kontext"
-
- block_classes = FluxKontextVaeEncoderBlocks.values()
- block_names = FluxKontextVaeEncoderBlocks.keys()
-
- @property
- def description(self) -> str:
- return "Vae encoder step that preprocess andencode the image inputs into their latent representations."
-
-
-class FluxKontextAutoVaeEncoderStep(AutoPipelineBlocks):
- block_classes = [FluxKontextVaeEncoderStep]
- block_names = ["img2img"]
- block_trigger_inputs = ["image"]
-
- @property
- def description(self):
- return (
- "Vae encoder step that encode the image inputs into their latent representations.\n"
- + "This is an auto pipeline block that works for img2img tasks.\n"
- + " - `FluxKontextVaeEncoderStep` (img2img) is used when only `image` is provided."
- + " - if `image` is not provided, step will be skipped."
- )
-
-
-# before_denoise: text2img
-FluxBeforeDenoiseBlocks = InsertableDict(
- [
- ("prepare_latents", FluxPrepareLatentsStep()),
- ("set_timesteps", FluxSetTimestepsStep()),
- ("prepare_rope_inputs", FluxRoPEInputsStep()),
- ]
-)
-
-
-class FluxBeforeDenoiseStep(SequentialPipelineBlocks):
- block_classes = FluxBeforeDenoiseBlocks.values()
- block_names = FluxBeforeDenoiseBlocks.keys()
-
- @property
- def description(self):
- return "Before denoise step that prepares the inputs for the denoise step in text-to-image generation."
-
-
-# before_denoise: img2img
-FluxImg2ImgBeforeDenoiseBlocks = InsertableDict(
- [
- ("prepare_latents", FluxPrepareLatentsStep()),
- ("set_timesteps", FluxImg2ImgSetTimestepsStep()),
- ("prepare_img2img_latents", FluxImg2ImgPrepareLatentsStep()),
- ("prepare_rope_inputs", FluxRoPEInputsStep()),
- ]
-)
-
-
-class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks):
- block_classes = FluxImg2ImgBeforeDenoiseBlocks.values()
- block_names = FluxImg2ImgBeforeDenoiseBlocks.keys()
-
- @property
- def description(self):
- return "Before denoise step that prepare the inputs for the denoise step for img2img task."
-
-
-# before_denoise: all task (text2img, img2img)
-class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks):
- model_name = "flux-kontext"
- block_classes = [FluxImg2ImgBeforeDenoiseStep, FluxBeforeDenoiseStep]
- block_names = ["img2img", "text2image"]
- block_trigger_inputs = ["image_latents", None]
-
- @property
- def description(self):
- return (
- "Before denoise step that prepare the inputs for the denoise step.\n"
- + "This is an auto pipeline block that works for text2image.\n"
- + " - `FluxBeforeDenoiseStep` (text2image) is used.\n"
- + " - `FluxImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n"
- )
-
-
-# before_denoise: FluxKontext
-
-FluxKontextBeforeDenoiseBlocks = InsertableDict(
- [
- ("prepare_latents", FluxPrepareLatentsStep()),
- ("set_timesteps", FluxSetTimestepsStep()),
- ("prepare_rope_inputs", FluxKontextRoPEInputsStep()),
- ]
-)
-
-
-class FluxKontextBeforeDenoiseStep(SequentialPipelineBlocks):
- block_classes = FluxKontextBeforeDenoiseBlocks.values()
- block_names = FluxKontextBeforeDenoiseBlocks.keys()
-
- @property
- def description(self):
- return (
- "Before denoise step that prepare the inputs for the denoise step\n"
- "for img2img/text2img task for Flux Kontext."
- )
-
-
-class FluxKontextAutoBeforeDenoiseStep(AutoPipelineBlocks):
- block_classes = [FluxKontextBeforeDenoiseStep, FluxBeforeDenoiseStep]
- block_names = ["img2img", "text2image"]
- block_trigger_inputs = ["image_latents", None]
-
- @property
- def description(self):
- return (
- "Before denoise step that prepare the inputs for the denoise step.\n"
- + "This is an auto pipeline block that works for text2image.\n"
- + " - `FluxBeforeDenoiseStep` (text2image) is used.\n"
- + " - `FluxKontextBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n"
- )
-
-
-# denoise: text2image
-class FluxAutoDenoiseStep(AutoPipelineBlocks):
- block_classes = [FluxDenoiseStep]
- block_names = ["denoise"]
- block_trigger_inputs = [None]
-
- @property
- def description(self) -> str:
- return (
- "Denoise step that iteratively denoise the latents. "
- "This is a auto pipeline block that works for text2image and img2img tasks."
- " - `FluxDenoiseStep` (denoise) for text2image and img2img tasks."
- )
-
-
-# denoise: Flux Kontext
-
-
-class FluxKontextAutoDenoiseStep(AutoPipelineBlocks):
- block_classes = [FluxKontextDenoiseStep]
- block_names = ["denoise"]
- block_trigger_inputs = [None]
-
- @property
- def description(self) -> str:
- return (
- "Denoise step that iteratively denoise the latents for Flux Kontext. "
- "This is a auto pipeline block that works for text2image and img2img tasks."
- " - `FluxDenoiseStep` (denoise) for text2image and img2img tasks."
- )
-
-
-# decode: all task (text2img, img2img)
-class FluxAutoDecodeStep(AutoPipelineBlocks):
- block_classes = [FluxDecodeStep]
- block_names = ["non-inpaint"]
- block_trigger_inputs = [None]
-
- @property
- def description(self):
- return "Decode step that decode the denoised latents into image outputs.\n - `FluxDecodeStep`"
-
-
-# inputs: text2image/img2img
-FluxImg2ImgBlocks = InsertableDict(
- [("text_inputs", FluxTextInputStep()), ("additional_inputs", FluxInputsDynamicStep())]
-)
-
-
-class FluxImg2ImgInputStep(SequentialPipelineBlocks):
- model_name = "flux"
- block_classes = FluxImg2ImgBlocks.values()
- block_names = FluxImg2ImgBlocks.keys()
-
- @property
- def description(self):
- return "Input step that prepares the inputs for the img2img denoising step. It:\n"
- " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n"
- " - update height/width based `image_latents`, patchify `image_latents`."
-
-
-class FluxAutoInputStep(AutoPipelineBlocks):
- block_classes = [FluxImg2ImgInputStep, FluxTextInputStep]
- block_names = ["img2img", "text2image"]
- block_trigger_inputs = ["image_latents", None]
-
- @property
- def description(self):
- return (
- "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n"
- " This is an auto pipeline block that works for text2image/img2img tasks.\n"
- + " - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided.\n"
- + " - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided.\n"
- )
-
-
-# inputs: Flux Kontext
-
-FluxKontextBlocks = InsertableDict(
- [
- ("set_resolution", FluxKontextSetResolutionStep()),
- ("text_inputs", FluxTextInputStep()),
- ("additional_inputs", FluxKontextInputsDynamicStep()),
- ]
-)
-
-
-class FluxKontextInputStep(SequentialPipelineBlocks):
- model_name = "flux-kontext"
- block_classes = FluxKontextBlocks.values()
- block_names = FluxKontextBlocks.keys()
-
- @property
- def description(self):
- return (
- "Input step that prepares the inputs for the both text2img and img2img denoising step. It:\n"
- " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n"
- " - update height/width based `image_latents`, patchify `image_latents`."
- )
-
-
-class FluxKontextAutoInputStep(AutoPipelineBlocks):
- block_classes = [FluxKontextInputStep, FluxTextInputStep]
- # block_classes = [FluxKontextInputStep]
- block_names = ["img2img", "text2img"]
- # block_names = ["img2img"]
- block_trigger_inputs = ["image_latents", None]
- # block_trigger_inputs = ["image_latents"]
-
- @property
- def description(self):
- return (
- "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n"
- " This is an auto pipeline block that works for text2image/img2img tasks.\n"
- + " - `FluxKontextInputStep` (img2img) is used when `image_latents` is provided.\n"
- + " - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present."
- )
-
-
-class FluxCoreDenoiseStep(SequentialPipelineBlocks):
- model_name = "flux"
- block_classes = [FluxAutoInputStep, FluxAutoBeforeDenoiseStep, FluxAutoDenoiseStep]
- block_names = ["input", "before_denoise", "denoise"]
-
- @property
- def description(self):
- return (
- "Core step that performs the denoising process. \n"
- + " - `FluxAutoInputStep` (input) standardizes the inputs for the denoising step.\n"
- + " - `FluxAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n"
- + " - `FluxAutoDenoiseStep` (denoise) iteratively denoises the latents.\n"
- + "This step supports text-to-image and image-to-image tasks for Flux:\n"
- + " - for image-to-image generation, you need to provide `image_latents`\n"
- + " - for text-to-image generation, all you need to provide is prompt embeddings."
- )
-
-
-class FluxKontextCoreDenoiseStep(SequentialPipelineBlocks):
- model_name = "flux-kontext"
- block_classes = [FluxKontextAutoInputStep, FluxKontextAutoBeforeDenoiseStep, FluxKontextAutoDenoiseStep]
- block_names = ["input", "before_denoise", "denoise"]
-
- @property
- def description(self):
- return (
- "Core step that performs the denoising process. \n"
- + " - `FluxKontextAutoInputStep` (input) standardizes the inputs for the denoising step.\n"
- + " - `FluxKontextAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n"
- + " - `FluxKontextAutoDenoiseStep` (denoise) iteratively denoises the latents.\n"
- + "This step supports text-to-image and image-to-image tasks for Flux:\n"
- + " - for image-to-image generation, you need to provide `image_latents`\n"
- + " - for text-to-image generation, all you need to provide is prompt embeddings."
- )
-
-
-# Auto blocks (text2image and img2img)
-AUTO_BLOCKS = InsertableDict(
- [
- ("text_encoder", FluxTextEncoderStep()),
- ("vae_encoder", FluxAutoVaeEncoderStep()),
- ("denoise", FluxCoreDenoiseStep()),
- ("decode", FluxDecodeStep()),
- ]
-)
-
-AUTO_BLOCKS_KONTEXT = InsertableDict(
- [
- ("text_encoder", FluxTextEncoderStep()),
- ("vae_encoder", FluxKontextAutoVaeEncoderStep()),
- ("denoise", FluxKontextCoreDenoiseStep()),
- ("decode", FluxDecodeStep()),
- ]
-)
-
-
-class FluxAutoBlocks(SequentialPipelineBlocks):
- model_name = "flux"
-
- block_classes = AUTO_BLOCKS.values()
- block_names = AUTO_BLOCKS.keys()
-
- @property
- def description(self):
- return (
- "Auto Modular pipeline for text-to-image and image-to-image using Flux.\n"
- + "- for text-to-image generation, all you need to provide is `prompt`\n"
- + "- for image-to-image generation, you need to provide either `image` or `image_latents`"
- )
-
-
-class FluxKontextAutoBlocks(FluxAutoBlocks):
- model_name = "flux-kontext"
-
- block_classes = AUTO_BLOCKS_KONTEXT.values()
- block_names = AUTO_BLOCKS_KONTEXT.keys()
-
-
-TEXT2IMAGE_BLOCKS = InsertableDict(
- [
- ("text_encoder", FluxTextEncoderStep()),
- ("input", FluxTextInputStep()),
- ("prepare_latents", FluxPrepareLatentsStep()),
- ("set_timesteps", FluxSetTimestepsStep()),
- ("prepare_rope_inputs", FluxRoPEInputsStep()),
- ("denoise", FluxDenoiseStep()),
- ("decode", FluxDecodeStep()),
- ]
-)
-
-IMAGE2IMAGE_BLOCKS = InsertableDict(
- [
- ("text_encoder", FluxTextEncoderStep()),
- ("vae_encoder", FluxVaeEncoderDynamicStep()),
- ("input", FluxImg2ImgInputStep()),
- ("prepare_latents", FluxPrepareLatentsStep()),
- ("set_timesteps", FluxImg2ImgSetTimestepsStep()),
- ("prepare_img2img_latents", FluxImg2ImgPrepareLatentsStep()),
- ("prepare_rope_inputs", FluxRoPEInputsStep()),
- ("denoise", FluxDenoiseStep()),
- ("decode", FluxDecodeStep()),
- ]
-)
-
-FLUX_KONTEXT_BLOCKS = InsertableDict(
- [
- ("text_encoder", FluxTextEncoderStep()),
- ("vae_encoder", FluxVaeEncoderDynamicStep(sample_mode="argmax")),
- ("input", FluxKontextInputStep()),
- ("prepare_latents", FluxPrepareLatentsStep()),
- ("set_timesteps", FluxSetTimestepsStep()),
- ("prepare_rope_inputs", FluxKontextRoPEInputsStep()),
- ("denoise", FluxKontextDenoiseStep()),
- ("decode", FluxDecodeStep()),
- ]
-)
-
-ALL_BLOCKS = {
- "text2image": TEXT2IMAGE_BLOCKS,
- "img2img": IMAGE2IMAGE_BLOCKS,
- "auto": AUTO_BLOCKS,
- "auto_kontext": AUTO_BLOCKS_KONTEXT,
- "kontext": FLUX_KONTEXT_BLOCKS,
-}
diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py b/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py
new file mode 100644
index 000000000000..f2e78e933448
--- /dev/null
+++ b/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py
@@ -0,0 +1,586 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ FluxImg2ImgPrepareLatentsStep,
+ FluxImg2ImgSetTimestepsStep,
+ FluxPrepareLatentsStep,
+ FluxRoPEInputsStep,
+ FluxSetTimestepsStep,
+)
+from .decoders import FluxDecodeStep
+from .denoise import FluxDenoiseStep
+from .encoders import (
+ FluxProcessImagesInputStep,
+ FluxTextEncoderStep,
+ FluxVaeEncoderStep,
+)
+from .inputs import (
+ FluxAdditionalInputsStep,
+ FluxTextInputStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# vae encoder (run before before_denoise)
+
+
+# auto_docstring
+class FluxImg2ImgVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Vae encoder step that preprocess andencode the image inputs into their latent representations.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ resized_image (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ processed_image (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "flux"
+
+ block_classes = [FluxProcessImagesInputStep(), FluxVaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "Vae encoder step that preprocess andencode the image inputs into their latent representations."
+
+
+# auto_docstring
+class FluxAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Vae encoder step that encode the image inputs into their latent representations.
+ This is an auto pipeline block that works for img2img tasks.
+ - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided. - if `image` is not provided,
+ step will be skipped.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ resized_image (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ processed_image (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "flux"
+ block_classes = [FluxImg2ImgVaeEncoderStep]
+ block_names = ["img2img"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae encoder step that encode the image inputs into their latent representations.\n"
+ + "This is an auto pipeline block that works for img2img tasks.\n"
+ + " - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided."
+ + " - if `image` is not provided, step will be skipped."
+ )
+
+
+# before_denoise: text2img
+# auto_docstring
+class FluxBeforeDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepares the inputs for the denoise step in text-to-image generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux"
+ block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxRoPEInputsStep()]
+ block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepares the inputs for the denoise step in text-to-image generation."
+
+
+# before_denoise: img2img
+# auto_docstring
+class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs for the denoise step for img2img task.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The image latents to use for the denoising process. Can be generated in vae encoder and packed in input
+ step.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ initial_noise (`Tensor`):
+ The initial random noised used for inpainting denoising.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux"
+ block_classes = [
+ FluxPrepareLatentsStep(),
+ FluxImg2ImgSetTimestepsStep(),
+ FluxImg2ImgPrepareLatentsStep(),
+ FluxRoPEInputsStep(),
+ ]
+ block_names = ["prepare_latents", "set_timesteps", "prepare_img2img_latents", "prepare_rope_inputs"]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepare the inputs for the denoise step for img2img task."
+
+
+# before_denoise: all task (text2img, img2img)
+# auto_docstring
+class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs for the denoise step.
+ This is an auto pipeline block that works for text2image.
+ - `FluxBeforeDenoiseStep` (text2image) is used.
+ - `FluxImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`):
+ TODO: Add description.
+ width (`int`):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ The image latents to use for the denoising process. Can be generated in vae encoder and packed in input
+ step.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ initial_noise (`Tensor`):
+ The initial random noised used for inpainting denoising.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux"
+ block_classes = [FluxImg2ImgBeforeDenoiseStep, FluxBeforeDenoiseStep]
+ block_names = ["img2img", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step.\n"
+ + "This is an auto pipeline block that works for text2image.\n"
+ + " - `FluxBeforeDenoiseStep` (text2image) is used.\n"
+ + " - `FluxImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n"
+ )
+
+
+# inputs: text2image/img2img
+
+
+# auto_docstring
+class FluxImg2ImgInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the img2img denoising step. It:
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `prompt_embeds`)
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation
+ pooled_prompt_embeds (`Tensor`):
+ pooled text embeddings used to guide the image generation
+ image_height (`int`):
+ The height of the image latents
+ image_width (`int`):
+ The width of the image latents
+ """
+
+ model_name = "flux"
+ block_classes = [FluxTextInputStep(), FluxAdditionalInputsStep()]
+ block_names = ["text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return "Input step that prepares the inputs for the img2img denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+
+
+# auto_docstring
+class FluxAutoInputStep(AutoPipelineBlocks):
+ """
+ Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size,
+ and patchified.
+ This is an auto pipeline block that works for text2image/img2img tasks.
+ - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided.
+ - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided.
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `prompt_embeds`)
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation
+ pooled_prompt_embeds (`Tensor`):
+ pooled text embeddings used to guide the image generation
+ image_height (`int`):
+ The height of the image latents
+ image_width (`int`):
+ The width of the image latents
+ """
+
+ model_name = "flux"
+
+ block_classes = [FluxImg2ImgInputStep, FluxTextInputStep]
+ block_names = ["img2img", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n"
+ " This is an auto pipeline block that works for text2image/img2img tasks.\n"
+ + " - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided.\n"
+ + " - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided.\n"
+ )
+
+
+# auto_docstring
+class FluxCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core step that performs the denoising process for Flux.
+ This step supports text-to-image and image-to-image tasks for Flux:
+ - for image-to-image generation, you need to provide `image_latents`
+ - for text-to-image generation, all you need to provide is prompt embeddings.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux"
+ block_classes = [FluxAutoInputStep, FluxAutoBeforeDenoiseStep, FluxDenoiseStep]
+ block_names = ["input", "before_denoise", "denoise"]
+
+ @property
+ def description(self):
+ return (
+ "Core step that performs the denoising process for Flux.\n"
+ + "This step supports text-to-image and image-to-image tasks for Flux:\n"
+ + " - for image-to-image generation, you need to provide `image_latents`\n"
+ + " - for text-to-image generation, all you need to provide is prompt embeddings."
+ )
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Auto blocks (text2image and img2img)
+AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", FluxTextEncoderStep()),
+ ("vae_encoder", FluxAutoVaeEncoderStep()),
+ ("denoise", FluxCoreDenoiseStep()),
+ ("decode", FluxDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class FluxAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image and image-to-image using Flux.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`CLIPTextModel`) tokenizer (`CLIPTokenizer`) text_encoder_2 (`T5EncoderModel`) tokenizer_2
+ (`T5TokenizerFast`) image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ prompt_2 (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ resized_image (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux"
+
+ block_classes = AUTO_BLOCKS.values()
+ block_names = AUTO_BLOCKS.keys()
+
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for text-to-image and image-to-image using Flux."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py b/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py
new file mode 100644
index 000000000000..b5a5dbf78c0e
--- /dev/null
+++ b/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py
@@ -0,0 +1,585 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ FluxKontextRoPEInputsStep,
+ FluxPrepareLatentsStep,
+ FluxRoPEInputsStep,
+ FluxSetTimestepsStep,
+)
+from .decoders import FluxDecodeStep
+from .denoise import FluxKontextDenoiseStep
+from .encoders import (
+ FluxKontextProcessImagesInputStep,
+ FluxTextEncoderStep,
+ FluxVaeEncoderStep,
+)
+from .inputs import (
+ FluxKontextAdditionalInputsStep,
+ FluxKontextSetResolutionStep,
+ FluxTextInputStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Flux Kontext vae encoder (run before before_denoise)
+# auto_docstring
+class FluxKontextVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Vae encoder step that preprocess andencode the image inputs into their latent representations.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ _auto_resize (`bool`, *optional*, defaults to True):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ processed_image (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxKontextProcessImagesInputStep(), FluxVaeEncoderStep(sample_mode="argmax")]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "Vae encoder step that preprocess andencode the image inputs into their latent representations."
+
+
+# auto_docstring
+class FluxKontextAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Vae encoder step that encode the image inputs into their latent representations.
+ This is an auto pipeline block that works for image-conditioned tasks.
+ - `FluxKontextVaeEncoderStep` (image_conditioned) is used when only `image` is provided. - if `image` is not
+ provided, step will be skipped.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ _auto_resize (`bool`, *optional*, defaults to True):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ processed_image (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxKontextVaeEncoderStep]
+ block_names = ["image_conditioned"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae encoder step that encode the image inputs into their latent representations.\n"
+ + "This is an auto pipeline block that works for image-conditioned tasks.\n"
+ + " - `FluxKontextVaeEncoderStep` (image_conditioned) is used when only `image` is provided."
+ + " - if `image` is not provided, step will be skipped."
+ )
+
+
+# before_denoise: text2img
+# auto_docstring
+class FluxKontextBeforeDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepares the inputs for the denoise step for Flux Kontext
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxRoPEInputsStep()]
+ block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepares the inputs for the denoise step for Flux Kontext\n"
+ "for text-to-image tasks."
+
+
+# before_denoise: image-conditioned
+# auto_docstring
+class FluxKontextImageConditionedBeforeDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs for the denoise step for Flux Kontext
+ for image-conditioned tasks.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ image_height (`None`, *optional*):
+ TODO: Add description.
+ image_width (`None`, *optional*):
+ TODO: Add description.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxKontextRoPEInputsStep()]
+ block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step for Flux Kontext\n"
+ "for image-conditioned tasks."
+ )
+
+
+# auto_docstring
+class FluxKontextAutoBeforeDenoiseStep(AutoPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs for the denoise step.
+ This is an auto pipeline block that works for text2image.
+ - `FluxKontextBeforeDenoiseStep` (text2image) is used.
+ - `FluxKontextImageConditionedBeforeDenoiseStep` (image_conditioned) is used when only `image_latents` is
+ provided.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ image_height (`None`, *optional*):
+ TODO: Add description.
+ image_width (`None`, *optional*):
+ TODO: Add description.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxKontextImageConditionedBeforeDenoiseStep, FluxKontextBeforeDenoiseStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step.\n"
+ + "This is an auto pipeline block that works for text2image.\n"
+ + " - `FluxKontextBeforeDenoiseStep` (text2image) is used.\n"
+ + " - `FluxKontextImageConditionedBeforeDenoiseStep` (image_conditioned) is used when only `image_latents` is provided.\n"
+ )
+
+
+# inputs: Flux Kontext
+# auto_docstring
+class FluxKontextInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the both text2img and img2img denoising step. It:
+ - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).
+ - update height/width based `image_latents`, patchify `image_latents`.
+
+ Inputs:
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ max_area (`int`, *optional*, defaults to 1048576):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ height (`int`):
+ The height of the initial noisy latents
+ width (`int`):
+ The width of the initial noisy latents
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `prompt_embeds`)
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation
+ pooled_prompt_embeds (`Tensor`):
+ pooled text embeddings used to guide the image generation
+ image_height (`int`):
+ The height of the image latents
+ image_width (`int`):
+ The width of the image latents
+ """
+
+ model_name = "flux-kontext"
+ block_classes = [FluxKontextSetResolutionStep(), FluxTextInputStep(), FluxKontextAdditionalInputsStep()]
+ block_names = ["set_resolution", "text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return (
+ "Input step that prepares the inputs for the both text2img and img2img denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+ )
+
+
+# auto_docstring
+class FluxKontextAutoInputStep(AutoPipelineBlocks):
+ """
+ Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size,
+ and patchified.
+ This is an auto pipeline block that works for text2image/img2img tasks.
+ - `FluxKontextInputStep` (image_conditioned) is used when `image_latents` is provided.
+ - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present.
+
+ Inputs:
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ max_area (`int`, *optional*, defaults to 1048576):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ height (`int`):
+ The height of the initial noisy latents
+ width (`int`):
+ The width of the initial noisy latents
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `prompt_embeds`)
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation
+ pooled_prompt_embeds (`Tensor`):
+ pooled text embeddings used to guide the image generation
+ image_height (`int`):
+ The height of the image latents
+ image_width (`int`):
+ The width of the image latents
+ """
+
+ model_name = "flux-kontext"
+ block_classes = [FluxKontextInputStep, FluxTextInputStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n"
+ " This is an auto pipeline block that works for text2image/img2img tasks.\n"
+ + " - `FluxKontextInputStep` (image_conditioned) is used when `image_latents` is provided.\n"
+ + " - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present."
+ )
+
+
+# auto_docstring
+class FluxKontextCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core step that performs the denoising process for Flux Kontext.
+ This step supports text-to-image and image-conditioned tasks for Flux Kontext:
+ - for image-conditioned generation, you need to provide `image_latents`
+ - for text-to-image generation, all you need to provide is prompt embeddings.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`)
+
+ Inputs:
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ max_area (`int`, *optional*, defaults to 1048576):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux-kontext"
+ block_classes = [FluxKontextAutoInputStep, FluxKontextAutoBeforeDenoiseStep, FluxKontextDenoiseStep]
+ block_names = ["input", "before_denoise", "denoise"]
+
+ @property
+ def description(self):
+ return (
+ "Core step that performs the denoising process for Flux Kontext.\n"
+ + "This step supports text-to-image and image-conditioned tasks for Flux Kontext:\n"
+ + " - for image-conditioned generation, you need to provide `image_latents`\n"
+ + " - for text-to-image generation, all you need to provide is prompt embeddings."
+ )
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+AUTO_BLOCKS_KONTEXT = InsertableDict(
+ [
+ ("text_encoder", FluxTextEncoderStep()),
+ ("vae_encoder", FluxKontextAutoVaeEncoderStep()),
+ ("denoise", FluxKontextCoreDenoiseStep()),
+ ("decode", FluxDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class FluxKontextAutoBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline for image-to-image using Flux Kontext.
+
+ Supported workflows:
+ - `image_conditioned`: requires `image`, `prompt`
+ - `text2image`: requires `prompt`
+
+ Components:
+ text_encoder (`CLIPTextModel`) tokenizer (`CLIPTokenizer`) text_encoder_2 (`T5EncoderModel`) tokenizer_2
+ (`T5TokenizerFast`) image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ prompt_2 (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ _auto_resize (`bool`, *optional*, defaults to True):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ max_area (`int`, *optional*, defaults to 1048576):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = AUTO_BLOCKS_KONTEXT.values()
+ block_names = AUTO_BLOCKS_KONTEXT.keys()
+ _workflow_map = {
+ "image_conditioned": {"image": True, "prompt": True},
+ "text2image": {"prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Modular pipeline for image-to-image using Flux Kontext."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/src/diffusers/modular_pipelines/flux2/__init__.py b/src/diffusers/modular_pipelines/flux2/__init__.py
index 74907a9af806..d7cc8badcaf7 100644
--- a/src/diffusers/modular_pipelines/flux2/__init__.py
+++ b/src/diffusers/modular_pipelines/flux2/__init__.py
@@ -21,44 +21,14 @@
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
- _import_structure["encoders"] = [
- "Flux2TextEncoderStep",
- "Flux2RemoteTextEncoderStep",
- "Flux2VaeEncoderStep",
- ]
- _import_structure["before_denoise"] = [
- "Flux2SetTimestepsStep",
- "Flux2PrepareLatentsStep",
- "Flux2RoPEInputsStep",
- "Flux2PrepareImageLatentsStep",
- ]
- _import_structure["denoise"] = [
- "Flux2LoopDenoiser",
- "Flux2LoopAfterDenoiser",
- "Flux2DenoiseLoopWrapper",
- "Flux2DenoiseStep",
- ]
- _import_structure["decoders"] = ["Flux2DecodeStep"]
- _import_structure["inputs"] = [
- "Flux2ProcessImagesInputStep",
- "Flux2TextInputStep",
- ]
- _import_structure["modular_blocks_flux2"] = [
- "ALL_BLOCKS",
- "AUTO_BLOCKS",
- "REMOTE_AUTO_BLOCKS",
- "TEXT2IMAGE_BLOCKS",
- "IMAGE_CONDITIONED_BLOCKS",
- "Flux2AutoBlocks",
- "Flux2AutoVaeEncoderStep",
- "Flux2CoreDenoiseStep",
- "Flux2VaeEncoderSequentialStep",
- ]
- _import_structure["modular_blocks_flux2_klein"] = ["Flux2KleinAutoBlocks", "Flux2KleinBaseAutoBlocks"]
+ _import_structure["encoders"] = ["Flux2RemoteTextEncoderStep"]
+ _import_structure["modular_blocks_flux2"] = ["Flux2AutoBlocks"]
+ _import_structure["modular_blocks_flux2_klein"] = ["Flux2KleinAutoBlocks"]
+ _import_structure["modular_blocks_flux2_klein_base"] = ["Flux2KleinBaseAutoBlocks"]
_import_structure["modular_pipeline"] = [
- "Flux2ModularPipeline",
- "Flux2KleinModularPipeline",
"Flux2KleinBaseModularPipeline",
+ "Flux2KleinModularPipeline",
+ "Flux2ModularPipeline",
]
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
@@ -68,43 +38,10 @@
except OptionalDependencyNotAvailable:
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
else:
- from .before_denoise import (
- Flux2PrepareImageLatentsStep,
- Flux2PrepareLatentsStep,
- Flux2RoPEInputsStep,
- Flux2SetTimestepsStep,
- )
- from .decoders import Flux2DecodeStep
- from .denoise import (
- Flux2DenoiseLoopWrapper,
- Flux2DenoiseStep,
- Flux2LoopAfterDenoiser,
- Flux2LoopDenoiser,
- )
- from .encoders import (
- Flux2RemoteTextEncoderStep,
- Flux2TextEncoderStep,
- Flux2VaeEncoderStep,
- )
- from .inputs import (
- Flux2ProcessImagesInputStep,
- Flux2TextInputStep,
- )
- from .modular_blocks_flux2 import (
- ALL_BLOCKS,
- AUTO_BLOCKS,
- IMAGE_CONDITIONED_BLOCKS,
- REMOTE_AUTO_BLOCKS,
- TEXT2IMAGE_BLOCKS,
- Flux2AutoBlocks,
- Flux2AutoVaeEncoderStep,
- Flux2CoreDenoiseStep,
- Flux2VaeEncoderSequentialStep,
- )
- from .modular_blocks_flux2_klein import (
- Flux2KleinAutoBlocks,
- Flux2KleinBaseAutoBlocks,
- )
+ from .encoders import Flux2RemoteTextEncoderStep
+ from .modular_blocks_flux2 import Flux2AutoBlocks
+ from .modular_blocks_flux2_klein import Flux2KleinAutoBlocks
+ from .modular_blocks_flux2_klein_base import Flux2KleinBaseAutoBlocks
from .modular_pipeline import Flux2KleinBaseModularPipeline, Flux2KleinModularPipeline, Flux2ModularPipeline
else:
import sys
diff --git a/src/diffusers/modular_pipelines/flux2/before_denoise.py b/src/diffusers/modular_pipelines/flux2/before_denoise.py
index d5bab16586d7..6b1b3bd96324 100644
--- a/src/diffusers/modular_pipelines/flux2/before_denoise.py
+++ b/src/diffusers/modular_pipelines/flux2/before_denoise.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Union
import numpy as np
import torch
@@ -52,10 +51,10 @@ def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -70,15 +69,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -113,7 +112,7 @@ class Flux2SetTimestepsStep(ModularPipelineBlocks):
model_name = "flux2"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
ComponentSpec("transformer", Flux2Transformer2DModel),
@@ -124,7 +123,7 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for Flux2 inference using empirical mu calculation"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_inference_steps", default=50),
InputParam("timesteps"),
@@ -135,7 +134,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
OutputParam(
@@ -192,7 +191,7 @@ class Flux2PrepareLatentsStep(ModularPipelineBlocks):
model_name = "flux2"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return []
@property
@@ -200,11 +199,11 @@ def description(self) -> str:
return "Prepare latents step that prepares the initial noise latents for Flux2 text-to-image generation"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("height", type_hint=int),
InputParam("width", type_hint=int),
- InputParam("latents", type_hint=Optional[torch.Tensor]),
+ InputParam("latents", type_hint=torch.Tensor | None),
InputParam("num_images_per_prompt", type_hint=int, default=1),
InputParam("generator"),
InputParam(
@@ -217,7 +216,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
@@ -336,13 +335,13 @@ def description(self) -> str:
return "Step that prepares the 4D RoPE position IDs for Flux2 denoising. Should be placed after text encoder and latent preparation steps."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(name="prompt_embeds", required=True),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="txt_ids",
@@ -353,7 +352,7 @@ def intermediate_outputs(self) -> List[OutputParam]:
]
@staticmethod
- def _prepare_text_ids(x: torch.Tensor, t_coord: Optional[torch.Tensor] = None):
+ def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None):
"""Prepare 4D position IDs for text tokens."""
B, L, _ = x.shape
out_ids = []
@@ -390,14 +389,14 @@ def description(self) -> str:
return "Step that prepares the 4D RoPE position IDs for Flux2-Klein base model denoising. Should be placed after text encoder and latent preparation steps."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(name="prompt_embeds", required=True),
InputParam(name="negative_prompt_embeds", required=False),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="txt_ids",
@@ -414,7 +413,7 @@ def intermediate_outputs(self) -> List[OutputParam]:
]
@staticmethod
- def _prepare_text_ids(x: torch.Tensor, t_coord: Optional[torch.Tensor] = None):
+ def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None):
"""Prepare 4D position IDs for text tokens."""
B, L, _ = x.shape
out_ids = []
@@ -456,15 +455,15 @@ def description(self) -> str:
return "Step that prepares image latents and their position IDs for Flux2 image conditioning."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
- InputParam("image_latents", type_hint=List[torch.Tensor]),
+ InputParam("image_latents", type_hint=list[torch.Tensor]),
InputParam("batch_size", required=True, type_hint=int),
InputParam("num_images_per_prompt", default=1, type_hint=int),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"image_latents",
@@ -479,7 +478,7 @@ def intermediate_outputs(self) -> List[OutputParam]:
]
@staticmethod
- def _prepare_image_ids(image_latents: List[torch.Tensor], scale: int = 10):
+ def _prepare_image_ids(image_latents: list[torch.Tensor], scale: int = 10):
"""
Generates 4D time-space coordinates (T, H, W, L) for a sequence of image latents.
@@ -561,7 +560,7 @@ def description(self) -> str:
return "Step that prepares the guidance scale tensor for Flux2 inference"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("guidance_scale", default=4.0),
InputParam("num_images_per_prompt", default=1),
@@ -574,7 +573,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("guidance", type_hint=torch.Tensor, description="Guidance scale tensor"),
]
diff --git a/src/diffusers/modular_pipelines/flux2/decoders.py b/src/diffusers/modular_pipelines/flux2/decoders.py
index c79375072037..c8ad9401efff 100644
--- a/src/diffusers/modular_pipelines/flux2/decoders.py
+++ b/src/diffusers/modular_pipelines/flux2/decoders.py
@@ -12,7 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, List, Tuple, Union
+from __future__ import annotations
+
+from typing import Any, Union
import numpy as np
import PIL
@@ -37,7 +39,7 @@ def description(self) -> str:
return "Step that unpacks the latents from the denoising step"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam(
"latents",
@@ -54,7 +56,7 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"latents",
@@ -113,7 +115,7 @@ class Flux2DecodeStep(ModularPipelineBlocks):
model_name = "flux2"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKLFlux2),
ComponentSpec(
@@ -129,7 +131,7 @@ def description(self) -> str:
return "Step that decodes the denoised latents into images using Flux2 VAE with batch norm denormalization"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("output_type", default="pil"),
InputParam(
@@ -141,11 +143,11 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"images",
- type_hint=Union[List[PIL.Image.Image], torch.Tensor, np.ndarray],
+ type_hint=Union[list[PIL.Image.Image], torch.Tensor, np.ndarray],
description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array",
)
]
diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py
index a726959a29e2..66783cc36953 100644
--- a/src/diffusers/modular_pipelines/flux2/denoise.py
+++ b/src/diffusers/modular_pipelines/flux2/denoise.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, List, Tuple
+from typing import Any
import torch
@@ -46,7 +46,7 @@ class Flux2LoopDenoiser(ModularPipelineBlocks):
model_name = "flux2"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", Flux2Transformer2DModel)]
@property
@@ -58,7 +58,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("joint_attention_kwargs"),
InputParam(
@@ -141,7 +141,7 @@ class Flux2KleinLoopDenoiser(ModularPipelineBlocks):
model_name = "flux2-klein"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", Flux2Transformer2DModel)]
@property
@@ -153,7 +153,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("joint_attention_kwargs"),
InputParam(
@@ -230,7 +230,7 @@ class Flux2KleinBaseLoopDenoiser(ModularPipelineBlocks):
model_name = "flux2-klein"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("transformer", Flux2Transformer2DModel),
ComponentSpec(
@@ -242,7 +242,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def expected_configs(self) -> List[ConfigSpec]:
+ def expected_configs(self) -> list[ConfigSpec]:
return [
ConfigSpec(name="is_distilled", default=False),
]
@@ -256,7 +256,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("joint_attention_kwargs"),
InputParam(
@@ -363,7 +363,7 @@ class Flux2LoopAfterDenoiser(ModularPipelineBlocks):
model_name = "flux2"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
@property
@@ -375,15 +375,15 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return []
@property
- def intermediate_inputs(self) -> List[str]:
+ def intermediate_inputs(self) -> list[str]:
return [InputParam("generator")]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
@torch.no_grad()
@@ -414,14 +414,14 @@ def description(self) -> str:
)
@property
- def loop_expected_components(self) -> List[ComponentSpec]:
+ def loop_expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
ComponentSpec("transformer", Flux2Transformer2DModel),
]
@property
- def loop_inputs(self) -> List[InputParam]:
+ def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
"timesteps",
diff --git a/src/diffusers/modular_pipelines/flux2/encoders.py b/src/diffusers/modular_pipelines/flux2/encoders.py
index 265fb387367c..81d20a8f4c65 100644
--- a/src/diffusers/modular_pipelines/flux2/encoders.py
+++ b/src/diffusers/modular_pipelines/flux2/encoders.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple, Union
import torch
from transformers import AutoProcessor, Mistral3ForConditionalGeneration, Qwen2TokenizerFast, Qwen3ForCausalLM
@@ -29,7 +28,7 @@
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-def format_text_input(prompts: List[str], system_message: str = None):
+def format_text_input(prompts: list[str], system_message: str = None):
"""Format prompts for Mistral3 chat template."""
cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts]
@@ -47,7 +46,7 @@ def format_text_input(prompts: List[str], system_message: str = None):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -71,22 +70,22 @@ def description(self) -> str:
return "Text Encoder step that generates text embeddings using Mistral3 to guide the image generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Mistral3ForConditionalGeneration),
ComponentSpec("tokenizer", AutoProcessor),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("prompt"),
InputParam("max_sequence_length", type_hint=int, default=512, required=False),
- InputParam("text_encoder_out_layers", type_hint=Tuple[int], default=(10, 20, 30), required=False),
+ InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(10, 20, 30), required=False),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
@@ -106,14 +105,14 @@ def check_inputs(block_state):
def _get_mistral_3_prompt_embeds(
text_encoder: Mistral3ForConditionalGeneration,
tokenizer: AutoProcessor,
- prompt: Union[str, List[str]],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
max_sequence_length: int = 512,
# fmt: off
system_message: str = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation.",
# fmt: on
- hidden_states_layers: Tuple[int] = (10, 20, 30),
+ hidden_states_layers: tuple[int] = (10, 20, 30),
):
dtype = text_encoder.dtype if dtype is None else dtype
device = text_encoder.device if device is None else device
@@ -187,17 +186,17 @@ def description(self) -> str:
return "Text Encoder step that generates text embeddings using a remote API endpoint"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return []
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("prompt"),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
@@ -255,28 +254,28 @@ def description(self) -> str:
return "Text Encoder step that generates text embeddings using Qwen3 to guide the image generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Qwen3ForCausalLM),
ComponentSpec("tokenizer", Qwen2TokenizerFast),
]
@property
- def expected_configs(self) -> List[ConfigSpec]:
+ def expected_configs(self) -> list[ConfigSpec]:
return [
ConfigSpec(name="is_distilled", default=True),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("prompt"),
InputParam("max_sequence_length", type_hint=int, default=512, required=False),
- InputParam("text_encoder_out_layers", type_hint=Tuple[int], default=(9, 18, 27), required=False),
+ InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(9, 18, 27), required=False),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
@@ -298,11 +297,11 @@ def check_inputs(block_state):
def _get_qwen3_prompt_embeds(
text_encoder: Qwen3ForCausalLM,
tokenizer: Qwen2TokenizerFast,
- prompt: Union[str, List[str]],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
max_sequence_length: int = 512,
- hidden_states_layers: List[int] = (9, 18, 27),
+ hidden_states_layers: list[int] = (9, 18, 27),
):
dtype = text_encoder.dtype if dtype is None else dtype
device = text_encoder.device if device is None else device
@@ -384,7 +383,7 @@ def description(self) -> str:
return "Text Encoder step that generates text embeddings using Qwen3 to guide the image generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Qwen3ForCausalLM),
ComponentSpec("tokenizer", Qwen2TokenizerFast),
@@ -397,21 +396,21 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def expected_configs(self) -> List[ConfigSpec]:
+ def expected_configs(self) -> list[ConfigSpec]:
return [
ConfigSpec(name="is_distilled", default=False),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("prompt"),
InputParam("max_sequence_length", type_hint=int, default=512, required=False),
- InputParam("text_encoder_out_layers", type_hint=Tuple[int], default=(9, 18, 27), required=False),
+ InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(9, 18, 27), required=False),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
@@ -439,11 +438,11 @@ def check_inputs(block_state):
def _get_qwen3_prompt_embeds(
text_encoder: Qwen3ForCausalLM,
tokenizer: Qwen2TokenizerFast,
- prompt: Union[str, List[str]],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
max_sequence_length: int = 512,
- hidden_states_layers: List[int] = (9, 18, 27),
+ hidden_states_layers: list[int] = (9, 18, 27),
):
dtype = text_encoder.dtype if dtype is None else dtype
device = text_encoder.device if device is None else device
@@ -538,22 +537,22 @@ def description(self) -> str:
return "VAE Encoder step that encodes preprocessed images into latent representations for Flux2."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("vae", AutoencoderKLFlux2)]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
- InputParam("condition_images", type_hint=List[torch.Tensor]),
+ InputParam("condition_images", type_hint=list[torch.Tensor]),
InputParam("generator"),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"image_latents",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
description="List of latent representations for each reference image",
),
]
diff --git a/src/diffusers/modular_pipelines/flux2/inputs.py b/src/diffusers/modular_pipelines/flux2/inputs.py
index 3463de1999c6..7c6881f70a46 100644
--- a/src/diffusers/modular_pipelines/flux2/inputs.py
+++ b/src/diffusers/modular_pipelines/flux2/inputs.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List
-
import torch
from ...configuration_utils import FrozenDict
@@ -39,7 +37,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_images_per_prompt", default=1),
InputParam(
@@ -52,7 +50,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"batch_size",
@@ -101,7 +99,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_images_per_prompt", default=1),
InputParam(
@@ -121,7 +119,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"batch_size",
@@ -181,7 +179,7 @@ def description(self) -> str:
return "Image preprocess step for Flux2. Validates and preprocesses reference images."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
@@ -192,7 +190,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("image"),
InputParam("height"),
@@ -200,8 +198,8 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
- return [OutputParam(name="condition_images", type_hint=List[torch.Tensor])]
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="condition_images", type_hint=list[torch.Tensor])]
@torch.no_grad()
def __call__(self, components: Flux2ModularPipeline, state: PipelineState):
diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py
index 41a0ff7dee28..b1033a7dff9e 100644
--- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py
+++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py
@@ -12,10 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List
-
-import PIL.Image
-import torch
from ...utils import logging
from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
@@ -30,7 +26,6 @@
from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep
from .denoise import Flux2DenoiseStep
from .encoders import (
- Flux2RemoteTextEncoderStep,
Flux2TextEncoderStep,
Flux2VaeEncoderStep,
)
@@ -43,26 +38,69 @@
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-Flux2VaeEncoderBlocks = InsertableDict(
- [
- ("preprocess", Flux2ProcessImagesInputStep()),
- ("encode", Flux2VaeEncoderStep()),
- ]
-)
-
-
+# auto_docstring
class Flux2VaeEncoderSequentialStep(SequentialPipelineBlocks):
+ """
+ VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
model_name = "flux2"
- block_classes = Flux2VaeEncoderBlocks.values()
- block_names = Flux2VaeEncoderBlocks.keys()
+ block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
@property
def description(self) -> str:
return "VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning."
+# auto_docstring
class Flux2AutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder step that encodes the image inputs into their latent representations.
+ This is an auto pipeline block that works for image conditioning tasks.
+ - `Flux2VaeEncoderSequentialStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
block_classes = [Flux2VaeEncoderSequentialStep]
block_names = ["img_conditioning"]
block_trigger_inputs = ["image"]
@@ -80,7 +118,6 @@ def description(self):
Flux2CoreDenoiseBlocks = InsertableDict(
[
("input", Flux2TextInputStep()),
- ("prepare_image_latents", Flux2PrepareImageLatentsStep()),
("prepare_latents", Flux2PrepareLatentsStep()),
("set_timesteps", Flux2SetTimestepsStep()),
("prepare_guidance", Flux2PrepareGuidanceStep()),
@@ -91,7 +128,47 @@ def description(self):
)
+# auto_docstring
class Flux2CoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-dev.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 4.0):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ Packed image latents for conditioning. Shape: (B, img_seq_len, C)
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
model_name = "flux2"
block_classes = Flux2CoreDenoiseBlocks.values()
@@ -99,108 +176,181 @@ class Flux2CoreDenoiseStep(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "Core denoise step that performs the denoising process for Flux2-dev.\n"
- " - `Flux2TextInputStep` (input) standardizes the text inputs (prompt_embeds) for the denoising step.\n"
- " - `Flux2PrepareImageLatentsStep` (prepare_image_latents) prepares the image latents and image_latent_ids for the denoising step.\n"
- " - `Flux2PrepareLatentsStep` (prepare_latents) prepares the initial latents (latents) and latent_ids for the denoising step.\n"
- " - `Flux2SetTimestepsStep` (set_timesteps) sets the timesteps for the denoising step.\n"
- " - `Flux2PrepareGuidanceStep` (prepare_guidance) prepares the guidance tensor for the denoising step.\n"
- " - `Flux2RoPEInputsStep` (prepare_rope_inputs) prepares the RoPE inputs (txt_ids) for the denoising step.\n"
- " - `Flux2DenoiseStep` (denoise) iteratively denoises the latents.\n"
- " - `Flux2UnpackLatentsStep` (after_denoise) unpacks the latents from the denoising step.\n"
- )
+ return "Core denoise step that performs the denoising process for Flux2-dev."
@property
def outputs(self):
return [
- OutputParam(
- name="latents",
- type_hint=torch.Tensor,
- description="The latents from the denoising step.",
- )
+ OutputParam.template("latents"),
]
-AUTO_BLOCKS = InsertableDict(
+Flux2ImageConditionedCoreDenoiseBlocks = InsertableDict(
[
- ("text_encoder", Flux2TextEncoderStep()),
- ("vae_encoder", Flux2AutoVaeEncoderStep()),
- ("denoise", Flux2CoreDenoiseStep()),
- ("decode", Flux2DecodeStep()),
+ ("input", Flux2TextInputStep()),
+ ("prepare_image_latents", Flux2PrepareImageLatentsStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_guidance", Flux2PrepareGuidanceStep()),
+ ("prepare_rope_inputs", Flux2RoPEInputsStep()),
+ ("denoise", Flux2DenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
]
)
-REMOTE_AUTO_BLOCKS = InsertableDict(
- [
- ("text_encoder", Flux2RemoteTextEncoderStep()),
- ("vae_encoder", Flux2AutoVaeEncoderStep()),
- ("denoise", Flux2CoreDenoiseStep()),
- ("decode", Flux2DecodeStep()),
- ]
-)
-
+# auto_docstring
+class Flux2ImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-dev with image conditioning.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 4.0):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
-class Flux2AutoBlocks(SequentialPipelineBlocks):
model_name = "flux2"
- block_classes = AUTO_BLOCKS.values()
- block_names = AUTO_BLOCKS.keys()
+ block_classes = Flux2ImageConditionedCoreDenoiseBlocks.values()
+ block_names = Flux2ImageConditionedCoreDenoiseBlocks.keys()
@property
def description(self):
- return (
- "Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2.\n"
- "- For text-to-image generation, all you need to provide is `prompt`.\n"
- "- For image-conditioned generation, you need to provide `image` (list of PIL images)."
- )
+ return "Core denoise step that performs the denoising process for Flux2-dev with image conditioning."
@property
def outputs(self):
return [
- OutputParam(
- name="images",
- type_hint=List[PIL.Image.Image],
- description="The images from the decoding step.",
- )
+ OutputParam.template("latents"),
]
-TEXT2IMAGE_BLOCKS = InsertableDict(
- [
- ("text_encoder", Flux2TextEncoderStep()),
- ("text_input", Flux2TextInputStep()),
- ("prepare_latents", Flux2PrepareLatentsStep()),
- ("set_timesteps", Flux2SetTimestepsStep()),
- ("prepare_guidance", Flux2PrepareGuidanceStep()),
- ("prepare_rope_inputs", Flux2RoPEInputsStep()),
- ("denoise", Flux2DenoiseStep()),
- ("after_denoise", Flux2UnpackLatentsStep()),
- ("decode", Flux2DecodeStep()),
- ]
-)
+class Flux2AutoCoreDenoiseStep(AutoPipelineBlocks):
+ model_name = "flux2"
+
+ block_classes = [Flux2ImageConditionedCoreDenoiseStep, Flux2CoreDenoiseStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Auto core denoise step that performs the denoising process for Flux2-dev."
+ "This is an auto pipeline block that works for text-to-image and image-conditioned generation."
+ " - `Flux2CoreDenoiseStep` is used for text-to-image generation.\n"
+ " - `Flux2ImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n"
+ )
+
-IMAGE_CONDITIONED_BLOCKS = InsertableDict(
+AUTO_BLOCKS = InsertableDict(
[
("text_encoder", Flux2TextEncoderStep()),
- ("text_input", Flux2TextInputStep()),
- ("preprocess_images", Flux2ProcessImagesInputStep()),
- ("vae_encoder", Flux2VaeEncoderStep()),
- ("prepare_image_latents", Flux2PrepareImageLatentsStep()),
- ("prepare_latents", Flux2PrepareLatentsStep()),
- ("set_timesteps", Flux2SetTimestepsStep()),
- ("prepare_guidance", Flux2PrepareGuidanceStep()),
- ("prepare_rope_inputs", Flux2RoPEInputsStep()),
- ("denoise", Flux2DenoiseStep()),
- ("after_denoise", Flux2UnpackLatentsStep()),
+ ("vae_encoder", Flux2AutoVaeEncoderStep()),
+ ("denoise", Flux2AutoCoreDenoiseStep()),
("decode", Flux2DecodeStep()),
]
)
-ALL_BLOCKS = {
- "text2image": TEXT2IMAGE_BLOCKS,
- "image_conditioned": IMAGE_CONDITIONED_BLOCKS,
- "auto": AUTO_BLOCKS,
- "remote": REMOTE_AUTO_BLOCKS,
-}
+
+# auto_docstring
+class Flux2AutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image_conditioned`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Mistral3ForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor
+ (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer
+ (`Flux2Transformer2DModel`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ text_encoder_out_layers (`tuple`, *optional*, defaults to (10, 20, 30)):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 4.0):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux2"
+
+ block_classes = AUTO_BLOCKS.values()
+ block_names = AUTO_BLOCKS.keys()
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image_conditioned": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("images"),
+ ]
diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py
index 984832d77be5..5dbae43a5a7f 100644
--- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py
+++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py
@@ -12,30 +12,23 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List
-
-import PIL.Image
-import torch
from ...utils import logging
from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
from ..modular_pipeline_utils import InsertableDict, OutputParam
from .before_denoise import (
- Flux2KleinBaseRoPEInputsStep,
Flux2PrepareImageLatentsStep,
Flux2PrepareLatentsStep,
Flux2RoPEInputsStep,
Flux2SetTimestepsStep,
)
from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep
-from .denoise import Flux2KleinBaseDenoiseStep, Flux2KleinDenoiseStep
+from .denoise import Flux2KleinDenoiseStep
from .encoders import (
- Flux2KleinBaseTextEncoderStep,
Flux2KleinTextEncoderStep,
Flux2VaeEncoderStep,
)
from .inputs import (
- Flux2KleinBaseTextInputStep,
Flux2ProcessImagesInputStep,
Flux2TextInputStep,
)
@@ -47,26 +40,72 @@
# VAE encoder
################
-Flux2KleinVaeEncoderBlocks = InsertableDict(
- [
- ("preprocess", Flux2ProcessImagesInputStep()),
- ("encode", Flux2VaeEncoderStep()),
- ]
-)
-
+# auto_docstring
class Flux2KleinVaeEncoderSequentialStep(SequentialPipelineBlocks):
- model_name = "flux2"
+ """
+ VAE encoder step that preprocesses and encodes the image inputs into their latent representations.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ model_name = "flux2-klein"
- block_classes = Flux2KleinVaeEncoderBlocks.values()
- block_names = Flux2KleinVaeEncoderBlocks.keys()
+ block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
@property
def description(self) -> str:
return "VAE encoder step that preprocesses and encodes the image inputs into their latent representations."
+# auto_docstring
class Flux2KleinAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder step that encodes the image inputs into their latent representations.
+ This is an auto pipeline block that works for image conditioning tasks.
+ - `Flux2KleinVaeEncoderSequentialStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ model_name = "flux2-klein"
+
block_classes = [Flux2KleinVaeEncoderSequentialStep]
block_names = ["img_conditioning"]
block_trigger_inputs = ["image"]
@@ -88,7 +127,6 @@ def description(self):
Flux2KleinCoreDenoiseBlocks = InsertableDict(
[
("input", Flux2TextInputStep()),
- ("prepare_image_latents", Flux2PrepareImageLatentsStep()),
("prepare_latents", Flux2PrepareLatentsStep()),
("set_timesteps", Flux2SetTimestepsStep()),
("prepare_rope_inputs", Flux2RoPEInputsStep()),
@@ -98,7 +136,46 @@ def description(self):
)
+# auto_docstring
class Flux2KleinCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-Klein (distilled model), for text-to-image
+ generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ Packed image latents for conditioning. Shape: (B, img_seq_len, C)
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
model_name = "flux2-klein"
block_classes = Flux2KleinCoreDenoiseBlocks.values()
@@ -106,127 +183,218 @@ class Flux2KleinCoreDenoiseStep(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "Core denoise step that performs the denoising process for Flux2-Klein (distilled model).\n"
- " - `Flux2KleinTextInputStep` (input) standardizes the text inputs (prompt_embeds) for the denoising step.\n"
- " - `Flux2PrepareImageLatentsStep` (prepare_image_latents) prepares the image latents and image_latent_ids for the denoising step.\n"
- " - `Flux2PrepareLatentsStep` (prepare_latents) prepares the initial latents (latents) and latent_ids for the denoising step.\n"
- " - `Flux2SetTimestepsStep` (set_timesteps) sets the timesteps for the denoising step.\n"
- " - `Flux2RoPEInputsStep` (prepare_rope_inputs) prepares the RoPE inputs (txt_ids) for the denoising step.\n"
- " - `Flux2KleinDenoiseStep` (denoise) iteratively denoises the latents.\n"
- " - `Flux2UnpackLatentsStep` (after_denoise) unpacks the latents from the denoising step.\n"
- )
+ return "Core denoise step that performs the denoising process for Flux2-Klein (distilled model), for text-to-image generation."
@property
def outputs(self):
return [
- OutputParam(
- name="latents",
- type_hint=torch.Tensor,
- description="The latents from the denoising step.",
- )
+ OutputParam.template("latents"),
]
-Flux2KleinBaseCoreDenoiseBlocks = InsertableDict(
+Flux2KleinImageConditionedCoreDenoiseBlocks = InsertableDict(
[
- ("input", Flux2KleinBaseTextInputStep()),
- ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("input", Flux2TextInputStep()),
("prepare_image_latents", Flux2PrepareImageLatentsStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
("set_timesteps", Flux2SetTimestepsStep()),
- ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()),
- ("denoise", Flux2KleinBaseDenoiseStep()),
+ ("prepare_rope_inputs", Flux2RoPEInputsStep()),
+ ("denoise", Flux2KleinDenoiseStep()),
("after_denoise", Flux2UnpackLatentsStep()),
]
)
-class Flux2KleinBaseCoreDenoiseStep(SequentialPipelineBlocks):
+# auto_docstring
+class Flux2KleinImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-Klein (distilled model) with image conditioning.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
model_name = "flux2-klein"
- block_classes = Flux2KleinBaseCoreDenoiseBlocks.values()
- block_names = Flux2KleinBaseCoreDenoiseBlocks.keys()
+
+ block_classes = Flux2KleinImageConditionedCoreDenoiseBlocks.values()
+ block_names = Flux2KleinImageConditionedCoreDenoiseBlocks.keys()
@property
def description(self):
- return "Core denoise step that performs the denoising process for Flux2-Klein (base model)."
- return (
- "Core denoise step that performs the denoising process for Flux2-Klein (base model).\n"
- " - `Flux2KleinBaseTextInputStep` (input) standardizes the text inputs (prompt_embeds + negative_prompt_embeds) for the denoising step.\n"
- " - `Flux2PrepareImageLatentsStep` (prepare_image_latents) prepares the image latents and image_latent_ids for the denoising step.\n"
- " - `Flux2PrepareLatentsStep` (prepare_latents) prepares the initial latents (latents) and latent_ids for the denoising step.\n"
- " - `Flux2SetTimestepsStep` (set_timesteps) sets the timesteps for the denoising step.\n"
- " - `Flux2KleinBaseRoPEInputsStep` (prepare_rope_inputs) prepares the RoPE inputs (txt_ids + negative_txt_ids) for the denoising step.\n"
- " - `Flux2KleinBaseDenoiseStep` (denoise) iteratively denoises the latents using Classifier-Free Guidance.\n"
- " - `Flux2UnpackLatentsStep` (after_denoise) unpacks the latents from the denoising step.\n"
- )
+ return "Core denoise step that performs the denoising process for Flux2-Klein (distilled model) with image conditioning."
@property
def outputs(self):
return [
- OutputParam(
- name="latents",
- type_hint=torch.Tensor,
- description="The latents from the denoising step.",
- )
+ OutputParam.template("latents"),
]
-###
-### Auto blocks
-###
-class Flux2KleinAutoBlocks(SequentialPipelineBlocks):
+# auto_docstring
+class Flux2KleinAutoCoreDenoiseStep(AutoPipelineBlocks):
+ """
+ Auto core denoise step that performs the denoising process for Flux2-Klein.
+ This is an auto pipeline block that works for text-to-image and image-conditioned generation.
+ - `Flux2KleinCoreDenoiseStep` is used for text-to-image generation.
+ - `Flux2KleinImageConditionedCoreDenoiseStep` is used for image-conditioned generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
model_name = "flux2-klein"
- block_classes = [
- Flux2KleinTextEncoderStep(),
- Flux2KleinAutoVaeEncoderStep(),
- Flux2KleinCoreDenoiseStep(),
- Flux2DecodeStep(),
- ]
- block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ block_classes = [Flux2KleinImageConditionedCoreDenoiseStep, Flux2KleinCoreDenoiseStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
@property
def description(self):
return (
- "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein.\n"
- + " - for image-conditioned generation, you need to provide `image` (list of PIL images).\n"
- + " - for text-to-image generation, all you need to provide is `prompt`.\n"
+ "Auto core denoise step that performs the denoising process for Flux2-Klein.\n"
+ "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n"
+ " - `Flux2KleinCoreDenoiseStep` is used for text-to-image generation.\n"
+ " - `Flux2KleinImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n"
)
- @property
- def outputs(self):
- return [
- OutputParam(
- name="images",
- type_hint=List[PIL.Image.Image],
- description="The images from the decoding step.",
- )
- ]
+
+###
+### Auto blocks
+###
-class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks):
+# auto_docstring
+class Flux2KleinAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image_conditioned`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2TokenizerFast`) image_processor (`Flux2ImageProcessor`)
+ vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer
+ (`Flux2Transformer2DModel`)
+
+ Configs:
+ is_distilled (default: True)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
model_name = "flux2-klein"
block_classes = [
- Flux2KleinBaseTextEncoderStep(),
+ Flux2KleinTextEncoderStep(),
Flux2KleinAutoVaeEncoderStep(),
- Flux2KleinBaseCoreDenoiseStep(),
+ Flux2KleinAutoCoreDenoiseStep(),
Flux2DecodeStep(),
]
block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image_conditioned": {"image": True, "prompt": True},
+ }
@property
def description(self):
- return (
- "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model).\n"
- + " - for image-conditioned generation, you need to provide `image` (list of PIL images).\n"
- + " - for text-to-image generation, all you need to provide is `prompt`.\n"
- )
+ return "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein."
@property
def outputs(self):
return [
- OutputParam(
- name="images",
- type_hint=List[PIL.Image.Image],
- description="The images from the decoding step.",
- )
+ OutputParam.template("images"),
]
diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py
new file mode 100644
index 000000000000..42e025c622b4
--- /dev/null
+++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py
@@ -0,0 +1,413 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ Flux2KleinBaseRoPEInputsStep,
+ Flux2PrepareImageLatentsStep,
+ Flux2PrepareLatentsStep,
+ Flux2SetTimestepsStep,
+)
+from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep
+from .denoise import Flux2KleinBaseDenoiseStep
+from .encoders import (
+ Flux2KleinBaseTextEncoderStep,
+ Flux2VaeEncoderStep,
+)
+from .inputs import (
+ Flux2KleinBaseTextInputStep,
+ Flux2ProcessImagesInputStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+################
+# VAE encoder
+################
+
+
+# auto_docstring
+class Flux2KleinBaseVaeEncoderSequentialStep(SequentialPipelineBlocks):
+ """
+ VAE encoder step that preprocesses and encodes the image inputs into their latent representations.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ model_name = "flux2"
+
+ block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "VAE encoder step that preprocesses and encodes the image inputs into their latent representations."
+
+
+# auto_docstring
+class Flux2KleinBaseAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder step that encodes the image inputs into their latent representations.
+ This is an auto pipeline block that works for image conditioning tasks.
+ - `Flux2KleinBaseVaeEncoderSequentialStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ block_classes = [Flux2KleinBaseVaeEncoderSequentialStep]
+ block_names = ["img_conditioning"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "VAE encoder step that encodes the image inputs into their latent representations.\n"
+ "This is an auto pipeline block that works for image conditioning tasks.\n"
+ " - `Flux2KleinBaseVaeEncoderSequentialStep` is used when `image` is provided.\n"
+ " - If `image` is not provided, step will be skipped."
+ )
+
+
+###
+### Core denoise
+###
+
+Flux2KleinBaseCoreDenoiseBlocks = InsertableDict(
+ [
+ ("input", Flux2KleinBaseTextInputStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()),
+ ("denoise", Flux2KleinBaseDenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2KleinBaseCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-Klein (base model), for text-to-image generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider
+ (`ClassifierFreeGuidance`)
+
+ Configs:
+ is_distilled (default: False)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ Packed image latents for conditioning. Shape: (B, img_seq_len, C)
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = Flux2KleinBaseCoreDenoiseBlocks.values()
+ block_names = Flux2KleinBaseCoreDenoiseBlocks.keys()
+
+ @property
+ def description(self):
+ return "Core denoise step that performs the denoising process for Flux2-Klein (base model), for text-to-image generation."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+Flux2KleinBaseImageConditionedCoreDenoiseBlocks = InsertableDict(
+ [
+ ("input", Flux2KleinBaseTextInputStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("prepare_image_latents", Flux2PrepareImageLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()),
+ ("denoise", Flux2KleinBaseDenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2KleinBaseImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-Klein (base model) with image conditioning.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider
+ (`ClassifierFreeGuidance`)
+
+ Configs:
+ is_distilled (default: False)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = Flux2KleinBaseImageConditionedCoreDenoiseBlocks.values()
+ block_names = Flux2KleinBaseImageConditionedCoreDenoiseBlocks.keys()
+
+ @property
+ def description(self):
+ return "Core denoise step that performs the denoising process for Flux2-Klein (base model) with image conditioning."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# auto_docstring
+class Flux2KleinBaseAutoCoreDenoiseStep(AutoPipelineBlocks):
+ """
+ Auto core denoise step that performs the denoising process for Flux2-Klein (base model).
+ This is an auto pipeline block that works for text-to-image and image-conditioned generation.
+ - `Flux2KleinBaseCoreDenoiseStep` is used for text-to-image generation.
+ - `Flux2KleinBaseImageConditionedCoreDenoiseStep` is used for image-conditioned generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider
+ (`ClassifierFreeGuidance`)
+
+ Configs:
+ is_distilled (default: False)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = [Flux2KleinBaseImageConditionedCoreDenoiseStep, Flux2KleinBaseCoreDenoiseStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Auto core denoise step that performs the denoising process for Flux2-Klein (base model).\n"
+ "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n"
+ " - `Flux2KleinBaseCoreDenoiseStep` is used for text-to-image generation.\n"
+ " - `Flux2KleinBaseImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n"
+ )
+
+
+###
+### Auto blocks
+###
+
+
+# auto_docstring
+class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model).
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image_conditioned`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2TokenizerFast`) guider (`ClassifierFreeGuidance`)
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Configs:
+ is_distilled (default: False)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = [
+ Flux2KleinBaseTextEncoderStep(),
+ Flux2KleinBaseAutoVaeEncoderStep(),
+ Flux2KleinBaseAutoCoreDenoiseStep(),
+ Flux2DecodeStep(),
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image_conditioned": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model)."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("images"),
+ ]
diff --git a/src/diffusers/modular_pipelines/mellon_node_utils.py b/src/diffusers/modular_pipelines/mellon_node_utils.py
index 5eb6319a6b7a..f65459dfc990 100644
--- a/src/diffusers/modular_pipelines/mellon_node_utils.py
+++ b/src/diffusers/modular_pipelines/mellon_node_utils.py
@@ -5,7 +5,7 @@
# Simple typed wrapper for parameter overrides
from dataclasses import asdict, dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
from huggingface_hub import create_repo, hf_hub_download, upload_file
from huggingface_hub.utils import (
@@ -278,19 +278,19 @@ class MellonParam(metaclass=MellonParamMeta):
name: str
label: str
type: str
- display: Optional[str] = None
+ display: str | None = None
default: Any = None
- min: Optional[float] = None
- max: Optional[float] = None
- step: Optional[float] = None
+ min: float | None = None
+ max: float | None = None
+ step: float | None = None
options: Any = None
value: Any = None
- fieldOptions: Optional[Dict[str, Any]] = None
+ fieldOptions: dict[str, Any] | None = None
onChange: Any = None
onSignal: Any = None
- required_block_params: Optional[Union[str, List[str]]] = None
+ required_block_params: str | list[str] | None = None
- def to_dict(self) -> Dict[str, Any]:
+ def to_dict(self) -> dict[str, Any]:
"""Convert to dict for Mellon schema, excluding None values and internal fields."""
data = asdict(self)
return {k: v for k, v in data.items() if v is not None and k not in ("name", "required_block_params")}
@@ -314,7 +314,7 @@ def textbox(cls, name: str, default: str = "") -> "MellonParam":
)
@classmethod
- def dropdown(cls, name: str, options: List[str] = None, default: str = None) -> "MellonParam":
+ def dropdown(cls, name: str, options: list[str] = None, default: str = None) -> "MellonParam":
"""dropdown selection."""
if options and not default:
default = options[0]
@@ -579,7 +579,7 @@ def mark_required(label: str, marker: str = " *") -> str:
return f"{label}{marker}"
-def node_spec_to_mellon_dict(node_spec: Dict[str, Any], node_type: str) -> Dict[str, Any]:
+def node_spec_to_mellon_dict(node_spec: dict[str, Any], node_type: str) -> dict[str, Any]:
"""
Convert a node spec dict into Mellon format.
@@ -736,7 +736,7 @@ class MellonPipelineConfig:
def __init__(
self,
- node_specs: Dict[str, Optional[Dict[str, Any]]],
+ node_specs: dict[str, dict[str, Any] | None],
label: str = "",
default_repo: str = "",
default_dtype: str = "",
@@ -758,7 +758,7 @@ def __init__(
self.default_dtype = default_dtype
@property
- def node_params(self) -> Dict[str, Any]:
+ def node_params(self) -> dict[str, Any]:
"""Lazily compute node_params from node_specs."""
if self.node_specs is None:
return self._node_params
@@ -788,7 +788,7 @@ def __repr__(self) -> str:
lines.append(f" outputs: {outputs}")
return "\n".join(lines)
- def to_dict(self) -> Dict[str, Any]:
+ def to_dict(self) -> dict[str, Any]:
"""Convert to a JSON-serializable dictionary."""
return {
"label": self.label,
@@ -798,7 +798,7 @@ def to_dict(self) -> Dict[str, Any]:
}
@classmethod
- def from_dict(cls, data: Dict[str, Any]) -> "MellonPipelineConfig":
+ def from_dict(cls, data: dict[str, Any]) -> "MellonPipelineConfig":
"""
Create from a dictionary (loaded from JSON).
@@ -816,19 +816,19 @@ def to_json_string(self) -> str:
"""Serialize to JSON string."""
return json.dumps(self.to_dict(), indent=2, sort_keys=False) + "\n"
- def to_json_file(self, json_file_path: Union[str, os.PathLike]):
+ def to_json_file(self, json_file_path: str | os.PathLike):
"""Save to a JSON file."""
with open(json_file_path, "w", encoding="utf-8") as writer:
writer.write(self.to_json_string())
@classmethod
- def from_json_file(cls, json_file_path: Union[str, os.PathLike]) -> "MellonPipelineConfig":
+ def from_json_file(cls, json_file_path: str | os.PathLike) -> "MellonPipelineConfig":
"""Load from a JSON file."""
with open(json_file_path, "r", encoding="utf-8") as reader:
data = json.load(reader)
return cls.from_dict(data)
- def save(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
+ def save(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
"""Save the mellon pipeline config to a directory."""
if os.path.isfile(save_directory):
raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
@@ -859,7 +859,7 @@ def save(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = Fals
@classmethod
def load(
cls,
- pretrained_model_name_or_path: Union[str, os.PathLike],
+ pretrained_model_name_or_path: str | os.PathLike,
**kwargs,
) -> "MellonPipelineConfig":
"""Load a pipeline config from a local path or Hugging Face Hub."""
@@ -942,7 +942,7 @@ def load(
def from_blocks(
cls,
blocks,
- template: Dict[str, Optional[Dict[str, Any]]] = None,
+ template: dict[str, dict[str, Any]] | None = None,
label: str = "",
default_repo: str = "",
default_dtype: str = "bfloat16",
@@ -955,7 +955,7 @@ def from_blocks(
sub_block_map = dict(blocks.sub_blocks)
- def filter_spec_for_block(template_spec: Dict[str, Any], block) -> Optional[Dict[str, Any]]:
+ def filter_spec_for_block(template_spec: dict[str, Any], block) -> dict[str, Any] | None:
"""Filter template spec params based on what the block actually supports."""
block_input_names = set(block.input_names)
block_output_names = set(block.intermediate_output_names)
@@ -1025,8 +1025,8 @@ def from_custom_block(
cls,
block,
node_label: str = None,
- input_types: Optional[Dict[str, str]] = None,
- output_types: Optional[Dict[str, str]] = None,
+ input_types: dict[str, Any] | None = None,
+ output_types: dict[str, Any] | None = None,
) -> "MellonPipelineConfig":
"""
Create a MellonPipelineConfig from a custom block.
diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py
index f0722ffe99f8..76a850b63c4e 100644
--- a/src/diffusers/modular_pipelines/modular_pipeline.py
+++ b/src/diffusers/modular_pipelines/modular_pipeline.py
@@ -19,7 +19,7 @@
from collections import OrderedDict
from copy import deepcopy
from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
from huggingface_hub import create_repo
@@ -40,8 +40,11 @@
InputParam,
InsertableDict,
OutputParam,
+ combine_inputs,
+ combine_outputs,
format_components,
format_configs,
+ format_workflow,
generate_modular_model_card_content,
make_doc_string,
)
@@ -119,8 +122,8 @@ class PipelineState:
[`PipelineState`] stores the state of a pipeline. It is used to pass data between pipeline blocks.
"""
- values: Dict[str, Any] = field(default_factory=dict)
- kwargs_mapping: Dict[str, List[str]] = field(default_factory=dict)
+ values: dict[str, Any] = field(default_factory=dict)
+ kwargs_mapping: dict[str, list[str]] = field(default_factory=dict)
def set(self, key: str, value: Any, kwargs_type: str = None):
"""
@@ -139,22 +142,22 @@ def set(self, key: str, value: Any, kwargs_type: str = None):
else:
self.kwargs_mapping[kwargs_type].append(key)
- def get(self, keys: Union[str, List[str]], default: Any = None) -> Union[Any, Dict[str, Any]]:
+ def get(self, keys: str | list[str], default: Any = None) -> Any | dict[str, Any]:
"""
Get one or multiple values from the pipeline state.
Args:
- keys (Union[str, List[str]]): Key or list of keys for the values
+ keys (str | list[str]): Key or list of keys for the values
default (Any): The default value to return if not found
Returns:
- Union[Any, Dict[str, Any]]: Single value if keys is str, dictionary of values if keys is list
+ Any | dict[str, Any]: Single value if keys is str, dictionary of values if keys is list
"""
if isinstance(keys, str):
return self.values.get(keys, default)
return {key: self.values.get(key, default) for key in keys}
- def get_by_kwargs(self, kwargs_type: str) -> Dict[str, Any]:
+ def get_by_kwargs(self, kwargs_type: str) -> dict[str, Any]:
"""
Get all values with matching kwargs_type.
@@ -162,12 +165,12 @@ def get_by_kwargs(self, kwargs_type: str) -> Dict[str, Any]:
kwargs_type (str): The kwargs_type to filter by
Returns:
- Dict[str, Any]: Dictionary of values with matching kwargs_type
+ dict[str, Any]: Dictionary of values with matching kwargs_type
"""
value_names = self.kwargs_mapping.get(kwargs_type, [])
return self.get(value_names)
- def to_dict(self) -> Dict[str, Any]:
+ def to_dict(self) -> dict[str, Any]:
"""
Convert PipelineState to a dictionary.
"""
@@ -226,7 +229,7 @@ def as_dict(self):
Convert BlockState to a dictionary.
Returns:
- Dict[str, Any]: Dictionary containing all attributes of the BlockState
+ dict[str, Any]: Dictionary containing all attributes of the BlockState
"""
return dict(self.__dict__.items())
@@ -240,14 +243,14 @@ def format_value(v):
elif isinstance(v, list):
if len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
shapes = [t.shape for t in v]
- return f"List[{len(v)}] of Tensors with shapes {shapes}"
+ return f"list[{len(v)}] of Tensors with shapes {shapes}"
return repr(v)
# Handle tuples of tensors
elif isinstance(v, tuple):
if len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
shapes = [t.shape for t in v]
- return f"Tuple[{len(v)}] of Tensors with shapes {shapes}"
+ return f"tuple[{len(v)}] of Tensors with shapes {shapes}"
return repr(v)
# Handle dicts with tensor values
@@ -263,7 +266,7 @@ def format_value(v):
and hasattr(val[0], "dtype")
):
shapes = [t.shape for t in val]
- formatted_dict[k] = f"List[{len(val)}] of Tensors with shapes {shapes}"
+ formatted_dict[k] = f"list[{len(val)}] of Tensors with shapes {shapes}"
else:
formatted_dict[k] = repr(val)
return formatted_dict
@@ -287,6 +290,7 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin):
config_name = "modular_config.json"
model_name = None
+ _workflow_map = None
@classmethod
def _get_signature_keys(cls, obj):
@@ -306,16 +310,16 @@ def description(self) -> str:
return ""
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return []
@property
- def expected_configs(self) -> List[ConfigSpec]:
+ def expected_configs(self) -> list[ConfigSpec]:
return []
@property
- def inputs(self) -> List[InputParam]:
- """List of input parameters. Must be implemented by subclasses."""
+ def inputs(self) -> list[InputParam]:
+ """list of input parameters. Must be implemented by subclasses."""
return []
def _get_required_inputs(self):
@@ -327,21 +331,50 @@ def _get_required_inputs(self):
return input_names
@property
- def required_inputs(self) -> List[InputParam]:
+ def required_inputs(self) -> list[InputParam]:
return self._get_required_inputs()
@property
- def intermediate_outputs(self) -> List[OutputParam]:
- """List of intermediate output parameters. Must be implemented by subclasses."""
+ def intermediate_outputs(self) -> list[OutputParam]:
+ """list of intermediate output parameters. Must be implemented by subclasses."""
return []
def _get_outputs(self):
return self.intermediate_outputs
@property
- def outputs(self) -> List[OutputParam]:
+ def outputs(self) -> list[OutputParam]:
return self._get_outputs()
+ # currentlyonly ConditionalPipelineBlocks and SequentialPipelineBlocks support `get_execution_blocks`
+ def get_execution_blocks(self, **kwargs):
+ """
+ Get the block(s) that would execute given the inputs. Must be implemented by subclasses that support
+ conditional block selection.
+
+ Args:
+ **kwargs: Input names and values. Only trigger inputs affect block selection.
+ """
+ raise NotImplementedError(f"`get_execution_blocks` is not implemented for {self.__class__.__name__}")
+
+ # currently only SequentialPipelineBlocks support workflows
+ @property
+ def available_workflows(self):
+ """
+ Returns a list of available workflow names. Must be implemented by subclasses that define `_workflow_map`.
+ """
+ raise NotImplementedError(f"`available_workflows` is not implemented for {self.__class__.__name__}")
+
+ def get_workflow(self, workflow_name: str):
+ """
+ Get the execution blocks for a specific workflow. Must be implemented by subclasses that define
+ `_workflow_map`.
+
+ Args:
+ workflow_name: Name of the workflow to retrieve.
+ """
+ raise NotImplementedError(f"`get_workflow` is not implemented for {self.__class__.__name__}")
+
@classmethod
def from_pretrained(
cls,
@@ -403,9 +436,9 @@ def save_pretrained(self, save_directory, push_to_hub=False, **kwargs):
def init_pipeline(
self,
- pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
- components_manager: Optional[ComponentsManager] = None,
- collection: Optional[str] = None,
+ pretrained_model_name_or_path: str | os.PathLike | None = None,
+ components_manager: ComponentsManager | None = None,
+ collection: str | None = None,
) -> "ModularPipeline":
"""
create a ModularPipeline, optionally accept pretrained_model_name_or_path to load from hub.
@@ -480,86 +513,20 @@ def set_block_state(self, state: PipelineState, block_state: BlockState):
if current_value is not param: # Using identity comparison to check if object was modified
state.set(param_name, param, input_param.kwargs_type)
- @staticmethod
- def combine_inputs(*named_input_lists: List[Tuple[str, List[InputParam]]]) -> List[InputParam]:
- """
- Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if
- current default value is None and new default value is not None. Warns if multiple non-None default values
- exist for the same input.
-
- Args:
- named_input_lists: List of tuples containing (block_name, input_param_list) pairs
-
- Returns:
- List[InputParam]: Combined list of unique InputParam objects
- """
- combined_dict = {} # name -> InputParam
- value_sources = {} # name -> block_name
-
- for block_name, inputs in named_input_lists:
- for input_param in inputs:
- if input_param.name is None and input_param.kwargs_type is not None:
- input_name = "*_" + input_param.kwargs_type
- else:
- input_name = input_param.name
- if input_name in combined_dict:
- current_param = combined_dict[input_name]
- if (
- current_param.default is not None
- and input_param.default is not None
- and current_param.default != input_param.default
- ):
- warnings.warn(
- f"Multiple different default values found for input '{input_name}': "
- f"{current_param.default} (from block '{value_sources[input_name]}') and "
- f"{input_param.default} (from block '{block_name}'). Using {current_param.default}."
- )
- if current_param.default is None and input_param.default is not None:
- combined_dict[input_name] = input_param
- value_sources[input_name] = block_name
- else:
- combined_dict[input_name] = input_param
- value_sources[input_name] = block_name
-
- return list(combined_dict.values())
-
- @staticmethod
- def combine_outputs(*named_output_lists: List[Tuple[str, List[OutputParam]]]) -> List[OutputParam]:
- """
- Combines multiple lists of OutputParam objects from different blocks. For duplicate outputs, keeps the first
- occurrence of each output name.
-
- Args:
- named_output_lists: List of tuples containing (block_name, output_param_list) pairs
-
- Returns:
- List[OutputParam]: Combined list of unique OutputParam objects
- """
- combined_dict = {} # name -> OutputParam
-
- for block_name, outputs in named_output_lists:
- for output_param in outputs:
- if (output_param.name not in combined_dict) or (
- combined_dict[output_param.name].kwargs_type is None and output_param.kwargs_type is not None
- ):
- combined_dict[output_param.name] = output_param
-
- return list(combined_dict.values())
-
@property
- def input_names(self) -> List[str]:
+ def input_names(self) -> list[str]:
return [input_param.name for input_param in self.inputs if input_param.name is not None]
@property
- def intermediate_output_names(self) -> List[str]:
+ def intermediate_output_names(self) -> list[str]:
return [output_param.name for output_param in self.intermediate_outputs if output_param.name is not None]
@property
- def output_names(self) -> List[str]:
+ def output_names(self) -> list[str]:
return [output_param.name for output_param in self.outputs if output_param.name is not None]
@property
- def component_names(self) -> List[str]:
+ def component_names(self) -> list[str]:
return [component.name for component in self.expected_components]
@property
@@ -577,7 +544,8 @@ def doc(self):
class ConditionalPipelineBlocks(ModularPipelineBlocks):
"""
A Pipeline Blocks that conditionally selects a block to run based on the inputs. Subclasses must implement the
- `select_block` method to define the logic for selecting the block.
+ `select_block` method to define the logic for selecting the block. Currently, we only support selection logic based
+ on the presence or absence of inputs (i.e., whether they are `None` or not)
This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the
library implements for all the pipeline blocks (such as loading or saving etc.)
@@ -585,15 +553,20 @@ class ConditionalPipelineBlocks(ModularPipelineBlocks):
> [!WARNING] > This is an experimental feature and is likely to change in the future.
Attributes:
- block_classes: List of block classes to be used
- block_names: List of prefixes for each block
- block_trigger_inputs: List of input names that select_block() uses to determine which block to run
+ block_classes: List of block classes to be used. Must have the same length as `block_names`.
+ block_names: List of names for each block. Must have the same length as `block_classes`.
+ block_trigger_inputs: List of input names that `select_block()` uses to determine which block to run.
+ For `ConditionalPipelineBlocks`, this does not need to correspond to `block_names` and `block_classes`. For
+ `AutoPipelineBlocks`, this must have the same length as `block_names` and `block_classes`, where each
+ element specifies the trigger input for the corresponding block.
+ default_block_name: Name of the default block to run when no trigger inputs match.
+ If None, this block can be skipped entirely when no trigger inputs are provided.
"""
block_classes = []
block_names = []
block_trigger_inputs = []
- default_block_name = None # name of the default block if no trigger inputs are provided, if None, this block can be skipped if no trigger inputs are provided
+ default_block_name = None
def __init__(self):
sub_blocks = InsertableDict()
@@ -639,7 +612,7 @@ def expected_configs(self):
return expected_configs
@property
- def required_inputs(self) -> List[str]:
+ def required_inputs(self) -> list[str]:
# no default block means this conditional block can be skipped entirely
if self.default_block_name is None:
return []
@@ -655,9 +628,9 @@ def required_inputs(self) -> List[str]:
return list(required_by_all)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
named_inputs = [(name, block.inputs) for name, block in self.sub_blocks.items()]
- combined_inputs = self.combine_inputs(*named_inputs)
+ combined_inputs = combine_inputs(*named_inputs)
# mark Required inputs only if that input is required by all the blocks
for input_param in combined_inputs:
if input_param.name in self.required_inputs:
@@ -667,17 +640,18 @@ def inputs(self) -> List[Tuple[str, Any]]:
return combined_inputs
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
named_outputs = [(name, block.intermediate_outputs) for name, block in self.sub_blocks.items()]
- combined_outputs = self.combine_outputs(*named_outputs)
+ combined_outputs = combine_outputs(*named_outputs)
return combined_outputs
@property
- def outputs(self) -> List[str]:
+ def outputs(self) -> list[str]:
named_outputs = [(name, block.outputs) for name, block in self.sub_blocks.items()]
- combined_outputs = self.combine_outputs(*named_outputs)
+ combined_outputs = combine_outputs(*named_outputs)
return combined_outputs
+ # used for `__repr__`
def _get_trigger_inputs(self) -> set:
"""
Returns a set of all unique trigger input values found in this block and nested blocks.
@@ -706,21 +680,21 @@ def fn_recursive_get_trigger(blocks):
return all_triggers
- @property
- def trigger_inputs(self):
- """All trigger inputs including from nested blocks."""
- return self._get_trigger_inputs()
-
- def select_block(self, **kwargs) -> Optional[str]:
+ def select_block(self, **kwargs) -> str | None:
"""
Select the block to run based on the trigger inputs. Subclasses must implement this method to define the logic
for selecting the block.
+ Note: When trigger inputs include intermediate outputs from earlier blocks, the selection logic should only
+ depend on the presence or absence of the input (i.e., whether it is None or not), not on its actual value. This
+ is because `get_execution_blocks()` resolves conditions statically by propagating intermediate output names
+ without their runtime values.
+
Args:
**kwargs: Trigger input names and their values from the state.
Returns:
- Optional[str]: The name of the block to run, or None to use default/skip.
+ str | None: The name of the block to run, or None to use default/skip.
"""
raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement the `select_block` method.")
@@ -750,6 +724,39 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState:
logger.error(error_msg)
raise
+ def get_execution_blocks(self, **kwargs) -> ModularPipelineBlocks | None:
+ """
+ Get the block(s) that would execute given the inputs.
+
+ Recursively resolves nested ConditionalPipelineBlocks until reaching either:
+ - A leaf block (no sub_blocks or LoopSequentialPipelineBlocks) → returns single `ModularPipelineBlocks`
+ - A `SequentialPipelineBlocks` → delegates to its `get_execution_blocks()` which returns
+ a `SequentialPipelineBlocks` containing the resolved execution blocks
+
+ Args:
+ **kwargs: Input names and values. Only trigger inputs affect block selection.
+
+ Returns:
+ - `ModularPipelineBlocks`: A leaf block or resolved `SequentialPipelineBlocks`
+ - `None`: If this block would be skipped (no trigger matched and no default)
+ """
+ trigger_kwargs = {name: kwargs.get(name) for name in self.block_trigger_inputs if name is not None}
+ block_name = self.select_block(**trigger_kwargs)
+
+ if block_name is None:
+ block_name = self.default_block_name
+
+ if block_name is None:
+ return None
+
+ block = self.sub_blocks[block_name]
+
+ # Recursively resolve until we hit a leaf block
+ if block.sub_blocks and not isinstance(block, LoopSequentialPipelineBlocks):
+ return block.get_execution_blocks(**kwargs)
+
+ return block
+
def __repr__(self):
class_name = self.__class__.__name__
base_class = self.__class__.__bases__[0].__name__
@@ -757,11 +764,11 @@ def __repr__(self):
f"{class_name}(\n Class: {base_class}\n" if base_class and base_class != "object" else f"{class_name}(\n"
)
- if self.trigger_inputs:
+ if self._get_trigger_inputs():
header += "\n"
header += " " + "=" * 100 + "\n"
header += " This pipeline contains blocks that are selected at runtime based on inputs.\n"
- header += f" Trigger Inputs: {sorted(self.trigger_inputs)}\n"
+ header += f" Trigger Inputs: {sorted(self._get_trigger_inputs())}\n"
header += " " + "=" * 100 + "\n\n"
# Format description with proper indentation
@@ -828,26 +835,58 @@ def doc(self):
class AutoPipelineBlocks(ConditionalPipelineBlocks):
"""
- A Pipeline Blocks that automatically selects a block to run based on the presence of trigger inputs.
+ A Pipeline Blocks that automatically selects a block to run based on the presence of trigger inputs.
+
+ This is a specialized version of `ConditionalPipelineBlocks` where:
+ - Each block has one corresponding trigger input (1:1 mapping)
+ - Block selection is automatic: the first block whose trigger input is present gets selected
+ - `block_trigger_inputs` must have the same length as `block_names` and `block_classes`
+ - Use `None` in `block_trigger_inputs` to specify the default block, i.e the block that will run if no trigger
+ inputs are present
+
+ Attributes:
+ block_classes:
+ List of block classes to be used. Must have the same length as `block_names` and
+ `block_trigger_inputs`.
+ block_names:
+ List of names for each block. Must have the same length as `block_classes` and `block_trigger_inputs`.
+ block_trigger_inputs:
+ List of input names where each element specifies the trigger input for the corresponding block. Use
+ `None` to mark the default block.
+
+ Example:
+ ```python
+ class MyAutoBlock(AutoPipelineBlocks):
+ block_classes = [InpaintEncoderBlock, ImageEncoderBlock, TextEncoderBlock]
+ block_names = ["inpaint", "img2img", "text2img"]
+ block_trigger_inputs = ["mask_image", "image", None] # text2img is the default
+ ```
+
+ With this definition:
+ - As long as `mask_image` is provided, "inpaint" block runs (regardless of `image` being provided or not)
+ - If `mask_image` is not provided but `image` is provided, "img2img" block runs
+ - Otherwise, "text2img" block runs (default, trigger is `None`)
"""
def __init__(self):
super().__init__()
+ if self.default_block_name is not None:
+ raise ValueError(
+ f"In {self.__class__.__name__}, do not set `default_block_name` for AutoPipelineBlocks. "
+ f"Use `None` in `block_trigger_inputs` to specify the default block."
+ )
+
if not (len(self.block_classes) == len(self.block_names) == len(self.block_trigger_inputs)):
raise ValueError(
f"In {self.__class__.__name__}, the number of block_classes, block_names, and block_trigger_inputs must be the same."
)
- @property
- def default_block_name(self) -> Optional[str]:
- """Derive default_block_name from block_trigger_inputs (None entry)."""
if None in self.block_trigger_inputs:
idx = self.block_trigger_inputs.index(None)
- return self.block_names[idx]
- return None
+ self.default_block_name = self.block_names[idx]
- def select_block(self, **kwargs) -> Optional[str]:
+ def select_block(self, **kwargs) -> str | None:
"""Select block based on which trigger input is present (not None)."""
for trigger_input, block_name in zip(self.block_trigger_inputs, self.block_names):
if trigger_input is not None and kwargs.get(trigger_input) is not None:
@@ -866,8 +905,8 @@ class SequentialPipelineBlocks(ModularPipelineBlocks):
> [!WARNING] > This is an experimental feature and is likely to change in the future.
Attributes:
- block_classes: List of block classes to be used
- block_names: List of prefixes for each block
+ block_classes: list of block classes to be used
+ block_names: list of prefixes for each block
"""
block_classes = []
@@ -899,9 +938,32 @@ def expected_configs(self):
expected_configs.append(config)
return expected_configs
+ @property
+ def available_workflows(self):
+ if self._workflow_map is None:
+ raise NotImplementedError(
+ f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}"
+ )
+
+ return list(self._workflow_map.keys())
+
+ def get_workflow(self, workflow_name: str):
+ if self._workflow_map is None:
+ raise NotImplementedError(
+ f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}"
+ )
+
+ if workflow_name not in self._workflow_map:
+ raise ValueError(f"Workflow {workflow_name} not found in {self.__class__.__name__}")
+
+ trigger_inputs = self._workflow_map[workflow_name]
+ workflow_blocks = self.get_execution_blocks(**trigger_inputs)
+
+ return workflow_blocks
+
@classmethod
def from_blocks_dict(
- cls, blocks_dict: Dict[str, Any], description: Optional[str] = None
+ cls, blocks_dict: dict[str, Any], description: str | None = None
) -> "SequentialPipelineBlocks":
"""Creates a SequentialPipelineBlocks instance from a dictionary of blocks.
@@ -969,11 +1031,11 @@ def _get_inputs(self):
# YiYi TODO: add test for this
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return self._get_inputs()
@property
- def required_inputs(self) -> List[str]:
+ def required_inputs(self) -> list[str]:
# Get the first block from the dictionary
first_block = next(iter(self.sub_blocks.values()))
required_by_any = set(getattr(first_block, "required_inputs", set()))
@@ -986,7 +1048,7 @@ def required_inputs(self) -> List[str]:
return list(required_by_any)
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
named_outputs = []
for name, block in self.sub_blocks.items():
inp_names = {inp.name for inp in block.inputs}
@@ -994,12 +1056,12 @@ def intermediate_outputs(self) -> List[str]:
# filter out them here so they do not end up as intermediate_outputs
if name not in inp_names:
named_outputs.append((name, block.intermediate_outputs))
- combined_outputs = self.combine_outputs(*named_outputs)
+ combined_outputs = combine_outputs(*named_outputs)
return combined_outputs
# YiYi TODO: I think we can remove the outputs property
@property
- def outputs(self) -> List[str]:
+ def outputs(self) -> list[str]:
# return next(reversed(self.sub_blocks.values())).intermediate_outputs
return self.intermediate_outputs
@@ -1018,6 +1080,7 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState:
raise
return pipeline, state
+ # used for `__repr__`
def _get_trigger_inputs(self):
"""
Returns a set of all unique trigger input values found in the blocks.
@@ -1041,89 +1104,56 @@ def fn_recursive_get_trigger(blocks):
return fn_recursive_get_trigger(self.sub_blocks)
- @property
- def trigger_inputs(self):
- return self._get_trigger_inputs()
-
- def _traverse_trigger_blocks(self, active_inputs):
+ def get_execution_blocks(self, **kwargs) -> "SequentialPipelineBlocks":
"""
- Traverse blocks and select which ones would run given the active inputs.
+ Get the blocks that would execute given the specified inputs.
+
+ As the traversal walks through sequential blocks, intermediate outputs from resolved blocks are added to the
+ active inputs. This means conditional blocks that depend on intermediates (e.g., "run img2img if image_latents
+ is present") will resolve correctly, as long as the condition is based on presence/absence (None or not None),
+ not on the actual value.
+
Args:
- active_inputs: Dict of input names to values that are "present"
+ **kwargs: Input names and values. Only trigger inputs affect block selection.
Returns:
- OrderedDict of block_name -> block that would execute
+ SequentialPipelineBlocks containing only the blocks that would execute
"""
+ # Copy kwargs so we can add outputs as we traverse
+ active_inputs = dict(kwargs)
def fn_recursive_traverse(block, block_name, active_inputs):
result_blocks = OrderedDict()
# ConditionalPipelineBlocks (includes AutoPipelineBlocks)
if isinstance(block, ConditionalPipelineBlocks):
- trigger_kwargs = {name: active_inputs.get(name) for name in block.block_trigger_inputs}
- selected_block_name = block.select_block(**trigger_kwargs)
-
- if selected_block_name is None:
- selected_block_name = block.default_block_name
-
- if selected_block_name is None:
+ block = block.get_execution_blocks(**active_inputs)
+ if block is None:
return result_blocks
- selected_block = block.sub_blocks[selected_block_name]
-
- if selected_block.sub_blocks:
- result_blocks.update(fn_recursive_traverse(selected_block, block_name, active_inputs))
- else:
- result_blocks[block_name] = selected_block
- if hasattr(selected_block, "outputs"):
- for out in selected_block.outputs:
- active_inputs[out.name] = True
-
- return result_blocks
-
- # SequentialPipelineBlocks or LoopSequentialPipelineBlocks
- if block.sub_blocks:
+ # Has sub_blocks (SequentialPipelineBlocks/ConditionalPipelineBlocks)
+ if block.sub_blocks and not isinstance(block, LoopSequentialPipelineBlocks):
for sub_block_name, sub_block in block.sub_blocks.items():
- blocks_to_update = fn_recursive_traverse(sub_block, sub_block_name, active_inputs)
- blocks_to_update = {f"{block_name}.{k}": v for k, v in blocks_to_update.items()}
- result_blocks.update(blocks_to_update)
+ nested_blocks = fn_recursive_traverse(sub_block, sub_block_name, active_inputs)
+ nested_blocks = {f"{block_name}.{k}": v for k, v in nested_blocks.items()}
+ result_blocks.update(nested_blocks)
else:
+ # Leaf block: single ModularPipelineBlocks or LoopSequentialPipelineBlocks
result_blocks[block_name] = block
- if hasattr(block, "outputs"):
- for out in block.outputs:
+ # Add outputs to active_inputs so subsequent blocks can use them as triggers
+ if hasattr(block, "intermediate_outputs"):
+ for out in block.intermediate_outputs:
active_inputs[out.name] = True
return result_blocks
all_blocks = OrderedDict()
for block_name, block in self.sub_blocks.items():
- blocks_to_update = fn_recursive_traverse(block, block_name, active_inputs)
- all_blocks.update(blocks_to_update)
- return all_blocks
+ nested_blocks = fn_recursive_traverse(block, block_name, active_inputs)
+ all_blocks.update(nested_blocks)
- def get_execution_blocks(self, **kwargs):
- """
- Get the blocks that would execute given the specified inputs.
-
- Args:
- **kwargs: Input names and values. Only trigger inputs affect block selection.
- Pass any inputs that would be non-None at runtime.
-
- Returns:
- SequentialPipelineBlocks containing only the blocks that would execute
-
- Example:
- # Get blocks for inpainting workflow blocks = pipeline.get_execution_blocks(prompt="a cat", mask=mask,
- image=image)
-
- # Get blocks for text2image workflow blocks = pipeline.get_execution_blocks(prompt="a cat")
- """
- # Filter out None values
- active_inputs = {k: v for k, v in kwargs.items() if v is not None}
-
- blocks_triggered = self._traverse_trigger_blocks(active_inputs)
- return SequentialPipelineBlocks.from_blocks_dict(blocks_triggered)
+ return SequentialPipelineBlocks.from_blocks_dict(all_blocks)
def __repr__(self):
class_name = self.__class__.__name__
@@ -1132,18 +1162,23 @@ def __repr__(self):
f"{class_name}(\n Class: {base_class}\n" if base_class and base_class != "object" else f"{class_name}(\n"
)
- if self.trigger_inputs:
+ if self._workflow_map is None and self._get_trigger_inputs():
header += "\n"
header += " " + "=" * 100 + "\n"
header += " This pipeline contains blocks that are selected at runtime based on inputs.\n"
- header += f" Trigger Inputs: {[inp for inp in self.trigger_inputs if inp is not None]}\n"
+ header += f" Trigger Inputs: {[inp for inp in self._get_trigger_inputs() if inp is not None]}\n"
# Get first trigger input as example
- example_input = next(t for t in self.trigger_inputs if t is not None)
+ example_input = next(t for t in self._get_trigger_inputs() if t is not None)
header += f" Use `get_execution_blocks()` to see selected blocks (e.g. `get_execution_blocks({example_input}=...)`).\n"
header += " " + "=" * 100 + "\n\n"
+ description = self.description
+ if self._workflow_map is not None:
+ workflow_str = format_workflow(self._workflow_map)
+ description = f"{self.description}\n\n{workflow_str}"
+
# Format description with proper indentation
- desc_lines = self.description.split("\n")
+ desc_lines = description.split("\n")
desc = []
# First line with "Description:" label
desc.append(f" Description: {desc_lines[0]}")
@@ -1191,10 +1226,15 @@ def __repr__(self):
@property
def doc(self):
+ description = self.description
+ if self._workflow_map is not None:
+ workflow_str = format_workflow(self._workflow_map)
+ description = f"{self.description}\n\n{workflow_str}"
+
return make_doc_string(
self.inputs,
self.outputs,
- self.description,
+ description=description,
class_name=self.__class__.__name__,
expected_components=self.expected_components,
expected_configs=self.expected_configs,
@@ -1212,8 +1252,8 @@ class LoopSequentialPipelineBlocks(ModularPipelineBlocks):
> [!WARNING] > This is an experimental feature and is likely to change in the future.
Attributes:
- block_classes: List of block classes to be used
- block_names: List of prefixes for each block
+ block_classes: list of block classes to be used
+ block_names: list of prefixes for each block
"""
model_name = None
@@ -1226,20 +1266,20 @@ def description(self) -> str:
raise NotImplementedError("description method must be implemented in subclasses")
@property
- def loop_expected_components(self) -> List[ComponentSpec]:
+ def loop_expected_components(self) -> list[ComponentSpec]:
return []
@property
- def loop_expected_configs(self) -> List[ConfigSpec]:
+ def loop_expected_configs(self) -> list[ConfigSpec]:
return []
@property
- def loop_inputs(self) -> List[InputParam]:
- """List of input parameters. Must be implemented by subclasses."""
+ def loop_inputs(self) -> list[InputParam]:
+ """list of input parameters. Must be implemented by subclasses."""
return []
@property
- def loop_required_inputs(self) -> List[str]:
+ def loop_required_inputs(self) -> list[str]:
input_names = []
for input_param in self.loop_inputs:
if input_param.required:
@@ -1247,8 +1287,8 @@ def loop_required_inputs(self) -> List[str]:
return input_names
@property
- def loop_intermediate_outputs(self) -> List[OutputParam]:
- """List of intermediate output parameters. Must be implemented by subclasses."""
+ def loop_intermediate_outputs(self) -> list[OutputParam]:
+ """list of intermediate output parameters. Must be implemented by subclasses."""
return []
# modified from SequentialPipelineBlocks to include loop_expected_components
@@ -1307,7 +1347,7 @@ def inputs(self):
# modified from SequentialPipelineBlocks, if any additionan input required by the loop is required by the block
@property
- def required_inputs(self) -> List[str]:
+ def required_inputs(self) -> list[str]:
# Get the first block from the dictionary
first_block = next(iter(self.sub_blocks.values()))
required_by_any = set(getattr(first_block, "required_inputs", set()))
@@ -1325,9 +1365,9 @@ def required_inputs(self) -> List[str]:
# YiYi TODO: this need to be thought about more
# modified from SequentialPipelineBlocks to include loop_intermediate_outputs
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
named_outputs = [(name, block.intermediate_outputs) for name, block in self.sub_blocks.items()]
- combined_outputs = self.combine_outputs(*named_outputs)
+ combined_outputs = combine_outputs(*named_outputs)
for output in self.loop_intermediate_outputs:
if output.name not in {output.name for output in combined_outputs}:
combined_outputs.append(output)
@@ -1335,7 +1375,7 @@ def intermediate_outputs(self) -> List[str]:
# YiYi TODO: this need to be thought about more
@property
- def outputs(self) -> List[str]:
+ def outputs(self) -> list[str]:
return next(reversed(self.sub_blocks.values())).intermediate_outputs
def __init__(self):
@@ -1356,7 +1396,7 @@ def __init__(self):
)
@classmethod
- def from_blocks_dict(cls, blocks_dict: Dict[str, Any]) -> "LoopSequentialPipelineBlocks":
+ def from_blocks_dict(cls, blocks_dict: dict[str, Any]) -> "LoopSequentialPipelineBlocks":
"""
Creates a LoopSequentialPipelineBlocks instance from a dictionary of blocks.
@@ -1507,12 +1547,12 @@ class ModularPipeline(ConfigMixin, PushToHubMixin):
# YiYi TODO: add warning for passing multiple ComponentSpec/ConfigSpec with the same name
def __init__(
self,
- blocks: Optional[ModularPipelineBlocks] = None,
- pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
- components_manager: Optional[ComponentsManager] = None,
- collection: Optional[str] = None,
- modular_config_dict: Optional[Dict[str, Any]] = None,
- config_dict: Optional[Dict[str, Any]] = None,
+ blocks: ModularPipelineBlocks | None = None,
+ pretrained_model_name_or_path: str | os.PathLike | None = None,
+ components_manager: ComponentsManager | None = None,
+ collection: str | None = None,
+ modular_config_dict: dict[str, Any] | None = None,
+ config_dict: dict[str, Any] | None = None,
**kwargs,
):
"""
@@ -1654,7 +1694,7 @@ def __init__(
)
@property
- def default_call_parameters(self) -> Dict[str, Any]:
+ def default_call_parameters(self) -> dict[str, Any]:
"""
Returns:
- Dictionary mapping input names to their default values
@@ -1667,7 +1707,7 @@ def default_call_parameters(self) -> Dict[str, Any]:
@classmethod
def _load_pipeline_config(
cls,
- pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
+ pretrained_model_name_or_path: str | os.PathLike | None,
**load_config_kwargs,
):
try:
@@ -1697,10 +1737,10 @@ def _load_pipeline_config(
@validate_hf_hub_args
def from_pretrained(
cls,
- pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
- trust_remote_code: Optional[bool] = None,
- components_manager: Optional[ComponentsManager] = None,
- collection: Optional[str] = None,
+ pretrained_model_name_or_path: str | os.PathLike | None,
+ trust_remote_code: bool | None = None,
+ components_manager: ComponentsManager | None = None,
+ collection: str | None = None,
**kwargs,
):
"""
@@ -1779,7 +1819,7 @@ def from_pretrained(
)
return pipeline
- def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
+ def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
"""
Save the pipeline to a directory. It does not save components, you need to save them separately.
@@ -1997,26 +2037,26 @@ def dtype(self) -> torch.dtype:
return torch.float32
@property
- def null_component_names(self) -> List[str]:
+ def null_component_names(self) -> list[str]:
"""
Returns:
- - List of names for components that needs to be loaded
+ - list of names for components that needs to be loaded
"""
return [name for name in self._component_specs.keys() if hasattr(self, name) and getattr(self, name) is None]
@property
- def component_names(self) -> List[str]:
+ def component_names(self) -> list[str]:
"""
Returns:
- - List of names for all components
+ - list of names for all components
"""
return list(self.components.keys())
@property
- def pretrained_component_names(self) -> List[str]:
+ def pretrained_component_names(self) -> list[str]:
"""
Returns:
- - List of names for from_pretrained components
+ - list of names for from_pretrained components
"""
return [
name
@@ -2025,10 +2065,10 @@ def pretrained_component_names(self) -> List[str]:
]
@property
- def config_component_names(self) -> List[str]:
+ def config_component_names(self) -> list[str]:
"""
Returns:
- - List of names for from_config components
+ - list of names for from_config components
"""
return [
name
@@ -2037,7 +2077,7 @@ def config_component_names(self) -> List[str]:
]
@property
- def components(self) -> Dict[str, Any]:
+ def components(self) -> dict[str, Any]:
"""
Returns:
- Dictionary mapping component names to their objects (include both from_pretrained and from_config
@@ -2067,58 +2107,29 @@ def update_components(self, **kwargs):
- the `config` dict, which will be saved as `modular_model_index.json` during `save_pretrained`
Args:
- **kwargs: Component objects, ComponentSpec objects, or configuration values to update:
- - Component objects: Only supports components we can extract specs using
- `ComponentSpec.from_component()` method i.e. components created with ComponentSpec.load() or
- ConfigMixin subclasses that aren't nn.Modules (e.g., `unet=new_unet, text_encoder=new_encoder`)
- - ComponentSpec objects: Only supports default_creation_method == "from_config", will call create()
- method to create a new component (e.g., `guider=ComponentSpec(name="guider",
- type_hint=ClassifierFreeGuidance, config={...}, default_creation_method="from_config")`)
- - Configuration values: Simple values to update configuration settings (e.g.,
- `requires_safety_checker=False`)
-
- Raises:
- ValueError: If a component object is not supported in ComponentSpec.from_component() method:
- - nn.Module components without a valid `_diffusers_load_id` attribute
- - Non-ConfigMixin components without a valid `_diffusers_load_id` attribute
+ **kwargs: Component objects or configuration values to update:
+ - Component objects: Models loaded with `AutoModel.from_pretrained()` or `ComponentSpec.load()`
+ are automatically tagged with loading information. ConfigMixin objects without weights (e.g.,
+ schedulers, guiders) can be passed directly.
+ - Configuration values: Simple values to update configuration settings
+ (e.g., `requires_safety_checker=False`)
Examples:
```python
- # Update multiple components at once
+ # Update pre-trained model
pipeline.update_components(unet=new_unet_model, text_encoder=new_text_encoder)
# Update configuration values
pipeline.update_components(requires_safety_checker=False)
-
- # Update both components and configs together
- pipeline.update_components(unet=new_unet_model, requires_safety_checker=False)
-
- # Update with ComponentSpec objects (from_config only)
- pipeline.update_components(
- guider=ComponentSpec(
- name="guider",
- type_hint=ClassifierFreeGuidance,
- config={"guidance_scale": 5.0},
- default_creation_method="from_config",
- )
- )
```
Notes:
- - Components with trained weights must be created using ComponentSpec.load(). If the component has not been
- shared in huggingface hub and you don't have loading specs, you can upload it using `push_to_hub()`
- - ConfigMixin objects without weights (e.g., schedulers, guiders) can be passed directly
- - ComponentSpec objects with default_creation_method="from_pretrained" are not supported in
- update_components()
+ - Components with trained weights should be loaded with `AutoModel.from_pretrained()` or
+ `ComponentSpec.load()` so that loading specs are preserved for serialization.
+ - ConfigMixin objects without weights (e.g., schedulers, guiders) can be passed directly.
"""
- # extract component_specs_updates & config_specs_updates from `specs`
- passed_component_specs = {
- k: kwargs.pop(k) for k in self._component_specs if k in kwargs and isinstance(kwargs[k], ComponentSpec)
- }
- passed_components = {
- k: kwargs.pop(k) for k in self._component_specs if k in kwargs and not isinstance(kwargs[k], ComponentSpec)
- }
+ passed_components = {k: kwargs.pop(k) for k in self._component_specs if k in kwargs}
passed_config_values = {k: kwargs.pop(k) for k in self._config_specs if k in kwargs}
for name, component in passed_components.items():
@@ -2157,39 +2168,20 @@ def update_components(self, **kwargs):
if len(kwargs) > 0:
logger.warning(f"Unexpected keyword arguments, will be ignored: {kwargs.keys()}")
- created_components = {}
- for name, component_spec in passed_component_specs.items():
- if component_spec.default_creation_method == "from_pretrained":
- raise ValueError(
- "ComponentSpec object with default_creation_method == 'from_pretrained' is not supported in update_components() method"
- )
- created_components[name] = component_spec.create()
- current_component_spec = self._component_specs[name]
- # warn if type changed
- if current_component_spec.type_hint is not None and not isinstance(
- created_components[name], current_component_spec.type_hint
- ):
- logger.info(
- f"ModularPipeline.update_components: adding {name} with new type: {created_components[name].__class__.__name__}, previous type: {current_component_spec.type_hint.__name__}"
- )
- # update _component_specs based on the user passed component_spec
- self._component_specs[name] = component_spec
- self.register_components(**passed_components, **created_components)
+ self.register_components(**passed_components)
config_to_register = {}
for name, new_value in passed_config_values.items():
- # e.g. requires_aesthetics_score = False
self._config_specs[name].default = new_value
config_to_register[name] = new_value
self.register_to_config(**config_to_register)
- # YiYi TODO: support map for additional from_pretrained kwargs
- def load_components(self, names: Optional[Union[List[str], str]] = None, **kwargs):
+ def load_components(self, names: list[str] | str | None = None, **kwargs):
"""
Load selected components from specs.
Args:
- names: List of component names to load. If None, will load all components with
+ names: list of component names to load. If None, will load all components with
default_creation_method == "from_pretrained". If provided as a list or string, will load only the
specified components.
**kwargs: additional kwargs to be passed to `from_pretrained()`.Can be:
@@ -2251,7 +2243,7 @@ def load_components(self, names: Optional[Union[List[str], str]] = None, **kwarg
# Copied from diffusers.pipelines.pipeline_utils.DiffusionPipeline._maybe_raise_error_if_group_offload_active
def _maybe_raise_error_if_group_offload_active(
- self, raise_error: bool = False, module: Optional[torch.nn.Module] = None
+ self, raise_error: bool = False, module: torch.nn.Module | None = None
) -> bool:
from ..hooks.group_offloading import _is_group_offload_enabled
@@ -2463,16 +2455,16 @@ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any:
the `default_creation_method` is not `from_pretrained`, return None.
This dict contains:
- - "type_hint": Tuple[str, str]
+ - "type_hint": tuple[str, str]
Library name and class name of the component. (e.g. ("diffusers", "UNet2DConditionModel"))
- All loading fields defined by `component_spec.loading_fields()`, typically:
- - "pretrained_model_name_or_path": Optional[str]
+ - "pretrained_model_name_or_path": str | None
The model pretrained_model_name_or_pathsitory (e.g., "stabilityai/stable-diffusion-xl").
- - "subfolder": Optional[str]
+ - "subfolder": str | None
A subfolder within the pretrained_model_name_or_path where this component lives.
- - "variant": Optional[str]
+ - "variant": str | None
An optional variant identifier for the model.
- - "revision": Optional[str]
+ - "revision": str | None
A specific git revision (commit hash, tag, or branch).
- ... any other loading fields defined on the spec.
@@ -2481,7 +2473,7 @@ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any:
The spec object describing one pipeline component.
Returns:
- Dict[str, Any]: A mapping suitable for JSON serialization.
+ dict[str, Any]: A mapping suitable for JSON serialization.
Example:
>>> from diffusers.pipelines.modular_pipeline_utils import ComponentSpec >>> from diffusers import
@@ -2511,31 +2503,28 @@ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any:
}
@staticmethod
- def _dict_to_component_spec(
- name: str,
- spec_dict: Dict[str, Any],
- ) -> ComponentSpec:
+ def _dict_to_component_spec(name: str, spec_dict: dict[str, Any]) -> ComponentSpec:
"""
Reconstruct a ComponentSpec from a loading specdict.
This method converts a dictionary representation back into a ComponentSpec object. The dict should contain:
- - "type_hint": Tuple[str, str]
+ - "type_hint": tuple[str, str]
Library name and class name of the component. (e.g. ("diffusers", "UNet2DConditionModel"))
- All loading fields defined by `component_spec.loading_fields()`, typically:
- - "pretrained_model_name_or_path": Optional[str]
+ - "pretrained_model_name_or_path": str | None
The model repository (e.g., "stabilityai/stable-diffusion-xl").
- - "subfolder": Optional[str]
+ - "subfolder": str | None
A subfolder within the pretrained_model_name_or_path where this component lives.
- - "variant": Optional[str]
+ - "variant": str | None
An optional variant identifier for the model.
- - "revision": Optional[str]
+ - "revision": str | None
A specific git revision (commit hash, tag, or branch).
- ... any other loading fields defined on the spec.
Args:
name (str):
The name of the component.
- specdict (Dict[str, Any]):
+ specdict (dict[str, Any]):
A dictionary containing the component specification data.
Returns:
@@ -2579,7 +2568,7 @@ def set_progress_bar_config(self, **kwargs):
if hasattr(sub_block, "set_progress_bar_config"):
sub_block.set_progress_bar_config(**kwargs)
- def __call__(self, state: PipelineState = None, output: Union[str, List[str]] = None, **kwargs):
+ def __call__(self, state: PipelineState = None, output: str | list[str] = None, **kwargs):
"""
Execute the pipeline by running the pipeline blocks with the given inputs.
@@ -2587,11 +2576,11 @@ def __call__(self, state: PipelineState = None, output: Union[str, List[str]] =
state (`PipelineState`, optional):
PipelineState instance contains inputs and intermediate values. If None, a new `PipelineState` will be
created based on the user inputs and the pipeline blocks's requirement.
- output (`str` or `List[str]`, optional):
+ output (`str` or `list[str]`, optional):
Optional specification of what to return:
- None: Returns the complete `PipelineState` with all inputs and intermediates (default)
- str: Returns a specific intermediate value from the state (e.g. `output="image"`)
- - List[str]: Returns a dictionary of specific intermediate values (e.g. `output=["image",
+ - list[str]: Returns a dictionary of specific intermediate values (e.g. `output=["image",
"latents"]`)
@@ -2616,7 +2605,7 @@ def __call__(self, state: PipelineState = None, output: Union[str, List[str]] =
Returns:
- If `output` is None: Complete `PipelineState` containing all inputs and intermediates
- If `output` is str: The specific intermediate value from the state (e.g. `output="image"`)
- - If `output` is List[str]: Dictionary mapping output names to their values from the state (e.g.
+ - If `output` is list[str]: Dictionary mapping output names to their values from the state (e.g.
`output=["image", "latents"]`)
"""
if state is None:
diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py
index 9e11fb7ef79b..cab17c2aed5c 100644
--- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py
+++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py
@@ -14,9 +14,11 @@
import inspect
import re
+import warnings
from collections import OrderedDict
from dataclasses import dataclass, field
-from typing import Any, Dict, List, Literal, Optional, Type, Union
+from types import UnionType
+from typing import Any, Literal, Type, Union, get_args, get_origin
import PIL.Image
import torch
@@ -113,18 +115,18 @@ class ComponentSpec:
default_creation_method: Preferred creation method - "from_config" or "from_pretrained"
"""
- name: Optional[str] = None
- type_hint: Optional[Type] = None
- description: Optional[str] = None
- config: Optional[FrozenDict] = None
- pretrained_model_name_or_path: Optional[Union[str, List[str]]] = field(default=None, metadata={"loading": True})
- subfolder: Optional[str] = field(default="", metadata={"loading": True})
- variant: Optional[str] = field(default=None, metadata={"loading": True})
- revision: Optional[str] = field(default=None, metadata={"loading": True})
+ name: str | None = None
+ type_hint: Type | None = None
+ description: str | None = None
+ config: FrozenDict | None = None
+ pretrained_model_name_or_path: str | list[str] | None = field(default=None, metadata={"loading": True})
+ subfolder: str | None = field(default="", metadata={"loading": True})
+ variant: str | None = field(default=None, metadata={"loading": True})
+ revision: str | None = field(default=None, metadata={"loading": True})
default_creation_method: Literal["from_config", "from_pretrained"] = "from_pretrained"
# Deprecated
- repo: Optional[Union[str, List[str]]] = field(default=None, metadata={"loading": False})
+ repo: str | list[str] | None = field(default=None, metadata={"loading": False})
def __post_init__(self):
repo_value = self.repo
@@ -206,7 +208,7 @@ def from_component(cls, name: str, component: Any) -> Any:
)
@classmethod
- def loading_fields(cls) -> List[str]:
+ def loading_fields(cls) -> list[str]:
"""
Return the names of all loading‐related fields (i.e. those whose field.metadata["loading"] is True).
"""
@@ -225,7 +227,7 @@ def load_id(self) -> str:
return "|".join(parts)
@classmethod
- def decode_load_id(cls, load_id: str) -> Dict[str, Optional[str]]:
+ def decode_load_id(cls, load_id: str) -> dict[str, str | None]:
"""
Decode a load_id string back into a dictionary of loading fields and values.
@@ -263,7 +265,7 @@ def decode_load_id(cls, load_id: str) -> Dict[str, Optional[str]]:
# otherwise we cannot do spec -> spec.create() -> component -> ComponentSpec.from_component(component)
# the config info is lost in the process
# remove error check in from_component spec and ModularPipeline.update_components() if we remove support for non configmixin in `create()` method
- def create(self, config: Optional[Union[FrozenDict, Dict[str, Any]]] = None, **kwargs) -> Any:
+ def create(self, config: FrozenDict | dict[str, Any] | None = None, **kwargs) -> Any:
"""Create component using from_config with config."""
if self.type_hint is None or not isinstance(self.type_hint, type):
@@ -345,7 +347,7 @@ class ConfigSpec:
name: str
default: Any
- description: Optional[str] = None
+ description: str | None = None
# ======================================================
@@ -390,7 +392,7 @@ class ConfigSpec:
"description": "Torch generator for deterministic generation.",
},
"sigmas": {
- "type_hint": List[float],
+ "type_hint": list[float],
"description": "Custom sigmas for the denoising process.",
},
"strength": {
@@ -399,7 +401,7 @@ class ConfigSpec:
"description": "Strength for img2img/inpainting.",
},
"image": {
- "type_hint": Union[PIL.Image.Image, List[PIL.Image.Image]],
+ "type_hint": PIL.Image.Image | list[PIL.Image.Image],
"required": True,
"description": "Reference image(s) for denoising. Can be a single image or list of images.",
},
@@ -417,7 +419,7 @@ class ConfigSpec:
"description": "Output format: 'pil', 'np', 'pt'.",
},
"attention_kwargs": {
- "type_hint": Dict[str, Any],
+ "type_hint": dict[str, Any],
"description": "Additional kwargs for attention processors.",
},
"denoiser_input_fields": {
@@ -499,9 +501,13 @@ class ConfigSpec:
OUTPUT_PARAM_TEMPLATES = {
"images": {
- "type_hint": List[PIL.Image.Image],
+ "type_hint": list[PIL.Image.Image],
"description": "Generated images.",
},
+ "videos": {
+ "type_hint": list[PIL.Image.Image],
+ "description": "The generated videos.",
+ },
"latents": {
"type_hint": torch.Tensor,
"description": "Denoised latents.",
@@ -544,7 +550,7 @@ class InputParam:
required: bool = False
description: str = ""
kwargs_type: str = None
- metadata: Dict[str, Any] = None
+ metadata: dict[str, Any] = None
def __repr__(self):
return f"<{self.name}: {'required' if self.required else 'optional'}, default={self.default}>"
@@ -578,7 +584,7 @@ class OutputParam:
type_hint: Any = None
description: str = ""
kwargs_type: str = None
- metadata: Dict[str, Any] = None
+ metadata: dict[str, Any] = None
def __repr__(self):
return (
@@ -611,7 +617,7 @@ def format_inputs_short(inputs):
Format input parameters into a string representation, with required params first followed by optional ones.
Args:
- inputs: List of input parameters with 'required' and 'name' attributes, and 'default' for optional params
+ inputs: list of input parameters with 'required' and 'name' attributes, and 'default' for optional params
Returns:
str: Formatted string of input parameters
@@ -640,9 +646,9 @@ def format_intermediates_short(intermediate_inputs, required_intermediate_inputs
Formats intermediate inputs and outputs of a block into a string representation.
Args:
- intermediate_inputs: List of intermediate input parameters
- required_intermediate_inputs: List of required intermediate input names
- intermediate_outputs: List of intermediate output parameters
+ intermediate_inputs: list of intermediate input parameters
+ required_intermediate_inputs: list of required intermediate input names
+ intermediate_outputs: list of intermediate output parameters
Returns:
str: Formatted string like:
@@ -689,7 +695,7 @@ def format_params(params, header="Args", indent_level=4, max_line_length=115):
"""Format a list of InputParam or OutputParam objects into a readable string representation.
Args:
- params: List of InputParam or OutputParam objects to format
+ params: list of InputParam or OutputParam objects to format
header: Header text to use (e.g. "Args" or "Returns")
indent_level: Number of spaces to indent each parameter line (default: 4)
max_line_length: Maximum length for each line before wrapping (default: 115)
@@ -706,9 +712,9 @@ def format_params(params, header="Args", indent_level=4, max_line_length=115):
formatted_params = []
def get_type_str(type_hint):
- if hasattr(type_hint, "__origin__") and type_hint.__origin__ is Union:
- types = [t.__name__ if hasattr(t, "__name__") else str(t) for t in type_hint.__args__]
- return f"Union[{', '.join(types)}]"
+ if isinstance(type_hint, UnionType) or get_origin(type_hint) is Union:
+ type_strs = [t.__name__ if hasattr(t, "__name__") else str(t) for t in get_args(type_hint)]
+ return " | ".join(type_strs)
return type_hint.__name__ if hasattr(type_hint, "__name__") else str(type_hint)
def wrap_text(text, indent, max_length):
@@ -769,7 +775,7 @@ def format_input_params(input_params, indent_level=4, max_line_length=115):
"""Format a list of InputParam objects into a readable string representation.
Args:
- input_params: List of InputParam objects to format
+ input_params: list of InputParam objects to format
indent_level: Number of spaces to indent each parameter line (default: 4)
max_line_length: Maximum length for each line before wrapping (default: 115)
@@ -783,7 +789,7 @@ def format_output_params(output_params, indent_level=4, max_line_length=115):
"""Format a list of OutputParam objects into a readable string representation.
Args:
- output_params: List of OutputParam objects to format
+ output_params: list of OutputParam objects to format
indent_level: Number of spaces to indent each parameter line (default: 4)
max_line_length: Maximum length for each line before wrapping (default: 115)
@@ -797,7 +803,7 @@ def format_components(components, indent_level=4, max_line_length=115, add_empty
"""Format a list of ComponentSpec objects into a readable string representation.
Args:
- components: List of ComponentSpec objects to format
+ components: list of ComponentSpec objects to format
indent_level: Number of spaces to indent each component line (default: 4)
max_line_length: Maximum length for each line before wrapping (default: 115)
add_empty_lines: Whether to add empty lines between components (default: True)
@@ -852,7 +858,7 @@ def format_configs(configs, indent_level=4, max_line_length=115, add_empty_lines
"""Format a list of ConfigSpec objects into a readable string representation.
Args:
- configs: List of ConfigSpec objects to format
+ configs: list of ConfigSpec objects to format
indent_level: Number of spaces to indent each config line (default: 4)
max_line_length: Maximum length for each line before wrapping (default: 115)
add_empty_lines: Whether to add empty lines between configs (default: True)
@@ -886,6 +892,30 @@ def format_configs(configs, indent_level=4, max_line_length=115, add_empty_lines
return "\n".join(formatted_configs)
+def format_workflow(workflow_map):
+ """Format a workflow map into a readable string representation.
+
+ Args:
+ workflow_map: Dictionary mapping workflow names to trigger inputs
+
+ Returns:
+ A formatted string representing all workflows
+ """
+ if workflow_map is None:
+ return ""
+
+ lines = ["Supported workflows:"]
+ for workflow_name, trigger_inputs in workflow_map.items():
+ required_inputs = [k for k, v in trigger_inputs.items() if v]
+ if required_inputs:
+ inputs_str = ", ".join(f"`{t}`" for t in required_inputs)
+ lines.append(f" - `{workflow_name}`: requires {inputs_str}")
+ else:
+ lines.append(f" - `{workflow_name}`: default (no additional inputs required)")
+
+ return "\n".join(lines)
+
+
def make_doc_string(
inputs,
outputs,
@@ -898,13 +928,13 @@ def make_doc_string(
Generates a formatted documentation string describing the pipeline block's parameters and structure.
Args:
- inputs: List of input parameters
- intermediate_inputs: List of intermediate input parameters
- outputs: List of output parameters
+ inputs: list of input parameters
+ intermediate_inputs: list of intermediate input parameters
+ outputs: list of output parameters
description (str, *optional*): Description of the block
class_name (str, *optional*): Name of the class to include in the documentation
- expected_components (List[ComponentSpec], *optional*): List of expected components
- expected_configs (List[ConfigSpec], *optional*): List of expected configurations
+ expected_components (list[ComponentSpec], *optional*): list of expected components
+ expected_configs (list[ConfigSpec], *optional*): list of expected configurations
Returns:
str: A formatted string containing information about components, configs, call parameters,
@@ -942,7 +972,73 @@ def make_doc_string(
return output
-def generate_modular_model_card_content(blocks) -> Dict[str, Any]:
+def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> list[InputParam]:
+ """
+ Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if current
+ default value is None and new default value is not None. Warns if multiple non-None default values exist for the
+ same input.
+
+ Args:
+ named_input_lists: List of tuples containing (block_name, input_param_list) pairs
+
+ Returns:
+ List[InputParam]: Combined list of unique InputParam objects
+ """
+ combined_dict = {} # name -> InputParam
+ value_sources = {} # name -> block_name
+
+ for block_name, inputs in named_input_lists:
+ for input_param in inputs:
+ if input_param.name is None and input_param.kwargs_type is not None:
+ input_name = "*_" + input_param.kwargs_type
+ else:
+ input_name = input_param.name
+ if input_name in combined_dict:
+ current_param = combined_dict[input_name]
+ if (
+ current_param.default is not None
+ and input_param.default is not None
+ and current_param.default != input_param.default
+ ):
+ warnings.warn(
+ f"Multiple different default values found for input '{input_name}': "
+ f"{current_param.default} (from block '{value_sources[input_name]}') and "
+ f"{input_param.default} (from block '{block_name}'). Using {current_param.default}."
+ )
+ if current_param.default is None and input_param.default is not None:
+ combined_dict[input_name] = input_param
+ value_sources[input_name] = block_name
+ else:
+ combined_dict[input_name] = input_param
+ value_sources[input_name] = block_name
+
+ return list(combined_dict.values())
+
+
+def combine_outputs(*named_output_lists: list[tuple[str, list[OutputParam]]]) -> list[OutputParam]:
+ """
+ Combines multiple lists of OutputParam objects from different blocks. For duplicate outputs, keeps the first
+ occurrence of each output name.
+
+ Args:
+ named_output_lists: List of tuples containing (block_name, output_param_list) pairs
+
+ Returns:
+ List[OutputParam]: Combined list of unique OutputParam objects
+ """
+ combined_dict = {} # name -> OutputParam
+
+ for block_name, outputs in named_output_lists:
+ for output_param in outputs:
+ if (output_param.name not in combined_dict) or (
+ combined_dict[output_param.name].kwargs_type is None and output_param.kwargs_type is not None
+ ):
+ combined_dict[output_param.name] = output_param
+
+ return list(combined_dict.values())
+
+
+def generate_modular_model_card_content(blocks) -> dict[str, Any]:
"""
Generate model card content for a modular pipeline.
diff --git a/src/diffusers/modular_pipelines/qwenimage/__init__.py b/src/diffusers/modular_pipelines/qwenimage/__init__.py
index 2b01a5b5a4b5..2e6af4495b37 100644
--- a/src/diffusers/modular_pipelines/qwenimage/__init__.py
+++ b/src/diffusers/modular_pipelines/qwenimage/__init__.py
@@ -21,27 +21,15 @@
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
- _import_structure["modular_blocks_qwenimage"] = [
- "AUTO_BLOCKS",
- "QwenImageAutoBlocks",
- ]
- _import_structure["modular_blocks_qwenimage_edit"] = [
- "EDIT_AUTO_BLOCKS",
- "QwenImageEditAutoBlocks",
- ]
- _import_structure["modular_blocks_qwenimage_edit_plus"] = [
- "EDIT_PLUS_AUTO_BLOCKS",
- "QwenImageEditPlusAutoBlocks",
- ]
- _import_structure["modular_blocks_qwenimage_layered"] = [
- "LAYERED_AUTO_BLOCKS",
- "QwenImageLayeredAutoBlocks",
- ]
+ _import_structure["modular_blocks_qwenimage"] = ["QwenImageAutoBlocks"]
+ _import_structure["modular_blocks_qwenimage_edit"] = ["QwenImageEditAutoBlocks"]
+ _import_structure["modular_blocks_qwenimage_edit_plus"] = ["QwenImageEditPlusAutoBlocks"]
+ _import_structure["modular_blocks_qwenimage_layered"] = ["QwenImageLayeredAutoBlocks"]
_import_structure["modular_pipeline"] = [
"QwenImageEditModularPipeline",
"QwenImageEditPlusModularPipeline",
- "QwenImageModularPipeline",
"QwenImageLayeredModularPipeline",
+ "QwenImageModularPipeline",
]
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
@@ -51,22 +39,10 @@
except OptionalDependencyNotAvailable:
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
else:
- from .modular_blocks_qwenimage import (
- AUTO_BLOCKS,
- QwenImageAutoBlocks,
- )
- from .modular_blocks_qwenimage_edit import (
- EDIT_AUTO_BLOCKS,
- QwenImageEditAutoBlocks,
- )
- from .modular_blocks_qwenimage_edit_plus import (
- EDIT_PLUS_AUTO_BLOCKS,
- QwenImageEditPlusAutoBlocks,
- )
- from .modular_blocks_qwenimage_layered import (
- LAYERED_AUTO_BLOCKS,
- QwenImageLayeredAutoBlocks,
- )
+ from .modular_blocks_qwenimage import QwenImageAutoBlocks
+ from .modular_blocks_qwenimage_edit import QwenImageEditAutoBlocks
+ from .modular_blocks_qwenimage_edit_plus import QwenImageEditPlusAutoBlocks
+ from .modular_blocks_qwenimage_layered import QwenImageLayeredAutoBlocks
from .modular_pipeline import (
QwenImageEditModularPipeline,
QwenImageEditPlusModularPipeline,
diff --git a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py
index 80a379da6be0..51b5c6ac8c3d 100644
--- a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py
+++ b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Tuple, Union
import numpy as np
import torch
@@ -43,10 +42,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -61,15 +60,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -159,13 +158,13 @@ def description(self) -> str:
return "Prepare initial random noise for the generation process"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("latents"),
InputParam.template("height"),
@@ -177,7 +176,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(name="height", type_hint=int, description="if not set, updated to default value"),
OutputParam(name="width", type_hint=int, description="if not set, updated to default value"),
@@ -277,13 +276,13 @@ def description(self) -> str:
return "Prepare initial random noise (B, layers+1, C, H, W) for the generation process"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("latents"),
InputParam.template("height"),
@@ -296,7 +295,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(name="height", type_hint=int, description="if not set, updated to default value"),
OutputParam(name="width", type_hint=int, description="if not set, updated to default value"),
@@ -385,13 +384,13 @@ def description(self) -> str:
return "Step that adds noise to image latents for image-to-image/inpainting. Should be run after set_timesteps, prepare_latents. Both noise and image latents should alreadybe patchified."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="latents",
@@ -409,7 +408,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="initial_noise",
@@ -488,13 +487,13 @@ def description(self) -> str:
return "Step that creates mask latents from preprocessed mask_image by interpolating to latent space."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="processed_mask_image",
@@ -508,7 +507,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="mask", type_hint=torch.Tensor, description="The mask to use for the inpainting process."
@@ -559,7 +558,7 @@ class QwenImageSetTimestepsStep(ModularPipelineBlocks):
Inputs:
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
latents (`Tensor`):
The initial random noised latents for the denoising process. Can be generated in prepare latents step.
@@ -576,13 +575,13 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for text-to-image generation. Should be run after prepare latents step."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("num_inference_steps"),
InputParam.template("sigmas"),
@@ -595,7 +594,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process"
@@ -645,7 +644,7 @@ class QwenImageLayeredSetTimestepsStep(ModularPipelineBlocks):
Inputs:
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
image_latents (`Tensor`):
image latents used to guide the image generation. Can be generated from vae_encoder step.
@@ -662,13 +661,13 @@ def description(self) -> str:
return "Set timesteps step for QwenImage Layered with custom mu calculation based on image_latents."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("num_inference_steps"),
InputParam.template("sigmas"),
@@ -676,7 +675,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process."
@@ -726,7 +725,7 @@ class QwenImageSetTimestepsWithStrengthStep(ModularPipelineBlocks):
Inputs:
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
latents (`Tensor`):
The latents to use for the denoising process. Can be generated in prepare latents step.
@@ -747,13 +746,13 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for image-to-image generation, and inpainting. Should be run after prepare latents step."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("num_inference_steps"),
InputParam.template("sigmas"),
@@ -767,7 +766,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="timesteps",
@@ -843,7 +842,7 @@ class QwenImageRoPEInputsStep(ModularPipelineBlocks):
mask for the negative text embeddings. Can be generated from text_encoder step.
Outputs:
- img_shapes (`List`):
+ img_shapes (`list`):
The shapes of the images latents, used for RoPE calculation
"""
@@ -856,7 +855,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("batch_size"),
InputParam.template("height", required=True),
@@ -866,12 +865,12 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="img_shapes",
kwargs_type="denoiser_input_fields",
- type_hint=List[List[Tuple[int, int, int]]],
+ type_hint=list[list[tuple[int, int, int]]],
description="The shapes of the images latents, used for RoPE calculation",
),
]
@@ -918,7 +917,7 @@ class QwenImageEditRoPEInputsStep(ModularPipelineBlocks):
mask for the negative text embeddings. Can be generated from text_encoder step.
Outputs:
- img_shapes (`List`):
+ img_shapes (`list`):
The shapes of the images latents, used for RoPE calculation
"""
@@ -929,7 +928,7 @@ def description(self) -> str:
return "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be placed after prepare_latents step"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("batch_size"),
InputParam(
@@ -951,12 +950,12 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="img_shapes",
kwargs_type="denoiser_input_fields",
- type_hint=List[List[Tuple[int, int, int]]],
+ type_hint=list[list[tuple[int, int, int]]],
description="The shapes of the images latents, used for RoPE calculation",
),
]
@@ -996,9 +995,9 @@ class QwenImageEditPlusRoPEInputsStep(ModularPipelineBlocks):
batch_size (`int`, *optional*, defaults to 1):
Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
be generated in input step.
- image_height (`List`):
+ image_height (`list`):
The heights of the reference images. Can be generated in input step.
- image_width (`List`):
+ image_width (`list`):
The widths of the reference images. Can be generated in input step.
height (`int`):
The height in pixels of the generated image.
@@ -1010,11 +1009,11 @@ class QwenImageEditPlusRoPEInputsStep(ModularPipelineBlocks):
mask for the negative text embeddings. Can be generated from text_encoder step.
Outputs:
- img_shapes (`List`):
+ img_shapes (`list`):
The shapes of the image latents, used for RoPE calculation
- txt_seq_lens (`List`):
+ txt_seq_lens (`list`):
The sequence lengths of the prompt embeds, used for RoPE calculation
- negative_txt_seq_lens (`List`):
+ negative_txt_seq_lens (`list`):
The sequence lengths of the negative prompt embeds, used for RoPE calculation
"""
@@ -1029,19 +1028,19 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("batch_size"),
InputParam(
name="image_height",
required=True,
- type_hint=List[int],
+ type_hint=list[int],
description="The heights of the reference images. Can be generated in input step.",
),
InputParam(
name="image_width",
required=True,
- type_hint=List[int],
+ type_hint=list[int],
description="The widths of the reference images. Can be generated in input step.",
),
InputParam.template("height", required=True),
@@ -1051,24 +1050,24 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="img_shapes",
kwargs_type="denoiser_input_fields",
- type_hint=List[List[Tuple[int, int, int]]],
+ type_hint=list[list[tuple[int, int, int]]],
description="The shapes of the image latents, used for RoPE calculation",
),
OutputParam(
name="txt_seq_lens",
kwargs_type="denoiser_input_fields",
- type_hint=List[int],
+ type_hint=list[int],
description="The sequence lengths of the prompt embeds, used for RoPE calculation",
),
OutputParam(
name="negative_txt_seq_lens",
kwargs_type="denoiser_input_fields",
- type_hint=List[int],
+ type_hint=list[int],
description="The sequence lengths of the negative prompt embeds, used for RoPE calculation",
),
]
@@ -1124,11 +1123,11 @@ class QwenImageLayeredRoPEInputsStep(ModularPipelineBlocks):
mask for the negative text embeddings. Can be generated from text_encoder step.
Outputs:
- img_shapes (`List`):
+ img_shapes (`list`):
The shapes of the image latents, used for RoPE calculation
- txt_seq_lens (`List`):
+ txt_seq_lens (`list`):
The sequence lengths of the prompt embeds, used for RoPE calculation
- negative_txt_seq_lens (`List`):
+ negative_txt_seq_lens (`list`):
The sequence lengths of the negative prompt embeds, used for RoPE calculation
additional_t_cond (`Tensor`):
The additional t cond, used for RoPE calculation
@@ -1143,7 +1142,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("batch_size"),
InputParam.template("layers"),
@@ -1154,23 +1153,23 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="img_shapes",
- type_hint=List[List[Tuple[int, int, int]]],
+ type_hint=list[list[tuple[int, int, int]]],
kwargs_type="denoiser_input_fields",
description="The shapes of the image latents, used for RoPE calculation",
),
OutputParam(
name="txt_seq_lens",
- type_hint=List[int],
+ type_hint=list[int],
kwargs_type="denoiser_input_fields",
description="The sequence lengths of the prompt embeds, used for RoPE calculation",
),
OutputParam(
name="negative_txt_seq_lens",
- type_hint=List[int],
+ type_hint=list[int],
kwargs_type="denoiser_input_fields",
description="The sequence lengths of the negative prompt embeds, used for RoPE calculation",
),
@@ -1239,14 +1238,14 @@ class QwenImageControlNetBeforeDenoiserStep(ModularPipelineBlocks):
The timesteps to use for the denoising process. Can be generated in set_timesteps step.
Outputs:
- controlnet_keep (`List`):
+ controlnet_keep (`list`):
The controlnet keep values
"""
model_name = "qwenimage"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("controlnet", QwenImageControlNetModel),
]
@@ -1256,7 +1255,7 @@ def description(self) -> str:
return "step that prepare inputs for controlnet. Insert before the Denoise Step, after set_timesteps step."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("control_guidance_start"),
InputParam.template("control_guidance_end"),
@@ -1276,9 +1275,9 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
- OutputParam("controlnet_keep", type_hint=List[float], description="The controlnet keep values"),
+ OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"),
]
@torch.no_grad()
diff --git a/src/diffusers/modular_pipelines/qwenimage/decoders.py b/src/diffusers/modular_pipelines/qwenimage/decoders.py
index 1adbf6bdd355..e4ccb6b8e047 100644
--- a/src/diffusers/modular_pipelines/qwenimage/decoders.py
+++ b/src/diffusers/modular_pipelines/qwenimage/decoders.py
@@ -12,7 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List
+
+from typing import Any
import torch
@@ -60,7 +61,7 @@ def description(self) -> str:
return "Step that unpack the latents from 3D tensor (batch_size, sequence_length, channels) into 5D tensor (batch_size, channels, 1, height, width)"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
components = [
ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
]
@@ -68,7 +69,7 @@ def expected_components(self) -> List[ComponentSpec]:
return components
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("height", required=True),
InputParam.template("width", required=True),
@@ -81,7 +82,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="latents", type_hint=torch.Tensor, description="The denoisedlatents unpacked to B, C, 1, H, W"
@@ -131,13 +132,13 @@ def description(self) -> str:
return "Unpack latents from (B, seq, C*4) to (B, C, layers+1, H, W) after denoising."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="latents",
@@ -151,7 +152,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam.template("latents", note="unpacked to B, C, layers+1, H, W"),
]
@@ -190,7 +191,7 @@ class QwenImageDecoderStep(ModularPipelineBlocks):
step.
Outputs:
- images (`List`):
+ images (`list`):
Generated images. (tensor output of the vae decoder.)
"""
@@ -201,7 +202,7 @@ def description(self) -> str:
return "Step that decodes the latents to images"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
components = [
ComponentSpec("vae", AutoencoderKLQwenImage),
]
@@ -209,7 +210,7 @@ def expected_components(self) -> List[ComponentSpec]:
return components
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="latents",
@@ -220,7 +221,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam.template("images", note="tensor output of the vae decoder.")]
@torch.no_grad()
@@ -267,7 +268,7 @@ class QwenImageLayeredDecoderStep(ModularPipelineBlocks):
Output format: 'pil', 'np', 'pt'.
Outputs:
- images (`List`):
+ images (`list`):
Generated images.
"""
@@ -278,7 +279,7 @@ def description(self) -> str:
return "Decode unpacked latents (B, C, layers+1, H, W) into layer images."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKLQwenImage),
ComponentSpec(
@@ -290,7 +291,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="latents",
@@ -302,10 +303,8 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
- return [
- OutputParam.template("images"),
- ]
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
@@ -367,7 +366,7 @@ class QwenImageProcessImagesOutputStep(ModularPipelineBlocks):
Output format: 'pil', 'np', 'pt'.
Outputs:
- images (`List`):
+ images (`list`):
Generated images.
"""
@@ -378,7 +377,7 @@ def description(self) -> str:
return "postprocess the generated image"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
@@ -389,7 +388,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="images",
@@ -401,7 +400,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam.template("images")]
@staticmethod
@@ -437,12 +436,12 @@ class QwenImageInpaintProcessImagesOutputStep(ModularPipelineBlocks):
the generated image tensor from decoders step
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
- mask_overlay_kwargs (`Dict`, *optional*):
+ mask_overlay_kwargs (`dict`, *optional*):
The kwargs for the postprocess step to apply the mask overlay. generated in
InpaintProcessImagesInputStep.
Outputs:
- images (`List`):
+ images (`list`):
Generated images.
"""
@@ -453,7 +452,7 @@ def description(self) -> str:
return "postprocess the generated image, optional apply the mask overally to the original image.."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_mask_processor",
@@ -464,7 +463,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="images",
@@ -475,13 +474,13 @@ def inputs(self) -> List[InputParam]:
InputParam.template("output_type"),
InputParam(
name="mask_overlay_kwargs",
- type_hint=Dict[str, Any],
+ type_hint=dict[str, Any],
description="The kwargs for the postprocess step to apply the mask overlay. generated in InpaintProcessImagesInputStep.",
),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam.template("images")]
@staticmethod
diff --git a/src/diffusers/modular_pipelines/qwenimage/denoise.py b/src/diffusers/modular_pipelines/qwenimage/denoise.py
index 8579c9843a89..de8ea05c5047 100644
--- a/src/diffusers/modular_pipelines/qwenimage/denoise.py
+++ b/src/diffusers/modular_pipelines/qwenimage/denoise.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Tuple
import torch
@@ -47,7 +46,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="latents",
@@ -77,7 +76,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="latents",
@@ -101,7 +100,7 @@ class QwenImageLoopBeforeDenoiserControlNet(ModularPipelineBlocks):
model_name = "qwenimage"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -121,7 +120,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
"control_image_latents",
@@ -133,7 +132,7 @@ def inputs(self) -> List[InputParam]:
InputParam(
name="controlnet_keep",
required=True,
- type_hint=List[float],
+ type_hint=list[float],
description="The controlnet keep values. Can be generated in prepare_controlnet_inputs step.",
),
]
@@ -181,7 +180,7 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -193,14 +192,14 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("attention_kwargs"),
InputParam.template("denoiser_input_fields"),
InputParam(
"img_shapes",
required=True,
- type_hint=List[Tuple[int, int]],
+ type_hint=list[tuple[int, int]],
description="The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.",
),
]
@@ -266,7 +265,7 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -278,14 +277,14 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("attention_kwargs"),
InputParam.template("denoiser_input_fields"),
InputParam(
"img_shapes",
required=True,
- type_hint=List[Tuple[int, int]],
+ type_hint=list[tuple[int, int]],
description="The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.",
),
]
@@ -355,13 +354,13 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam.template("latents"),
]
@@ -396,7 +395,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
"mask",
@@ -414,7 +413,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam.template("latents"),
]
@@ -449,13 +448,13 @@ def description(self) -> str:
)
@property
- def loop_expected_components(self) -> List[ComponentSpec]:
+ def loop_expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@property
- def loop_inputs(self) -> List[InputParam]:
+ def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
name="timesteps",
@@ -519,11 +518,11 @@ class QwenImageDenoiseStep(QwenImageDenoiseLoopWrapper):
The number of denoising steps.
latents (`Tensor`):
The initial latents to use for the denoising process. Can be generated in prepare_latent step.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
- img_shapes (`List`):
+ img_shapes (`list`):
The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.
Outputs:
@@ -577,11 +576,11 @@ class QwenImageInpaintDenoiseStep(QwenImageDenoiseLoopWrapper):
The number of denoising steps.
latents (`Tensor`):
The initial latents to use for the denoising process. Can be generated in prepare_latent step.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
- img_shapes (`List`):
+ img_shapes (`list`):
The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.
mask (`Tensor`):
The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.
@@ -646,13 +645,13 @@ class QwenImageControlNetDenoiseStep(QwenImageDenoiseLoopWrapper):
The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.
controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
Scale for ControlNet conditioning. (updated in prepare_controlnet_inputs step.)
- controlnet_keep (`List`):
+ controlnet_keep (`list`):
The controlnet keep values. Can be generated in prepare_controlnet_inputs step.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
- img_shapes (`List`):
+ img_shapes (`list`):
The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.
Outputs:
@@ -712,13 +711,13 @@ class QwenImageInpaintControlNetDenoiseStep(QwenImageDenoiseLoopWrapper):
The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.
controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
Scale for ControlNet conditioning. (updated in prepare_controlnet_inputs step.)
- controlnet_keep (`List`):
+ controlnet_keep (`list`):
The controlnet keep values. Can be generated in prepare_controlnet_inputs step.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
- img_shapes (`List`):
+ img_shapes (`list`):
The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.
mask (`Tensor`):
The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.
@@ -788,11 +787,11 @@ class QwenImageEditDenoiseStep(QwenImageDenoiseLoopWrapper):
The initial latents to use for the denoising process. Can be generated in prepare_latent step.
image_latents (`Tensor`):
image latents used to guide the image generation. Can be generated from vae_encoder step.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
- img_shapes (`List`):
+ img_shapes (`list`):
The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.
Outputs:
@@ -847,11 +846,11 @@ class QwenImageEditInpaintDenoiseStep(QwenImageDenoiseLoopWrapper):
The initial latents to use for the denoising process. Can be generated in prepare_latent step.
image_latents (`Tensor`):
image latents used to guide the image generation. Can be generated from vae_encoder step.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
- img_shapes (`List`):
+ img_shapes (`list`):
The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.
mask (`Tensor`):
The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.
@@ -911,11 +910,11 @@ class QwenImageLayeredDenoiseStep(QwenImageDenoiseLoopWrapper):
The initial latents to use for the denoising process. Can be generated in prepare_latent step.
image_latents (`Tensor`):
image latents used to guide the image generation. Can be generated from vae_encoder step.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
- img_shapes (`List`):
+ img_shapes (`list`):
The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.
Outputs:
diff --git a/src/diffusers/modular_pipelines/qwenimage/encoders.py b/src/diffusers/modular_pipelines/qwenimage/encoders.py
index 5e1821cca5c0..527267dc0d6e 100644
--- a/src/diffusers/modular_pipelines/qwenimage/encoders.py
+++ b/src/diffusers/modular_pipelines/qwenimage/encoders.py
@@ -16,8 +16,6 @@
Text and VAE encoder blocks for QwenImage pipelines.
"""
-from typing import Dict, List, Optional, Union
-
import PIL
import torch
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
@@ -59,11 +57,11 @@ def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor):
def get_qwen_prompt_embeds(
text_encoder,
tokenizer,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
prompt_template_encode: str = QWENIMAGE_PROMPT_TEMPLATE,
prompt_template_encode_start_idx: int = QWENIMAGE_PROMPT_TEMPLATE_START_IDX,
tokenizer_max_length: int = 1024,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -99,11 +97,11 @@ def get_qwen_prompt_embeds(
def get_qwen_prompt_embeds_edit(
text_encoder,
processor,
- prompt: Union[str, List[str]] = None,
- image: Optional[torch.Tensor] = None,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | None = None,
prompt_template_encode: str = QWENIMAGE_EDIT_PROMPT_TEMPLATE,
prompt_template_encode_start_idx: int = QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -146,12 +144,12 @@ def get_qwen_prompt_embeds_edit(
def get_qwen_prompt_embeds_edit_plus(
text_encoder,
processor,
- prompt: Union[str, List[str]] = None,
- image: Optional[Union[torch.Tensor, List[PIL.Image.Image], PIL.Image.Image]] = None,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | list[PIL.Image.Image, PIL.Image.Image] | None = None,
prompt_template_encode: str = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE,
img_template_encode: str = QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE,
prompt_template_encode_start_idx: int = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
if isinstance(image, list):
@@ -200,7 +198,7 @@ def get_qwen_prompt_embeds_edit_plus(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -287,11 +285,11 @@ class QwenImageEditResizeStep(ModularPipelineBlocks):
image_resize_processor (`VaeImageProcessor`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized images
"""
@@ -302,7 +300,7 @@ def description(self) -> str:
return "Image Resize step that resize the image to target area while maintaining the aspect ratio."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_resize_processor",
@@ -313,15 +311,15 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [InputParam.template("image")]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="resized_image",
- type_hint=List[PIL.Image.Image],
+ type_hint=list[PIL.Image.Image],
description="The resized images",
),
]
@@ -361,13 +359,13 @@ class QwenImageLayeredResizeStep(ModularPipelineBlocks):
image_resize_processor (`VaeImageProcessor`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
resolution (`int`, *optional*, defaults to 640):
The target area to resize the image to, can be 1024 or 640
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized images
"""
@@ -378,7 +376,7 @@ def description(self) -> str:
return "Image Resize step that resize the image to a target area (defined by the resolution parameter from user) while maintaining the aspect ratio."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_resize_processor",
@@ -389,7 +387,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("image"),
InputParam(
@@ -401,11 +399,11 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="resized_image",
- type_hint=List[PIL.Image.Image],
+ type_hint=list[PIL.Image.Image],
description="The resized images",
)
]
@@ -454,13 +452,13 @@ class QwenImageEditPlusResizeStep(ModularPipelineBlocks):
image_resize_processor (`VaeImageProcessor`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
Images resized to 1024x1024 target area for VAE encoding
- resized_cond_image (`List`):
+ resized_cond_image (`list`):
Images resized to 384x384 target area for VL text encoding
"""
@@ -476,7 +474,7 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_resize_processor",
@@ -487,21 +485,21 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
# image
return [InputParam.template("image")]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="resized_image",
- type_hint=List[PIL.Image.Image],
+ type_hint=list[PIL.Image.Image],
description="Images resized to 1024x1024 target area for VAE encoding",
),
OutputParam(
name="resized_cond_image",
- type_hint=List[PIL.Image.Image],
+ type_hint=list[PIL.Image.Image],
description="Images resized to 384x384 target area for VL text encoding",
),
]
@@ -584,14 +582,14 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
ComponentSpec("processor", Qwen2VLProcessor),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template(
"prompt", required=False
@@ -611,7 +609,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="prompt",
@@ -700,7 +698,7 @@ def description(self) -> str:
return "Text Encoder step that generates text embeddings to guide the image generation."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration, description="The text encoder to use"),
ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer to use"),
@@ -713,7 +711,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("prompt"),
InputParam.template("negative_prompt"),
@@ -721,7 +719,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam.template("prompt_embeds"),
OutputParam.template("prompt_embeds_mask"),
@@ -829,7 +827,7 @@ def description(self) -> str:
return "Text Encoder step that processes both prompt and image together to generate text embeddings for guiding image generation."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
ComponentSpec("processor", Qwen2VLProcessor),
@@ -842,7 +840,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("prompt"),
InputParam.template("negative_prompt"),
@@ -855,7 +853,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam.template("prompt_embeds"),
OutputParam.template("prompt_embeds_mask"),
@@ -957,7 +955,7 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
ComponentSpec("processor", Qwen2VLProcessor),
@@ -970,7 +968,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("prompt"),
InputParam.template("negative_prompt"),
@@ -983,7 +981,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam.template("prompt_embeds"),
OutputParam.template("prompt_embeds_mask"),
@@ -1060,7 +1058,7 @@ class QwenImageInpaintProcessImagesInputStep(ModularPipelineBlocks):
Inputs:
mask_image (`Image`):
Mask image for inpainting.
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
height (`int`, *optional*):
The height in pixels of the generated image.
@@ -1074,7 +1072,7 @@ class QwenImageInpaintProcessImagesInputStep(ModularPipelineBlocks):
The processed image
processed_mask_image (`Tensor`):
The processed mask image
- mask_overlay_kwargs (`Dict`):
+ mask_overlay_kwargs (`dict`):
The kwargs for the postprocess step to apply the mask overlay
"""
@@ -1085,7 +1083,7 @@ def description(self) -> str:
return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images will be resized to the given height and width."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_mask_processor",
@@ -1096,7 +1094,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("mask_image"),
InputParam.template("image"),
@@ -1106,7 +1104,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="processed_image",
@@ -1120,7 +1118,7 @@ def intermediate_outputs(self) -> List[OutputParam]:
),
OutputParam(
name="mask_overlay_kwargs",
- type_hint=Dict,
+ type_hint=dict,
description="The kwargs for the postprocess step to apply the mask overlay",
),
]
@@ -1179,7 +1177,7 @@ class QwenImageEditInpaintProcessImagesInputStep(ModularPipelineBlocks):
The processed image
processed_mask_image (`Tensor`):
The processed mask image
- mask_overlay_kwargs (`Dict`):
+ mask_overlay_kwargs (`dict`):
The kwargs for the postprocess step to apply the mask overlay
"""
@@ -1190,7 +1188,7 @@ def description(self) -> str:
return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images should be resized first."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_mask_processor",
@@ -1201,7 +1199,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("mask_image"),
InputParam(
@@ -1214,7 +1212,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(name="processed_image", type_hint=torch.Tensor, description="The processed image"),
OutputParam(
@@ -1224,7 +1222,7 @@ def intermediate_outputs(self) -> List[OutputParam]:
),
OutputParam(
name="mask_overlay_kwargs",
- type_hint=Dict,
+ type_hint=dict,
description="The kwargs for the postprocess step to apply the mask overlay",
),
]
@@ -1258,7 +1256,7 @@ class QwenImageProcessImagesInputStep(ModularPipelineBlocks):
image_processor (`VaeImageProcessor`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
height (`int`, *optional*):
The height in pixels of the generated image.
@@ -1277,7 +1275,7 @@ def description(self) -> str:
return "Image Preprocess step. will resize the image to the given height and width."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
@@ -1288,7 +1286,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("image"),
InputParam.template("height"),
@@ -1296,7 +1294,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="processed_image",
@@ -1342,7 +1340,7 @@ class QwenImageEditProcessImagesInputStep(ModularPipelineBlocks):
image_processor (`VaeImageProcessor`)
Inputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized image. should be generated using a resize step
Outputs:
@@ -1357,7 +1355,7 @@ def description(self) -> str:
return "Image Preprocess step. Images needs to be resized first."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
@@ -1368,18 +1366,18 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="resized_image",
required=True,
- type_hint=List[PIL.Image.Image],
+ type_hint=list[PIL.Image.Image],
description="The resized image. should be generated using a resize step",
),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="processed_image",
@@ -1414,7 +1412,7 @@ class QwenImageEditPlusProcessImagesInputStep(ModularPipelineBlocks):
image_processor (`VaeImageProcessor`)
Inputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized image. should be generated using a resize step
Outputs:
@@ -1429,7 +1427,7 @@ def description(self) -> str:
return "Image Preprocess step. Images can be resized first. If a list of images is provided, will return a list of processed images."
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
@@ -1440,18 +1438,18 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="resized_image",
required=True,
- type_hint=List[PIL.Image.Image],
+ type_hint=list[PIL.Image.Image],
description="The resized image. should be generated using a resize step",
)
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="processed_image",
@@ -1513,11 +1511,7 @@ class QwenImageVaeEncoderStep(ModularPipelineBlocks):
model_name = "qwenimage"
- def __init__(
- self,
- input: Optional[InputParam] = None,
- output: Optional[OutputParam] = None,
- ):
+ def __init__(self, input: InputParam | None = None, output: OutputParam | None = None):
"""Initialize a VAE encoder step for converting images to latent representations.
Handles both single images and lists of images. When input is a list, outputs a list of latents. When input is
@@ -1554,18 +1548,18 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("vae", AutoencoderKLQwenImage)]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
self._input, # default is "processed_image"
InputParam.template("generator"),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [self._output] # default is "image_latents"
@torch.no_grad()
@@ -1634,7 +1628,7 @@ def description(self) -> str:
return "VAE Encoder step that converts `control_image` into latent representations control_image_latents.\n"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
components = [
ComponentSpec("vae", AutoencoderKLQwenImage),
ComponentSpec("controlnet", QwenImageControlNetModel),
@@ -1648,7 +1642,7 @@ def expected_components(self) -> List[ComponentSpec]:
return components
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [
InputParam.template("control_image"),
InputParam.template("height"),
@@ -1658,7 +1652,7 @@ def inputs(self) -> List[InputParam]:
return inputs
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"control_image_latents",
@@ -1763,13 +1757,13 @@ def description(self) -> str:
return "Permute image latents from (B, C, 1, H, W) to (B, 1, C, H, W) for Layered packing."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("image_latents"),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam.template("image_latents", note="permuted from [B, C, 1, H, W] to [B, 1, C, H, W]"),
]
diff --git a/src/diffusers/modular_pipelines/qwenimage/inputs.py b/src/diffusers/modular_pipelines/qwenimage/inputs.py
index 818bbca5ed0a..faec7db245df 100644
--- a/src/diffusers/modular_pipelines/qwenimage/inputs.py
+++ b/src/diffusers/modular_pipelines/qwenimage/inputs.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple
import torch
@@ -78,7 +77,7 @@ def repeat_tensor_to_batch_size(
return input_tensor
-def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: int) -> Tuple[int, int]:
+def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: int) -> tuple[int, int]:
"""Calculate image dimensions from latent tensor dimensions.
This function converts latent space dimensions to image space dimensions by multiplying the latent height and width
@@ -91,7 +90,7 @@ def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: in
Typically 8 for most VAEs (image is 8x larger than latents in each dimension)
Returns:
- Tuple[int, int]: The calculated image dimensions as (height, width)
+ tuple[int, int]: The calculated image dimensions as (height, width)
Raises:
ValueError: If latents tensor doesn't have 4 or 5 dimensions
@@ -164,7 +163,7 @@ def description(self) -> str:
return summary_section + placement_section
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam.template("num_images_per_prompt"),
InputParam.template("prompt_embeds"),
@@ -174,7 +173,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(name="batch_size", type_hint=int, description="The batch size of the prompt embeddings"),
OutputParam(name="dtype", type_hint=torch.dtype, description="The data type of the prompt embeddings"),
@@ -300,8 +299,8 @@ class QwenImageAdditionalInputsStep(ModularPipelineBlocks):
def __init__(
self,
- image_latent_inputs: Optional[List[InputParam]] = None,
- additional_batch_inputs: Optional[List[InputParam]] = None,
+ image_latent_inputs: list[InputParam] | None = None,
+ additional_batch_inputs: list[InputParam] | None = None,
):
# by default, process `image_latents`
if image_latent_inputs is None:
@@ -350,13 +349,13 @@ def description(self) -> str:
return summary_section + inputs_info + placement_section
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [
InputParam.template("num_images_per_prompt"),
InputParam.template("batch_size"),
@@ -369,7 +368,7 @@ def inputs(self) -> List[InputParam]:
return inputs
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
outputs = [
OutputParam(
name="image_height",
@@ -497,9 +496,9 @@ class QwenImageEditPlusAdditionalInputsStep(ModularPipelineBlocks):
image latents used to guide the image generation. Can be generated from vae_encoder step.
Outputs:
- image_height (`List`):
+ image_height (`list`):
The image heights calculated from the image latents dimension
- image_width (`List`):
+ image_width (`list`):
The image widths calculated from the image latents dimension
height (`int`):
if not provided, updated to image height
@@ -514,8 +513,8 @@ class QwenImageEditPlusAdditionalInputsStep(ModularPipelineBlocks):
def __init__(
self,
- image_latent_inputs: Optional[List[InputParam]] = None,
- additional_batch_inputs: Optional[List[InputParam]] = None,
+ image_latent_inputs: list[InputParam] | None = None,
+ additional_batch_inputs: list[InputParam] | None = None,
):
if image_latent_inputs is None:
image_latent_inputs = [InputParam.template("image_latents")]
@@ -564,13 +563,13 @@ def description(self) -> str:
return summary_section + inputs_info + placement_section
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [
InputParam.template("num_images_per_prompt"),
InputParam.template("batch_size"),
@@ -584,16 +583,16 @@ def inputs(self) -> List[InputParam]:
return inputs
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
outputs = [
OutputParam(
name="image_height",
- type_hint=List[int],
+ type_hint=list[int],
description="The image heights calculated from the image latents dimension",
),
OutputParam(
name="image_width",
- type_hint=List[int],
+ type_hint=list[int],
description="The image widths calculated from the image latents dimension",
),
]
@@ -744,8 +743,8 @@ class QwenImageLayeredAdditionalInputsStep(ModularPipelineBlocks):
def __init__(
self,
- image_latent_inputs: Optional[List[InputParam]] = None,
- additional_batch_inputs: Optional[List[InputParam]] = None,
+ image_latent_inputs: list[InputParam] | None = None,
+ additional_batch_inputs: list[InputParam] | None = None,
):
if image_latent_inputs is None:
image_latent_inputs = [InputParam.template("image_latents")]
@@ -793,13 +792,13 @@ def description(self) -> str:
return summary_section + inputs_info + placement_section
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [
InputParam.template("num_images_per_prompt"),
InputParam.template("batch_size"),
@@ -811,7 +810,7 @@ def inputs(self) -> List[InputParam]:
return inputs
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
outputs = [
OutputParam(
name="image_height",
@@ -945,7 +944,7 @@ def description(self) -> str:
return "prepare the `control_image_latents` for controlnet. Insert after all the other inputs steps."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
name="control_image_latents",
@@ -960,7 +959,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="control_image_latents",
diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py
index 5837799d3431..bf87028b2f90 100644
--- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py
+++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py
@@ -119,7 +119,7 @@ class QwenImageInpaintVaeEncoderStep(SequentialPipelineBlocks):
Inputs:
mask_image (`Image`):
Mask image for inpainting.
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
height (`int`, *optional*):
The height in pixels of the generated image.
@@ -135,7 +135,7 @@ class QwenImageInpaintVaeEncoderStep(SequentialPipelineBlocks):
The processed image
processed_mask_image (`Tensor`):
The processed mask image
- mask_overlay_kwargs (`Dict`):
+ mask_overlay_kwargs (`dict`):
The kwargs for the postprocess step to apply the mask overlay
image_latents (`Tensor`):
The latent representation of the input image.
@@ -164,7 +164,7 @@ class QwenImageImg2ImgVaeEncoderStep(SequentialPipelineBlocks):
image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
height (`int`, *optional*):
The height in pixels of the generated image.
@@ -476,9 +476,9 @@ class QwenImageCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -553,11 +553,11 @@ class QwenImageInpaintCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
strength (`float`, *optional*, defaults to 0.9):
Strength for img2img/inpainting.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -632,11 +632,11 @@ class QwenImageImg2ImgCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
strength (`float`, *optional*, defaults to 0.9):
Strength for img2img/inpainting.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -712,7 +712,7 @@ class QwenImageControlNetCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
control_guidance_start (`float`, *optional*, defaults to 0.0):
When to start applying ControlNet.
@@ -720,7 +720,7 @@ class QwenImageControlNetCoreDenoiseStep(SequentialPipelineBlocks):
When to stop applying ControlNet.
controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
Scale for ControlNet conditioning.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -802,7 +802,7 @@ class QwenImageControlNetInpaintCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
strength (`float`, *optional*, defaults to 0.9):
Strength for img2img/inpainting.
@@ -812,7 +812,7 @@ class QwenImageControlNetInpaintCoreDenoiseStep(SequentialPipelineBlocks):
When to stop applying ControlNet.
controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
Scale for ControlNet conditioning.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -894,7 +894,7 @@ class QwenImageControlNetImg2ImgCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
strength (`float`, *optional*, defaults to 0.9):
Strength for img2img/inpainting.
@@ -904,7 +904,7 @@ class QwenImageControlNetImg2ImgCoreDenoiseStep(SequentialPipelineBlocks):
When to stop applying ControlNet.
controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
Scale for ControlNet conditioning.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -1032,7 +1032,7 @@ class QwenImageDecodeStep(SequentialPipelineBlocks):
Output format: 'pil', 'np', 'pt'.
Outputs:
- images (`List`):
+ images (`list`):
Generated images. (tensor output of the vae decoder.)
"""
@@ -1061,12 +1061,12 @@ class QwenImageInpaintDecodeStep(SequentialPipelineBlocks):
step.
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
- mask_overlay_kwargs (`Dict`, *optional*):
+ mask_overlay_kwargs (`dict`, *optional*):
The kwargs for the postprocess step to apply the mask overlay. generated in
InpaintProcessImagesInputStep.
Outputs:
- images (`List`):
+ images (`list`):
Generated images. (tensor output of the vae decoder.)
"""
@@ -1113,10 +1113,14 @@ def description(self):
class QwenImageAutoBlocks(SequentialPipelineBlocks):
"""
Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage.
- - for image-to-image generation, you need to provide `image`
- - for inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop`.
- - to run the controlnet workflow, you need to provide `control_image`
- - for text-to-image generation, all you need to provide is `prompt`
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `prompt`, `image`
+ - `inpainting`: requires `prompt`, `mask_image`, `image`
+ - `controlnet_text2image`: requires `prompt`, `control_image`
+ - `controlnet_image2image`: requires `prompt`, `image`, `control_image`
+ - `controlnet_inpainting`: requires `prompt`, `mask_image`, `image`, `control_image`
Components:
text_encoder (`Qwen2_5_VLForConditionalGeneration`): The text encoder to use tokenizer (`Qwen2Tokenizer`):
@@ -1134,7 +1138,7 @@ class QwenImageAutoBlocks(SequentialPipelineBlocks):
Maximum sequence length for prompt encoding.
mask_image (`Image`, *optional*):
Mask image for inpainting.
- image (`Union[Image, List]`, *optional*):
+ image (`Image | list`, *optional*):
Reference image(s) for denoising. Can be a single image or list of images.
height (`int`, *optional*):
The height in pixels of the generated image.
@@ -1160,9 +1164,9 @@ class QwenImageAutoBlocks(SequentialPipelineBlocks):
Pre-generated noisy latents for image generation.
num_inference_steps (`int`):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -1183,12 +1187,12 @@ class QwenImageAutoBlocks(SequentialPipelineBlocks):
Scale for ControlNet conditioning.
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
- mask_overlay_kwargs (`Dict`, *optional*):
+ mask_overlay_kwargs (`dict`, *optional*):
The kwargs for the postprocess step to apply the mask overlay. generated in
InpaintProcessImagesInputStep.
Outputs:
- images (`List`):
+ images (`list`):
Generated images.
"""
@@ -1197,18 +1201,24 @@ class QwenImageAutoBlocks(SequentialPipelineBlocks):
block_classes = AUTO_BLOCKS.values()
block_names = AUTO_BLOCKS.keys()
+ # Workflow map defines the trigger conditions for each workflow.
+ # How to define:
+ # - Only include required inputs and trigger inputs (inputs that determine which blocks run)
+ # - currently, only supports `True` means the workflow triggers when the input is not None
+
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"prompt": True, "image": True},
+ "inpainting": {"prompt": True, "mask_image": True, "image": True},
+ "controlnet_text2image": {"prompt": True, "control_image": True},
+ "controlnet_image2image": {"prompt": True, "image": True, "control_image": True},
+ "controlnet_inpainting": {"prompt": True, "mask_image": True, "image": True, "control_image": True},
+ }
+
@property
def description(self):
- return (
- "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage.\n"
- + "- for image-to-image generation, you need to provide `image`\n"
- + "- for inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop`.\n"
- + "- to run the controlnet workflow, you need to provide `control_image`\n"
- + "- for text-to-image generation, all you need to provide is `prompt`"
- )
+ return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage."
@property
def outputs(self):
- return [
- OutputParam.template("images"),
- ]
+ return [OutputParam.template("images")]
diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py
index e1e5c4335481..37b80b69ec7e 100644
--- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py
+++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional
import torch
@@ -68,7 +67,7 @@ class QwenImageEditVLEncoderStep(SequentialPipelineBlocks):
(`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
prompt (`str`):
The prompt or prompts to guide image generation.
@@ -76,7 +75,7 @@ class QwenImageEditVLEncoderStep(SequentialPipelineBlocks):
The prompt or prompts not to guide the image generation.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized images
prompt_embeds (`Tensor`):
The prompt embeddings.
@@ -116,13 +115,13 @@ class QwenImageEditVaeEncoderStep(SequentialPipelineBlocks):
(`AutoencoderKLQwenImage`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized images
processed_image (`Tensor`):
The processed image
@@ -157,7 +156,7 @@ class QwenImageEditInpaintVaeEncoderStep(SequentialPipelineBlocks):
(`AutoencoderKLQwenImage`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
mask_image (`Image`):
Mask image for inpainting.
@@ -167,13 +166,13 @@ class QwenImageEditInpaintVaeEncoderStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized images
processed_image (`Tensor`):
The processed image
processed_mask_image (`Tensor`):
The processed mask image
- mask_overlay_kwargs (`Dict`):
+ mask_overlay_kwargs (`dict`):
The kwargs for the postprocess step to apply the mask overlay
image_latents (`Tensor`):
The latent representation of the input image.
@@ -451,9 +450,9 @@ class QwenImageEditCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -527,11 +526,11 @@ class QwenImageEditInpaintCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
strength (`float`, *optional*, defaults to 0.9):
Strength for img2img/inpainting.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -583,7 +582,7 @@ class QwenImageEditAutoCoreDenoiseStep(ConditionalPipelineBlocks):
block_trigger_inputs = ["processed_mask_image", "image_latents"]
default_block_name = "edit"
- def select_block(self, processed_mask_image=None, image_latents=None) -> Optional[str]:
+ def select_block(self, processed_mask_image=None, image_latents=None) -> str | None:
if processed_mask_image is not None:
return "edit_inpaint"
elif image_latents is not None:
@@ -628,7 +627,7 @@ class QwenImageEditDecodeStep(SequentialPipelineBlocks):
Output format: 'pil', 'np', 'pt'.
Outputs:
- images (`List`):
+ images (`list`):
Generated images. (tensor output of the vae decoder.)
"""
@@ -657,12 +656,12 @@ class QwenImageEditInpaintDecodeStep(SequentialPipelineBlocks):
step.
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
- mask_overlay_kwargs (`Dict`, *optional*):
+ mask_overlay_kwargs (`dict`, *optional*):
The kwargs for the postprocess step to apply the mask overlay. generated in
InpaintProcessImagesInputStep.
Outputs:
- images (`List`):
+ images (`list`):
Generated images. (tensor output of the vae decoder.)
"""
@@ -719,6 +718,11 @@ class QwenImageEditAutoBlocks(SequentialPipelineBlocks):
- for edit inpainting, you need to provide `mask_image` and `image`, optionally you can provide
`padding_mask_crop`
+
+ Supported workflows:
+ - `image_conditioned`: requires `prompt`, `image`
+ - `image_conditioned_inpainting`: requires `prompt`, `mask_image`, `image`
+
Components:
image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor
(`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`) image_mask_processor (`InpaintProcessor`) vae
@@ -726,7 +730,7 @@ class QwenImageEditAutoBlocks(SequentialPipelineBlocks):
(`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImageTransformer2DModel`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
prompt (`str`):
The prompt or prompts to guide image generation.
@@ -752,28 +756,32 @@ class QwenImageEditAutoBlocks(SequentialPipelineBlocks):
Pre-generated noisy latents for image generation.
num_inference_steps (`int`):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
strength (`float`, *optional*, defaults to 0.9):
Strength for img2img/inpainting.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
- mask_overlay_kwargs (`Dict`, *optional*):
+ mask_overlay_kwargs (`dict`, *optional*):
The kwargs for the postprocess step to apply the mask overlay. generated in
InpaintProcessImagesInputStep.
Outputs:
- images (`List`):
+ images (`list`):
Generated images.
"""
model_name = "qwenimage-edit"
block_classes = EDIT_AUTO_BLOCKS.values()
block_names = EDIT_AUTO_BLOCKS.keys()
+ _workflow_map = {
+ "image_conditioned": {"prompt": True, "image": True},
+ "image_conditioned_inpainting": {"prompt": True, "mask_image": True, "image": True},
+ }
@property
def description(self):
@@ -785,6 +793,4 @@ def description(self):
@property
def outputs(self):
- return [
- OutputParam.template("images"),
- ]
+ return [OutputParam.template("images")]
diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py
index 37656cef5d76..4a1f418d7b45 100644
--- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py
+++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py
@@ -58,7 +58,7 @@ class QwenImageEditPlusVLEncoderStep(SequentialPipelineBlocks):
(`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
prompt (`str`):
The prompt or prompts to guide image generation.
@@ -66,9 +66,9 @@ class QwenImageEditPlusVLEncoderStep(SequentialPipelineBlocks):
The prompt or prompts not to guide the image generation.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
Images resized to 1024x1024 target area for VAE encoding
- resized_cond_image (`List`):
+ resized_cond_image (`list`):
Images resized to 384x384 target area for VL text encoding
prompt_embeds (`Tensor`):
The prompt embeddings.
@@ -108,15 +108,15 @@ class QwenImageEditPlusVaeEncoderStep(SequentialPipelineBlocks):
(`AutoencoderKLQwenImage`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
Images resized to 1024x1024 target area for VAE encoding
- resized_cond_image (`List`):
+ resized_cond_image (`list`):
Images resized to 384x384 target area for VL text encoding
processed_image (`Tensor`):
The processed image
@@ -189,9 +189,9 @@ class QwenImageEditPlusInputStep(SequentialPipelineBlocks):
The negative prompt embeddings. (batch-expanded)
negative_prompt_embeds_mask (`Tensor`):
The negative prompt embeddings mask. (batch-expanded)
- image_height (`List`):
+ image_height (`list`):
The image heights calculated from the image latents dimension
- image_width (`List`):
+ image_width (`list`):
The image widths calculated from the image latents dimension
height (`int`):
if not provided, updated to image height
@@ -253,9 +253,9 @@ class QwenImageEditPlusCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -315,7 +315,7 @@ class QwenImageEditPlusDecodeStep(SequentialPipelineBlocks):
Output format: 'pil', 'np', 'pt'.
Outputs:
- images (`List`):
+ images (`list`):
Generated images. (tensor output of the vae decoder.)
"""
@@ -357,7 +357,7 @@ class QwenImageEditPlusAutoBlocks(SequentialPipelineBlocks):
transformer (`QwenImageTransformer2DModel`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
prompt (`str`):
The prompt or prompts to guide image generation.
@@ -375,9 +375,9 @@ class QwenImageEditPlusAutoBlocks(SequentialPipelineBlocks):
Pre-generated noisy latents for image generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -385,7 +385,7 @@ class QwenImageEditPlusAutoBlocks(SequentialPipelineBlocks):
Output format: 'pil', 'np', 'pt'.
Outputs:
- images (`List`):
+ images (`list`):
Generated images.
"""
@@ -404,6 +404,4 @@ def description(self):
@property
def outputs(self):
- return [
- OutputParam.template("images"),
- ]
+ return [OutputParam.template("images")]
diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py
index fdfeab048835..a10454f1fb0c 100644
--- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py
+++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py
@@ -60,7 +60,7 @@ class QwenImageLayeredTextEncoderStep(SequentialPipelineBlocks):
(`Qwen2VLProcessor`) tokenizer (`Qwen2Tokenizer`): The tokenizer to use guider (`ClassifierFreeGuidance`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
resolution (`int`, *optional*, defaults to 640):
The target area to resize the image to, can be 1024 or 640
@@ -74,7 +74,7 @@ class QwenImageLayeredTextEncoderStep(SequentialPipelineBlocks):
Maximum sequence length for prompt encoding.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized images
prompt (`str`):
The prompt or prompts to guide image generation. If not provided, updated using image caption
@@ -117,7 +117,7 @@ class QwenImageLayeredVaeEncoderStep(SequentialPipelineBlocks):
(`AutoencoderKLQwenImage`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
resolution (`int`, *optional*, defaults to 640):
The target area to resize the image to, can be 1024 or 640
@@ -125,7 +125,7 @@ class QwenImageLayeredVaeEncoderStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
Outputs:
- resized_image (`List`):
+ resized_image (`list`):
The resized images
processed_image (`Tensor`):
The processed image
@@ -250,9 +250,9 @@ class QwenImageLayeredCoreDenoiseStep(SequentialPipelineBlocks):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -317,7 +317,7 @@ class QwenImageLayeredAutoBlocks(SequentialPipelineBlocks):
scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImageTransformer2DModel`)
Inputs:
- image (`Union[Image, List]`):
+ image (`Image | list`):
Reference image(s) for denoising. Can be a single image or list of images.
resolution (`int`, *optional*, defaults to 640):
The target area to resize the image to, can be 1024 or 640
@@ -339,9 +339,9 @@ class QwenImageLayeredAutoBlocks(SequentialPipelineBlocks):
Number of layers to extract from the image
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
- sigmas (`List`, *optional*):
+ sigmas (`list`, *optional*):
Custom sigmas for the denoising process.
- attention_kwargs (`Dict`, *optional*):
+ attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
@@ -349,7 +349,7 @@ class QwenImageLayeredAutoBlocks(SequentialPipelineBlocks):
Output format: 'pil', 'np', 'pt'.
Outputs:
- images (`List`):
+ images (`list`):
Generated images.
"""
@@ -363,6 +363,4 @@ def description(self):
@property
def outputs(self):
- return [
- OutputParam.template("images"),
- ]
+ return [OutputParam.template("images")]
diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py
index 59ec46dc6d36..44f1c555cef3 100644
--- a/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py
+++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py
@@ -21,21 +21,7 @@
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
- _import_structure["encoders"] = ["StableDiffusionXLTextEncoderStep"]
- _import_structure["modular_blocks"] = [
- "ALL_BLOCKS",
- "AUTO_BLOCKS",
- "CONTROLNET_BLOCKS",
- "IMAGE2IMAGE_BLOCKS",
- "INPAINT_BLOCKS",
- "IP_ADAPTER_BLOCKS",
- "TEXT2IMAGE_BLOCKS",
- "StableDiffusionXLAutoBlocks",
- "StableDiffusionXLAutoControlnetStep",
- "StableDiffusionXLAutoDecodeStep",
- "StableDiffusionXLAutoIPAdapterStep",
- "StableDiffusionXLAutoVaeEncoderStep",
- ]
+ _import_structure["modular_blocks_stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks"]
_import_structure["modular_pipeline"] = ["StableDiffusionXLModularPipeline"]
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
@@ -45,23 +31,7 @@
except OptionalDependencyNotAvailable:
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
else:
- from .encoders import (
- StableDiffusionXLTextEncoderStep,
- )
- from .modular_blocks import (
- ALL_BLOCKS,
- AUTO_BLOCKS,
- CONTROLNET_BLOCKS,
- IMAGE2IMAGE_BLOCKS,
- INPAINT_BLOCKS,
- IP_ADAPTER_BLOCKS,
- TEXT2IMAGE_BLOCKS,
- StableDiffusionXLAutoBlocks,
- StableDiffusionXLAutoControlnetStep,
- StableDiffusionXLAutoDecodeStep,
- StableDiffusionXLAutoIPAdapterStep,
- StableDiffusionXLAutoVaeEncoderStep,
- )
+ from .modular_blocks_stable_diffusion_xl import StableDiffusionXLAutoBlocks
from .modular_pipeline import StableDiffusionXLModularPipeline
else:
import sys
diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py
index 70cbf0c1c78d..4a393e7ce296 100644
--- a/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py
+++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any
import PIL
import torch
@@ -46,10 +46,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -64,15 +64,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -105,7 +105,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -210,7 +210,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_images_per_prompt", default=1),
InputParam(
@@ -236,18 +236,18 @@ def inputs(self) -> List[InputParam]:
),
InputParam(
"ip_adapter_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
description="Pre-generated image embeddings for IP-Adapter. Can be generated from ip_adapter step.",
),
InputParam(
"negative_ip_adapter_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
description="Pre-generated negative image embeddings for IP-Adapter. Can be generated from ip_adapter step.",
),
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"batch_size",
@@ -285,13 +285,13 @@ def intermediate_outputs(self) -> List[str]:
),
OutputParam(
"ip_adapter_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields
description="image embeddings for IP-Adapter",
),
OutputParam(
"negative_ip_adapter_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields
description="negative image embeddings for IP-Adapter",
),
@@ -393,7 +393,7 @@ class StableDiffusionXLImg2ImgSetTimestepsStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", EulerDiscreteScheduler),
]
@@ -406,7 +406,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_inference_steps", default=50),
InputParam("timesteps"),
@@ -425,7 +425,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
OutputParam(
@@ -537,7 +537,7 @@ class StableDiffusionXLSetTimestepsStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", EulerDiscreteScheduler),
]
@@ -547,7 +547,7 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for inference"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_inference_steps", default=50),
InputParam("timesteps"),
@@ -556,7 +556,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
OutputParam(
@@ -605,7 +605,7 @@ class StableDiffusionXLInpaintPrepareLatentsStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", EulerDiscreteScheduler),
]
@@ -615,7 +615,7 @@ def description(self) -> str:
return "Step that prepares the latents for the inpainting process"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("latents"),
InputParam("num_images_per_prompt", default=1),
@@ -664,7 +664,7 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
@@ -879,7 +879,7 @@ class StableDiffusionXLImg2ImgPrepareLatentsStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKL),
ComponentSpec("scheduler", EulerDiscreteScheduler),
@@ -890,7 +890,7 @@ def description(self) -> str:
return "Step that prepares the latents for the image-to-image generation process"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("latents"),
InputParam("num_images_per_prompt", default=1),
@@ -918,7 +918,7 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
@@ -955,7 +955,7 @@ class StableDiffusionXLPrepareLatentsStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", EulerDiscreteScheduler),
ComponentSpec("vae", AutoencoderKL),
@@ -966,7 +966,7 @@ def description(self) -> str:
return "Prepare latents step that prepares the latents for the text-to-image generation process"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("height"),
InputParam("width"),
@@ -983,7 +983,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
@@ -1061,13 +1061,13 @@ class StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep(ModularPipelineB
model_name = "stable-diffusion-xl"
@property
- def expected_configs(self) -> List[ConfigSpec]:
+ def expected_configs(self) -> list[ConfigSpec]:
return [
ConfigSpec("requires_aesthetics_score", False),
]
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("unet", UNet2DConditionModel),
ComponentSpec(
@@ -1083,7 +1083,7 @@ def description(self) -> str:
return "Step that prepares the additional conditioning for the image-to-image/inpainting generation process"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("original_size"),
InputParam("target_size"),
@@ -1115,7 +1115,7 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"add_time_ids",
@@ -1284,7 +1284,7 @@ def description(self) -> str:
return "Step that prepares the additional conditioning for the text-to-image generation process"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("unet", UNet2DConditionModel),
ComponentSpec(
@@ -1296,7 +1296,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("original_size"),
InputParam("target_size"),
@@ -1326,7 +1326,7 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"add_time_ids",
@@ -1458,7 +1458,7 @@ class StableDiffusionXLControlNetInputStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("controlnet", ControlNetModel),
ComponentSpec(
@@ -1474,7 +1474,7 @@ def description(self) -> str:
return "step that prepare inputs for controlnet"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("control_image", required=True),
InputParam("control_guidance_start", default=0.0),
@@ -1502,26 +1502,26 @@ def inputs(self) -> List[Tuple[str, Any]]:
),
InputParam(
"crops_coords",
- type_hint=Optional[Tuple[int]],
+ type_hint=tuple[int] | None,
description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.",
),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("controlnet_cond", type_hint=torch.Tensor, description="The processed control image"),
OutputParam(
- "control_guidance_start", type_hint=List[float], description="The controlnet guidance start values"
+ "control_guidance_start", type_hint=list[float], description="The controlnet guidance start values"
),
OutputParam(
- "control_guidance_end", type_hint=List[float], description="The controlnet guidance end values"
+ "control_guidance_end", type_hint=list[float], description="The controlnet guidance end values"
),
OutputParam(
- "conditioning_scale", type_hint=List[float], description="The controlnet conditioning scale values"
+ "conditioning_scale", type_hint=list[float], description="The controlnet conditioning scale values"
),
OutputParam("guess_mode", type_hint=bool, description="Whether guess mode is used"),
- OutputParam("controlnet_keep", type_hint=List[float], description="The controlnet keep values"),
+ OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"),
]
# Modified from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image
@@ -1672,7 +1672,7 @@ class StableDiffusionXLControlNetUnionInputStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("controlnet", ControlNetUnionModel),
ComponentSpec(
@@ -1688,7 +1688,7 @@ def description(self) -> str:
return "step that prepares inputs for the ControlNetUnion model"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("control_image", required=True),
InputParam("control_mode", required=True),
@@ -1723,18 +1723,18 @@ def inputs(self) -> List[Tuple[str, Any]]:
),
InputParam(
"crops_coords",
- type_hint=Optional[Tuple[int]],
+ type_hint=tuple[int] | None,
description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.",
),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
- OutputParam("controlnet_cond", type_hint=List[torch.Tensor], description="The processed control images"),
+ OutputParam("controlnet_cond", type_hint=list[torch.Tensor], description="The processed control images"),
OutputParam(
"control_type_idx",
- type_hint=List[int],
+ type_hint=list[int],
description="The control mode indices",
kwargs_type="controlnet_kwargs",
),
@@ -1747,10 +1747,10 @@ def intermediate_outputs(self) -> List[OutputParam]:
OutputParam("control_guidance_start", type_hint=float, description="The controlnet guidance start value"),
OutputParam("control_guidance_end", type_hint=float, description="The controlnet guidance end value"),
OutputParam(
- "conditioning_scale", type_hint=List[float], description="The controlnet conditioning scale values"
+ "conditioning_scale", type_hint=list[float], description="The controlnet conditioning scale values"
),
OutputParam("guess_mode", type_hint=bool, description="Whether guess mode is used"),
- OutputParam("controlnet_keep", type_hint=List[float], description="The controlnet keep values"),
+ OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"),
]
# Modified from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image
diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py
index 6e0307260d1d..7e505559f685 100644
--- a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py
+++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, List, Tuple, Union
+from typing import Any
import numpy as np
import PIL
@@ -36,7 +36,7 @@ class StableDiffusionXLDecodeStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKL),
ComponentSpec(
@@ -52,7 +52,7 @@ def description(self) -> str:
return "Step that decodes the denoised latents into images"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("output_type", default="pil"),
InputParam(
@@ -64,11 +64,11 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"images",
- type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]],
+ type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array],
description="The generated images, can be a PIL.Image.Image, torch.Tensor or a numpy array",
)
]
@@ -153,7 +153,7 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
@@ -164,19 +164,19 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("image"),
InputParam("mask_image"),
InputParam("padding_mask_crop"),
InputParam(
"images",
- type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]],
+ type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array],
description="The generated images from the decode step",
),
InputParam(
"crops_coords",
- type_hint=Tuple[int, int],
+ type_hint=tuple[int, int],
description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.",
),
]
diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py
index 862315e59169..0190bc3ea62f 100644
--- a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py
+++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, List, Optional, Tuple
+from typing import Any
import torch
@@ -41,7 +41,7 @@ class StableDiffusionXLLoopBeforeDenoiser(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", EulerDiscreteScheduler),
]
@@ -55,7 +55,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[str]:
+ def inputs(self) -> list[str]:
return [
InputParam(
"latents",
@@ -77,7 +77,7 @@ class StableDiffusionXLInpaintLoopBeforeDenoiser(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", EulerDiscreteScheduler),
ComponentSpec("unet", UNet2DConditionModel),
@@ -91,7 +91,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[str]:
+ def inputs(self) -> list[str]:
return [
InputParam(
"latents",
@@ -101,12 +101,12 @@ def inputs(self) -> List[str]:
),
InputParam(
"mask",
- type_hint=Optional[torch.Tensor],
+ type_hint=torch.Tensor | None,
description="The mask to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.",
),
InputParam(
"masked_image_latents",
- type_hint=Optional[torch.Tensor],
+ type_hint=torch.Tensor | None,
description="The masked image latents to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.",
),
]
@@ -148,7 +148,7 @@ class StableDiffusionXLLoopDenoiser(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -168,7 +168,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("cross_attention_kwargs"),
InputParam(
@@ -179,7 +179,7 @@ def inputs(self) -> List[Tuple[str, Any]]:
),
InputParam(
"timestep_cond",
- type_hint=Optional[torch.Tensor],
+ type_hint=torch.Tensor | None,
description="The guidance scale embedding to use for Latent Consistency Models(LCMs). Can be generated in prepare_additional_conditioning step.",
),
InputParam(
@@ -262,7 +262,7 @@ class StableDiffusionXLControlNetLoopDenoiser(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -283,7 +283,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("cross_attention_kwargs"),
InputParam(
@@ -306,12 +306,12 @@ def inputs(self) -> List[Tuple[str, Any]]:
InputParam(
"controlnet_keep",
required=True,
- type_hint=List[float],
+ type_hint=list[float],
description="The controlnet keep values to use for the denoising process. Can be generated in prepare_controlnet_inputs step.",
),
InputParam(
"timestep_cond",
- type_hint=Optional[torch.Tensor],
+ type_hint=torch.Tensor | None,
description="The guidance scale embedding to use for Latent Consistency Models(LCMs), can be generated by prepare_additional_conditioning step",
),
InputParam(
@@ -472,7 +472,7 @@ class StableDiffusionXLLoopAfterDenoiser(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", EulerDiscreteScheduler),
]
@@ -486,14 +486,14 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("eta", default=0.0),
InputParam("generator"),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
# YiYi TODO: move this out of here
@@ -537,7 +537,7 @@ class StableDiffusionXLInpaintLoopAfterDenoiser(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", EulerDiscreteScheduler),
ComponentSpec("unet", UNet2DConditionModel),
@@ -552,7 +552,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam("eta", default=0.0),
InputParam("generator"),
@@ -564,23 +564,23 @@ def inputs(self) -> List[Tuple[str, Any]]:
),
InputParam(
"mask",
- type_hint=Optional[torch.Tensor],
+ type_hint=torch.Tensor | None,
description="The mask to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.",
),
InputParam(
"noise",
- type_hint=Optional[torch.Tensor],
+ type_hint=torch.Tensor | None,
description="The noise added to the image latents, for inpainting task only. Can be generated in prepare_latent step.",
),
InputParam(
"image_latents",
- type_hint=Optional[torch.Tensor],
+ type_hint=torch.Tensor | None,
description="The image latents to use for the denoising process, for inpainting/image-to-image task only. Can be generated in vae_encode or prepare_latent step.",
),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
@staticmethod
@@ -654,7 +654,7 @@ def description(self) -> str:
)
@property
- def loop_expected_components(self) -> List[ComponentSpec]:
+ def loop_expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -667,7 +667,7 @@ def loop_expected_components(self) -> List[ComponentSpec]:
]
@property
- def loop_inputs(self) -> List[InputParam]:
+ def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
"timesteps",
diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py
index 90b254b6f5d4..8387ae7bd6b6 100644
--- a/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py
+++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple
-
import torch
from transformers import (
CLIPImageProcessor,
@@ -45,7 +43,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -71,7 +69,7 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("image_encoder", CLIPVisionModelWithProjection),
ComponentSpec(
@@ -90,7 +88,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
"ip_adapter_image",
@@ -101,7 +99,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("ip_adapter_embeds", type_hint=torch.Tensor, description="IP adapter image embeddings"),
OutputParam(
@@ -223,7 +221,7 @@ def description(self) -> str:
return "Text Encoder step that generate text_embeddings to guide the image generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", CLIPTextModel),
ComponentSpec("text_encoder_2", CLIPTextModelWithProjection),
@@ -238,11 +236,11 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def expected_configs(self) -> List[ConfigSpec]:
+ def expected_configs(self) -> list[ConfigSpec]:
return [ConfigSpec("force_zeros_for_empty_prompt", True)]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("prompt"),
InputParam("prompt_2"),
@@ -253,7 +251,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
@@ -296,26 +294,26 @@ def check_inputs(block_state):
def encode_prompt(
components,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
prepare_unconditional_embeds: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -324,11 +322,11 @@ def encode_prompt(
number of images that should be generated per prompt
prepare_unconditional_embeds (`bool`):
whether to use prepare unconditional embeddings or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -450,7 +448,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -584,7 +582,7 @@ def description(self) -> str:
return "Vae Encoder step that encode the input image into a latent representation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKL),
ComponentSpec(
@@ -596,7 +594,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("image", required=True),
InputParam("height"),
@@ -605,13 +603,13 @@ def inputs(self) -> List[InputParam]:
InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"),
InputParam(
"preprocess_kwargs",
- type_hint=Optional[dict],
+ type_hint=dict | None,
description="A kwargs dictionary that if specified is passed along to the `ImageProcessor` as defined under `self.image_processor` in [diffusers.image_processor.VaeImageProcessor]",
),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"image_latents",
@@ -687,7 +685,7 @@ class StableDiffusionXLInpaintVaeEncoderStep(ModularPipelineBlocks):
model_name = "stable-diffusion-xl"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKL),
ComponentSpec(
@@ -711,7 +709,7 @@ def description(self) -> str:
return "Vae encoder step that prepares the image and mask for the inpainting process"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("height"),
InputParam("width"),
@@ -723,7 +721,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"image_latents", type_hint=torch.Tensor, description="The latents representation of the input image"
@@ -736,7 +734,7 @@ def intermediate_outputs(self) -> List[OutputParam]:
),
OutputParam(
"crops_coords",
- type_hint=Optional[Tuple[int, int]],
+ type_hint=tuple[int, int] | None,
description="The crop coordinates to use for the preprocess/postprocess of the image and mask",
),
]
diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py
similarity index 55%
rename from src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py
rename to src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py
index 68b5e33755b5..a7a18e514777 100644
--- a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py
+++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py
@@ -14,7 +14,7 @@
from ...utils import logging
from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
-from ..modular_pipeline_utils import InsertableDict
+from ..modular_pipeline_utils import OutputParam
from .before_denoise import (
StableDiffusionXLControlNetInputStep,
StableDiffusionXLControlNetUnionInputStep,
@@ -277,7 +277,161 @@ def description(self):
# ip-adapter, controlnet, text2img, img2img, inpainting
+# auto_docstring
class StableDiffusionXLAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion
+ XL.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `image`, `prompt`
+ - `inpainting`: requires `mask_image`, `image`, `prompt`
+ - `controlnet_text2image`: requires `control_image`, `prompt`
+ - `controlnet_image2image`: requires `control_image`, `image`, `prompt`
+ - `controlnet_inpainting`: requires `control_image`, `mask_image`, `image`, `prompt`
+ - `controlnet_union_text2image`: requires `control_image`, `control_mode`, `prompt`
+ - `controlnet_union_image2image`: requires `control_image`, `control_mode`, `image`, `prompt`
+ - `controlnet_union_inpainting`: requires `control_image`, `control_mode`, `mask_image`, `image`, `prompt`
+ - `ip_adapter_text2image`: requires `ip_adapter_image`, `prompt`
+ - `ip_adapter_image2image`: requires `ip_adapter_image`, `image`, `prompt`
+ - `ip_adapter_inpainting`: requires `ip_adapter_image`, `mask_image`, `image`, `prompt`
+ - `ip_adapter_controlnet_text2image`: requires `ip_adapter_image`, `control_image`, `prompt`
+ - `ip_adapter_controlnet_image2image`: requires `ip_adapter_image`, `control_image`, `image`, `prompt`
+ - `ip_adapter_controlnet_inpainting`: requires `ip_adapter_image`, `control_image`, `mask_image`, `image`,
+ `prompt`
+ - `ip_adapter_controlnet_union_text2image`: requires `ip_adapter_image`, `control_image`, `control_mode`,
+ `prompt`
+ - `ip_adapter_controlnet_union_image2image`: requires `ip_adapter_image`, `control_image`, `control_mode`,
+ `image`, `prompt`
+ - `ip_adapter_controlnet_union_inpainting`: requires `ip_adapter_image`, `control_image`, `control_mode`,
+ `mask_image`, `image`, `prompt`
+
+ Components:
+ text_encoder (`CLIPTextModel`) text_encoder_2 (`CLIPTextModelWithProjection`) tokenizer (`CLIPTokenizer`)
+ tokenizer_2 (`CLIPTokenizer`) guider (`ClassifierFreeGuidance`) image_encoder
+ (`CLIPVisionModelWithProjection`) feature_extractor (`CLIPImageProcessor`) unet (`UNet2DConditionModel`) vae
+ (`AutoencoderKL`) image_processor (`VaeImageProcessor`) mask_processor (`VaeImageProcessor`) scheduler
+ (`EulerDiscreteScheduler`) controlnet (`ControlNetUnionModel`) control_image_processor (`VaeImageProcessor`)
+
+ Configs:
+ force_zeros_for_empty_prompt (default: True) requires_aesthetics_score (default: False)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ prompt_2 (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt_2 (`None`, *optional*):
+ TODO: Add description.
+ cross_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ clip_skip (`None`, *optional*):
+ TODO: Add description.
+ ip_adapter_image (`Image | ndarray | Tensor | list | list | list`, *optional*):
+ The image(s) to be used as ip adapter
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ mask_image (`None`, *optional*):
+ TODO: Add description.
+ padding_mask_crop (`None`, *optional*):
+ TODO: Add description.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ generator (`None`, *optional*):
+ TODO: Add description.
+ preprocess_kwargs (`dict | NoneType`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `ImageProcessor` as defined under
+ `self.image_processor` in [diffusers.image_processor.VaeImageProcessor]
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ ip_adapter_embeds (`list`, *optional*):
+ Pre-generated image embeddings for IP-Adapter. Can be generated from ip_adapter step.
+ negative_ip_adapter_embeds (`list`, *optional*):
+ Pre-generated negative image embeddings for IP-Adapter. Can be generated from ip_adapter step.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ denoising_end (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.3):
+ TODO: Add description.
+ denoising_start (`None`, *optional*):
+ TODO: Add description.
+ latents (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ The latents representing the reference image for image-to-image/inpainting generation. Can be generated
+ in vae_encode step.
+ mask (`Tensor`, *optional*):
+ The mask for the inpainting generation. Can be generated in vae_encode step.
+ masked_image_latents (`Tensor`, *optional*):
+ The masked image latents for the inpainting generation (only for inpainting-specific unet). Can be
+ generated in vae_encode step.
+ original_size (`None`, *optional*):
+ TODO: Add description.
+ target_size (`None`, *optional*):
+ TODO: Add description.
+ negative_original_size (`None`, *optional*):
+ TODO: Add description.
+ negative_target_size (`None`, *optional*):
+ TODO: Add description.
+ crops_coords_top_left (`None`, *optional*, defaults to (0, 0)):
+ TODO: Add description.
+ negative_crops_coords_top_left (`None`, *optional*, defaults to (0, 0)):
+ TODO: Add description.
+ aesthetic_score (`None`, *optional*, defaults to 6.0):
+ TODO: Add description.
+ negative_aesthetic_score (`None`, *optional*, defaults to 2.0):
+ TODO: Add description.
+ control_image (`None`, *optional*):
+ TODO: Add description.
+ control_mode (`None`, *optional*):
+ TODO: Add description.
+ control_guidance_start (`None`, *optional*, defaults to 0.0):
+ TODO: Add description.
+ control_guidance_end (`None`, *optional*, defaults to 1.0):
+ TODO: Add description.
+ controlnet_conditioning_scale (`None`, *optional*, defaults to 1.0):
+ TODO: Add description.
+ guess_mode (`None`, *optional*, defaults to False):
+ TODO: Add description.
+ crops_coords (`tuple | NoneType`, *optional*):
+ The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can
+ be generated in vae_encode step.
+ controlnet_cond (`Tensor`, *optional*):
+ The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.
+ conditioning_scale (`float`, *optional*):
+ The controlnet conditioning scale value to use for the denoising process. Can be generated in
+ prepare_controlnet_inputs step.
+ controlnet_keep (`list`, *optional*):
+ The controlnet keep values to use for the denoising process. Can be generated in
+ prepare_controlnet_inputs step.
+ **denoiser_input_fields (`None`, *optional*):
+ All conditional model inputs that need to be prepared with guider. It should contain
+ prompt_embeds/negative_prompt_embeds, add_time_ids/negative_add_time_ids,
+ pooled_prompt_embeds/negative_pooled_prompt_embeds, and ip_adapter_embeds/negative_ip_adapter_embeds
+ (optional).please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when
+ they are created and added to the pipeline state
+ eta (`None`, *optional*, defaults to 0.0):
+ TODO: Add description.
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
block_classes = [
StableDiffusionXLTextEncoderStep,
StableDiffusionXLAutoIPAdapterStep,
@@ -293,103 +447,66 @@ class StableDiffusionXLAutoBlocks(SequentialPipelineBlocks):
"decode",
]
- @property
- def description(self):
- return (
- "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion XL.\n"
- + "- for image-to-image generation, you need to provide either `image` or `image_latents`\n"
- + "- for inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop` \n"
- + "- to run the controlnet workflow, you need to provide `control_image`\n"
- + "- to run the controlnet_union workflow, you need to provide `control_image` and `control_mode`\n"
- + "- to run the ip_adapter workflow, you need to provide `ip_adapter_image`\n"
- + "- for text-to-image generation, all you need to provide is `prompt`"
- )
-
-
-# controlnet (input + denoise step)
-class StableDiffusionXLAutoControlnetStep(SequentialPipelineBlocks):
- block_classes = [
- StableDiffusionXLAutoControlNetInputStep,
- StableDiffusionXLAutoControlNetDenoiseStep,
- ]
- block_names = ["controlnet_input", "controlnet_denoise"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"image": True, "prompt": True},
+ "inpainting": {"mask_image": True, "image": True, "prompt": True},
+ "controlnet_text2image": {"control_image": True, "prompt": True},
+ "controlnet_image2image": {"control_image": True, "image": True, "prompt": True},
+ "controlnet_inpainting": {"control_image": True, "mask_image": True, "image": True, "prompt": True},
+ "controlnet_union_text2image": {"control_image": True, "control_mode": True, "prompt": True},
+ "controlnet_union_image2image": {"control_image": True, "control_mode": True, "image": True, "prompt": True},
+ "controlnet_union_inpainting": {
+ "control_image": True,
+ "control_mode": True,
+ "mask_image": True,
+ "image": True,
+ "prompt": True,
+ },
+ "ip_adapter_text2image": {"ip_adapter_image": True, "prompt": True},
+ "ip_adapter_image2image": {"ip_adapter_image": True, "image": True, "prompt": True},
+ "ip_adapter_inpainting": {"ip_adapter_image": True, "mask_image": True, "image": True, "prompt": True},
+ "ip_adapter_controlnet_text2image": {"ip_adapter_image": True, "control_image": True, "prompt": True},
+ "ip_adapter_controlnet_image2image": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "image": True,
+ "prompt": True,
+ },
+ "ip_adapter_controlnet_inpainting": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "mask_image": True,
+ "image": True,
+ "prompt": True,
+ },
+ "ip_adapter_controlnet_union_text2image": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "control_mode": True,
+ "prompt": True,
+ },
+ "ip_adapter_controlnet_union_image2image": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "control_mode": True,
+ "image": True,
+ "prompt": True,
+ },
+ "ip_adapter_controlnet_union_inpainting": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "control_mode": True,
+ "mask_image": True,
+ "image": True,
+ "prompt": True,
+ },
+ }
@property
def description(self):
- return (
- "Controlnet auto step that prepare the controlnet input and denoise the latents. "
- + "It works for both controlnet and controlnet_union and supports text2img, img2img and inpainting tasks."
- + " (it should be replace at 'denoise' step)"
- )
-
-
-TEXT2IMAGE_BLOCKS = InsertableDict(
- [
- ("text_encoder", StableDiffusionXLTextEncoderStep),
- ("input", StableDiffusionXLInputStep),
- ("set_timesteps", StableDiffusionXLSetTimestepsStep),
- ("prepare_latents", StableDiffusionXLPrepareLatentsStep),
- ("prepare_add_cond", StableDiffusionXLPrepareAdditionalConditioningStep),
- ("denoise", StableDiffusionXLDenoiseStep),
- ("decode", StableDiffusionXLDecodeStep),
- ]
-)
+ return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion XL."
-IMAGE2IMAGE_BLOCKS = InsertableDict(
- [
- ("text_encoder", StableDiffusionXLTextEncoderStep),
- ("vae_encoder", StableDiffusionXLVaeEncoderStep),
- ("input", StableDiffusionXLInputStep),
- ("set_timesteps", StableDiffusionXLImg2ImgSetTimestepsStep),
- ("prepare_latents", StableDiffusionXLImg2ImgPrepareLatentsStep),
- ("prepare_add_cond", StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep),
- ("denoise", StableDiffusionXLDenoiseStep),
- ("decode", StableDiffusionXLDecodeStep),
- ]
-)
-
-INPAINT_BLOCKS = InsertableDict(
- [
- ("text_encoder", StableDiffusionXLTextEncoderStep),
- ("vae_encoder", StableDiffusionXLInpaintVaeEncoderStep),
- ("input", StableDiffusionXLInputStep),
- ("set_timesteps", StableDiffusionXLImg2ImgSetTimestepsStep),
- ("prepare_latents", StableDiffusionXLInpaintPrepareLatentsStep),
- ("prepare_add_cond", StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep),
- ("denoise", StableDiffusionXLInpaintDenoiseStep),
- ("decode", StableDiffusionXLInpaintDecodeStep),
- ]
-)
-
-CONTROLNET_BLOCKS = InsertableDict(
- [
- ("denoise", StableDiffusionXLAutoControlnetStep),
- ]
-)
-
-
-IP_ADAPTER_BLOCKS = InsertableDict(
- [
- ("ip_adapter", StableDiffusionXLAutoIPAdapterStep),
- ]
-)
-
-AUTO_BLOCKS = InsertableDict(
- [
- ("text_encoder", StableDiffusionXLTextEncoderStep),
- ("ip_adapter", StableDiffusionXLAutoIPAdapterStep),
- ("vae_encoder", StableDiffusionXLAutoVaeEncoderStep),
- ("denoise", StableDiffusionXLCoreDenoiseStep),
- ("decode", StableDiffusionXLAutoDecodeStep),
- ]
-)
-
-
-ALL_BLOCKS = {
- "text2img": TEXT2IMAGE_BLOCKS,
- "img2img": IMAGE2IMAGE_BLOCKS,
- "inpaint": INPAINT_BLOCKS,
- "controlnet": CONTROLNET_BLOCKS,
- "ip_adapter": IP_ADAPTER_BLOCKS,
- "auto": AUTO_BLOCKS,
-}
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py
index f2a4c96073ea..209e2b11814f 100644
--- a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py
+++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import numpy as np
import PIL
import torch
@@ -94,30 +92,30 @@ def num_channels_latents(self):
# auto_docstring
SDXL_INPUTS_SCHEMA = {
"prompt": InputParam(
- "prompt", type_hint=Union[str, List[str]], description="The prompt or prompts to guide the image generation"
+ "prompt", type_hint=str | list[str], description="The prompt or prompts to guide the image generation"
),
"prompt_2": InputParam(
"prompt_2",
- type_hint=Union[str, List[str]],
+ type_hint=str | list[str],
description="The prompt or prompts to be sent to the tokenizer_2 and text_encoder_2",
),
"negative_prompt": InputParam(
"negative_prompt",
- type_hint=Union[str, List[str]],
+ type_hint=str | list[str],
description="The prompt or prompts not to guide the image generation",
),
"negative_prompt_2": InputParam(
"negative_prompt_2",
- type_hint=Union[str, List[str]],
+ type_hint=str | list[str],
description="The negative prompt or prompts for text_encoder_2",
),
"cross_attention_kwargs": InputParam(
"cross_attention_kwargs",
- type_hint=Optional[dict],
+ type_hint=dict | None,
description="Kwargs dictionary passed to the AttentionProcessor",
),
"clip_skip": InputParam(
- "clip_skip", type_hint=Optional[int], description="Number of layers to skip in CLIP text encoder"
+ "clip_skip", type_hint=int | None, description="Number of layers to skip in CLIP text encoder"
),
"image": InputParam(
"image",
@@ -133,11 +131,11 @@ def num_channels_latents(self):
),
"generator": InputParam(
"generator",
- type_hint=Optional[Union[torch.Generator, List[torch.Generator]]],
+ type_hint=torch.Generator | list[torch.Generator] | None,
description="Generator(s) for deterministic generation",
),
- "height": InputParam("height", type_hint=Optional[int], description="Height in pixels of the generated image"),
- "width": InputParam("width", type_hint=Optional[int], description="Width in pixels of the generated image"),
+ "height": InputParam("height", type_hint=int | None, description="Height in pixels of the generated image"),
+ "width": InputParam("width", type_hint=int | None, description="Width in pixels of the generated image"),
"num_images_per_prompt": InputParam(
"num_images_per_prompt", type_hint=int, default=1, description="Number of images to generate per prompt"
),
@@ -145,14 +143,14 @@ def num_channels_latents(self):
"num_inference_steps", type_hint=int, default=50, description="Number of denoising steps"
),
"timesteps": InputParam(
- "timesteps", type_hint=Optional[torch.Tensor], description="Custom timesteps for the denoising process"
+ "timesteps", type_hint=torch.Tensor | None, description="Custom timesteps for the denoising process"
),
"sigmas": InputParam(
- "sigmas", type_hint=Optional[torch.Tensor], description="Custom sigmas for the denoising process"
+ "sigmas", type_hint=torch.Tensor | None, description="Custom sigmas for the denoising process"
),
"denoising_end": InputParam(
"denoising_end",
- type_hint=Optional[float],
+ type_hint=float | None,
description="Fraction of denoising process to complete before termination",
),
# YiYi Notes: img2img defaults to 0.3, inpainting defaults to 0.9999
@@ -160,43 +158,43 @@ def num_channels_latents(self):
"strength", type_hint=float, default=0.3, description="How much to transform the reference image"
),
"denoising_start": InputParam(
- "denoising_start", type_hint=Optional[float], description="Starting point of the denoising process"
+ "denoising_start", type_hint=float | None, description="Starting point of the denoising process"
),
"latents": InputParam(
- "latents", type_hint=Optional[torch.Tensor], description="Pre-generated noisy latents for image generation"
+ "latents", type_hint=torch.Tensor | None, description="Pre-generated noisy latents for image generation"
),
"padding_mask_crop": InputParam(
"padding_mask_crop",
- type_hint=Optional[Tuple[int, int]],
+ type_hint=tuple[int, int] | None,
description="Size of margin in crop for image and mask",
),
"original_size": InputParam(
"original_size",
- type_hint=Optional[Tuple[int, int]],
+ type_hint=tuple[int, int] | None,
description="Original size of the image for SDXL's micro-conditioning",
),
"target_size": InputParam(
- "target_size", type_hint=Optional[Tuple[int, int]], description="Target size for SDXL's micro-conditioning"
+ "target_size", type_hint=tuple[int, int] | None, description="Target size for SDXL's micro-conditioning"
),
"negative_original_size": InputParam(
"negative_original_size",
- type_hint=Optional[Tuple[int, int]],
+ type_hint=tuple[int, int] | None,
description="Negative conditioning based on image resolution",
),
"negative_target_size": InputParam(
"negative_target_size",
- type_hint=Optional[Tuple[int, int]],
+ type_hint=tuple[int, int] | None,
description="Negative conditioning based on target resolution",
),
"crops_coords_top_left": InputParam(
"crops_coords_top_left",
- type_hint=Tuple[int, int],
+ type_hint=tuple[int, int],
default=(0, 0),
description="Top-left coordinates for SDXL's micro-conditioning",
),
"negative_crops_coords_top_left": InputParam(
"negative_crops_coords_top_left",
- type_hint=Tuple[int, int],
+ type_hint=tuple[int, int],
default=(0, 0),
description="Negative conditioning crop coordinates",
),
@@ -221,19 +219,19 @@ def num_channels_latents(self):
),
"control_guidance_start": InputParam(
"control_guidance_start",
- type_hint=Union[float, List[float]],
+ type_hint=float | list[float],
default=0.0,
description="When ControlNet starts applying",
),
"control_guidance_end": InputParam(
"control_guidance_end",
- type_hint=Union[float, List[float]],
+ type_hint=float | list[float],
default=1.0,
description="When ControlNet stops applying",
),
"controlnet_conditioning_scale": InputParam(
"controlnet_conditioning_scale",
- type_hint=Union[float, List[float]],
+ type_hint=float | list[float],
default=1.0,
description="Scale factor for ControlNet outputs",
),
@@ -244,7 +242,7 @@ def num_channels_latents(self):
description="Enables ControlNet encoder to recognize input without prompts",
),
"control_mode": InputParam(
- "control_mode", type_hint=List[int], required=True, description="Control mode for union controlnet"
+ "control_mode", type_hint=list[int], required=True, description="Control mode for union controlnet"
),
"prompt_embeds": InputParam(
"prompt_embeds",
@@ -264,7 +262,7 @@ def num_channels_latents(self):
"batch_size": InputParam("batch_size", type_hint=int, required=True, description="Number of prompts"),
"dtype": InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"),
"preprocess_kwargs": InputParam(
- "preprocess_kwargs", type_hint=Optional[dict], description="Kwargs for ImageProcessor"
+ "preprocess_kwargs", type_hint=dict | None, description="Kwargs for ImageProcessor"
),
"latent_timestep": InputParam(
"latent_timestep", type_hint=torch.Tensor, required=True, description="Initial noise level timestep"
@@ -284,18 +282,18 @@ def num_channels_latents(self):
),
"timestep_cond": InputParam("timestep_cond", type_hint=torch.Tensor, description="Timestep conditioning for LCM"),
"noise": InputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"),
- "crops_coords": InputParam("crops_coords", type_hint=Optional[Tuple[int]], description="Crop coordinates"),
+ "crops_coords": InputParam("crops_coords", type_hint=tuple[int] | None, description="Crop coordinates"),
"ip_adapter_embeds": InputParam(
- "ip_adapter_embeds", type_hint=List[torch.Tensor], description="Image embeddings for IP-Adapter"
+ "ip_adapter_embeds", type_hint=list[torch.Tensor], description="Image embeddings for IP-Adapter"
),
"negative_ip_adapter_embeds": InputParam(
"negative_ip_adapter_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
description="Negative image embeddings for IP-Adapter",
),
"images": InputParam(
"images",
- type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]],
+ type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array],
required=True,
description="Generated images",
),
@@ -324,7 +322,7 @@ def num_channels_latents(self):
"masked_image_latents": OutputParam(
"masked_image_latents", type_hint=torch.Tensor, description="Masked image latents for inpainting"
),
- "crops_coords": OutputParam("crops_coords", type_hint=Optional[Tuple[int]], description="Crop coordinates"),
+ "crops_coords": OutputParam("crops_coords", type_hint=tuple[int] | None, description="Crop coordinates"),
"timesteps": OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps for inference"),
"num_inference_steps": OutputParam("num_inference_steps", type_hint=int, description="Number of denoising steps"),
"latent_timestep": OutputParam(
@@ -338,16 +336,16 @@ def num_channels_latents(self):
"latents": OutputParam("latents", type_hint=torch.Tensor, description="Denoised latents"),
"noise": OutputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"),
"ip_adapter_embeds": OutputParam(
- "ip_adapter_embeds", type_hint=List[torch.Tensor], description="Image embeddings for IP-Adapter"
+ "ip_adapter_embeds", type_hint=list[torch.Tensor], description="Image embeddings for IP-Adapter"
),
"negative_ip_adapter_embeds": OutputParam(
"negative_ip_adapter_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
description="Negative image embeddings for IP-Adapter",
),
"images": OutputParam(
"images",
- type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]],
+ type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array],
description="Generated images",
),
}
@@ -356,9 +354,7 @@ def num_channels_latents(self):
SDXL_OUTPUTS_SCHEMA = {
"images": OutputParam(
"images",
- type_hint=Union[
- Tuple[Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]]], StableDiffusionXLPipelineOutput
- ],
+ type_hint=tuple[list[PIL.Image.Image] | list[torch.Tensor] | list[np.array]] | StableDiffusionXLPipelineOutput,
description="The final generated images",
)
}
diff --git a/src/diffusers/modular_pipelines/wan/before_denoise.py b/src/diffusers/modular_pipelines/wan/before_denoise.py
index 719ba4c21148..398b9665522c 100644
--- a/src/diffusers/modular_pipelines/wan/before_denoise.py
+++ b/src/diffusers/modular_pipelines/wan/before_denoise.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Tuple, Union
import torch
@@ -93,7 +92,7 @@ def repeat_tensor_to_batch_size(
def calculate_dimension_from_latents(
latents: torch.Tensor, vae_scale_factor_temporal: int, vae_scale_factor_spatial: int
-) -> Tuple[int, int]:
+) -> tuple[int, int]:
"""Calculate image dimensions from latent tensor dimensions.
This function converts latent temporal and spatial dimensions to image temporal and spatial dimensions by
@@ -108,7 +107,7 @@ def calculate_dimension_from_latents(
Typically 8 for most VAEs (image is 8x larger than latents in each dimension)
Returns:
- Tuple[int, int]: The calculated image dimensions as (height, width)
+ tuple[int, int]: The calculated image dimensions as (height, width)
Raises:
ValueError: If latents tensor doesn't have 4 or 5 dimensions
@@ -129,10 +128,10 @@ def calculate_dimension_from_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -147,15 +146,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -201,13 +200,13 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("transformer", WanTransformer3DModel),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_videos_per_prompt", default=1),
InputParam(
@@ -224,7 +223,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"batch_size",
@@ -280,8 +279,8 @@ class WanAdditionalInputsStep(ModularPipelineBlocks):
def __init__(
self,
- image_latent_inputs: List[str] = ["image_condition_latents"],
- additional_batch_inputs: List[str] = [],
+ image_latent_inputs: list[str] = ["image_condition_latents"],
+ additional_batch_inputs: list[str] = [],
):
"""Initialize a configurable step that standardizes the inputs for the denoising step. It:\n"
@@ -292,7 +291,7 @@ def __init__(
This is a dynamic block that allows you to configure which inputs to process.
Args:
- image_latent_inputs (List[str], optional): Names of image latent tensors to process.
+ image_latent_inputs (list[str], optional): Names of image latent tensors to process.
In additional to adjust batch size of these inputs, they will be used to determine height/width. Can be
a single string or list of strings. Defaults to ["image_condition_latents"].
additional_batch_inputs (List[str], optional):
@@ -339,7 +338,7 @@ def description(self) -> str:
return summary_section + inputs_info + placement_section
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [
InputParam(name="num_videos_per_prompt", default=1),
InputParam(name="batch_size", required=True),
@@ -409,7 +408,7 @@ class WanSetTimestepsStep(ModularPipelineBlocks):
model_name = "wan"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", UniPCMultistepScheduler),
]
@@ -419,7 +418,7 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for inference"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_inference_steps", default=50),
InputParam("timesteps"),
@@ -451,12 +450,12 @@ def description(self) -> str:
return "Prepare latents step that prepares the latents for the text-to-video generation process"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("height", type_hint=int),
InputParam("width", type_hint=int),
InputParam("num_frames", type_hint=int),
- InputParam("latents", type_hint=Optional[torch.Tensor]),
+ InputParam("latents", type_hint=torch.Tensor | None),
InputParam("num_videos_per_prompt", type_hint=int, default=1),
InputParam("generator"),
InputParam(
@@ -469,7 +468,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
@@ -500,10 +499,10 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
diff --git a/src/diffusers/modular_pipelines/wan/decoders.py b/src/diffusers/modular_pipelines/wan/decoders.py
index c26a8b11ba5c..9d1a4cf4f348 100644
--- a/src/diffusers/modular_pipelines/wan/decoders.py
+++ b/src/diffusers/modular_pipelines/wan/decoders.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, List, Tuple, Union
+from typing import Any
import numpy as np
import PIL
@@ -33,7 +33,7 @@ class WanVaeDecoderStep(ModularPipelineBlocks):
model_name = "wan"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKLWan),
ComponentSpec(
@@ -49,7 +49,7 @@ def description(self) -> str:
return "Step that decodes the denoised latents into images"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam(
"latents",
@@ -63,11 +63,11 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"videos",
- type_hint=Union[List[List[PIL.Image.Image]], List[torch.Tensor], List[np.ndarray]],
+ type_hint=list[list[PIL.Image.Image]] | list[torch.Tensor] | list[np.ndarray],
description="The generated videos, can be a PIL.Image.Image, torch.Tensor or a numpy array",
)
]
diff --git a/src/diffusers/modular_pipelines/wan/denoise.py b/src/diffusers/modular_pipelines/wan/denoise.py
index 7f44b0230d78..2f51f353012e 100644
--- a/src/diffusers/modular_pipelines/wan/denoise.py
+++ b/src/diffusers/modular_pipelines/wan/denoise.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List, Tuple
+from typing import Any
import torch
@@ -46,7 +46,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
"latents",
@@ -80,7 +80,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
"latents",
@@ -115,7 +115,7 @@ class WanLoopDenoiser(ModularPipelineBlocks):
def __init__(
self,
- guider_input_fields: Dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")},
+ guider_input_fields: dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")},
):
"""Initialize a denoiser block that calls the denoiser model. This block is used in Wan2.1.
@@ -136,7 +136,7 @@ def __init__(
super().__init__()
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -156,7 +156,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
inputs = [
InputParam("attention_kwargs"),
InputParam(
@@ -225,7 +225,7 @@ class Wan22LoopDenoiser(ModularPipelineBlocks):
def __init__(
self,
- guider_input_fields: Dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")},
+ guider_input_fields: dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")},
):
"""Initialize a denoiser block that calls the denoiser model. This block is used in Wan2.2.
@@ -246,7 +246,7 @@ def __init__(
super().__init__()
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -273,7 +273,7 @@ def description(self) -> str:
)
@property
- def expected_configs(self) -> List[ConfigSpec]:
+ def expected_configs(self) -> list[ConfigSpec]:
return [
ConfigSpec(
name="boundary_ratio",
@@ -283,7 +283,7 @@ def expected_configs(self) -> List[ConfigSpec]:
]
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
inputs = [
InputParam("attention_kwargs"),
InputParam(
@@ -359,7 +359,7 @@ class WanLoopAfterDenoiser(ModularPipelineBlocks):
model_name = "wan"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", UniPCMultistepScheduler),
]
@@ -400,13 +400,13 @@ def description(self) -> str:
)
@property
- def loop_expected_components(self) -> List[ComponentSpec]:
+ def loop_expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", UniPCMultistepScheduler),
]
@property
- def loop_inputs(self) -> List[InputParam]:
+ def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
"timesteps",
diff --git a/src/diffusers/modular_pipelines/wan/encoders.py b/src/diffusers/modular_pipelines/wan/encoders.py
index 22b62a601d34..3e675a66e4f2 100644
--- a/src/diffusers/modular_pipelines/wan/encoders.py
+++ b/src/diffusers/modular_pipelines/wan/encoders.py
@@ -13,7 +13,6 @@
# limitations under the License.
import html
-from typing import List, Optional, Union
import numpy as np
import PIL
@@ -62,7 +61,7 @@ def prompt_clean(text):
def get_t5_prompt_embeds(
text_encoder: UMT5EncoderModel,
tokenizer: AutoTokenizer,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
max_sequence_length: int,
device: torch.device,
):
@@ -95,7 +94,7 @@ def encode_image(
image: PipelineImageInput,
image_processor: CLIPImageProcessor,
image_encoder: CLIPVisionModel,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
image = image_processor(images=image, return_tensors="pt").to(device)
image_embeds = image_encoder(**image, output_hidden_states=True)
@@ -104,7 +103,7 @@ def encode_image(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -164,7 +163,7 @@ def description(self) -> str:
return "Text Encoder step that generate text_embeddings to guide the video generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", UMT5EncoderModel),
ComponentSpec("tokenizer", AutoTokenizer),
@@ -177,7 +176,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("prompt"),
InputParam("negative_prompt"),
@@ -185,7 +184,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
@@ -212,22 +211,22 @@ def check_inputs(block_state):
def encode_prompt(
components,
prompt: str,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
prepare_unconditional_embeds: bool = True,
- negative_prompt: Optional[str] = None,
+ negative_prompt: str | None = None,
max_sequence_length: int = 512,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
prepare_unconditional_embeds (`bool`):
whether to use prepare unconditional embeddings or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -307,7 +306,7 @@ def description(self) -> str:
return "Image Resize step that resize the image to the target area (height * width) while maintaining the aspect ratio."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("image", type_hint=PIL.Image.Image, required=True),
InputParam("height", type_hint=int, default=480),
@@ -315,7 +314,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("resized_image", type_hint=PIL.Image.Image),
]
@@ -343,7 +342,7 @@ def description(self) -> str:
return "Image Resize step that resize the last_image to the same size of first frame image with center crop."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
"resized_image", type_hint=PIL.Image.Image, required=True, description="The resized first frame image"
@@ -352,7 +351,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("resized_last_image", type_hint=PIL.Image.Image),
]
@@ -386,20 +385,20 @@ def description(self) -> str:
return "Image Encoder step that generate image_embeds based on first frame image to guide the video generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("image_processor", CLIPImageProcessor),
ComponentSpec("image_encoder", CLIPVisionModel),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("resized_image", type_hint=PIL.Image.Image, required=True),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("image_embeds", type_hint=torch.Tensor, description="The image embeddings"),
]
@@ -430,21 +429,21 @@ def description(self) -> str:
return "Image Encoder step that generate image_embeds based on first and last frame images to guide the video generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("image_processor", CLIPImageProcessor),
ComponentSpec("image_encoder", CLIPVisionModel),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("resized_image", type_hint=PIL.Image.Image, required=True),
InputParam("resized_last_image", type_hint=PIL.Image.Image, required=True),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("image_embeds", type_hint=torch.Tensor, description="The image embeddings"),
]
@@ -476,7 +475,7 @@ def description(self) -> str:
return "Vae Image Encoder step that generate condition_latents based on first frame image to guide the video generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKLWan),
ComponentSpec(
@@ -488,7 +487,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("resized_image", type_hint=PIL.Image.Image, required=True),
InputParam("height"),
@@ -498,7 +497,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"first_frame_latents",
@@ -572,16 +571,16 @@ def description(self) -> str:
return "step that prepares the masked first frame latents and add it to the latent condition"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
- InputParam("first_frame_latents", type_hint=Optional[torch.Tensor]),
+ InputParam("first_frame_latents", type_hint=torch.Tensor | None),
InputParam("num_frames", required=True),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
- OutputParam("image_condition_latents", type_hint=Optional[torch.Tensor]),
+ OutputParam("image_condition_latents", type_hint=torch.Tensor | None),
]
def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
@@ -616,7 +615,7 @@ def description(self) -> str:
return "Vae Image Encoder step that generate condition_latents based on first and last frame images to guide the video generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKLWan),
ComponentSpec(
@@ -628,7 +627,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("resized_image", type_hint=PIL.Image.Image, required=True),
InputParam("resized_last_image", type_hint=PIL.Image.Image, required=True),
@@ -639,7 +638,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"first_last_frame_latents",
@@ -721,16 +720,16 @@ def description(self) -> str:
return "step that prepares the masked latents with first and last frames and add it to the latent condition"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
- InputParam("first_last_frame_latents", type_hint=Optional[torch.Tensor]),
+ InputParam("first_last_frame_latents", type_hint=torch.Tensor | None),
InputParam("num_frames", type_hint=int, required=True),
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
- OutputParam("image_condition_latents", type_hint=Optional[torch.Tensor]),
+ OutputParam("image_condition_latents", type_hint=torch.Tensor | None),
]
def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py
index d01a86ca09b5..b641c6cd7fcc 100644
--- a/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py
+++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py
@@ -14,6 +14,7 @@
from ...utils import logging
from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
from .before_denoise import (
WanPrepareLatentsStep,
WanSetTimestepsStep,
@@ -37,7 +38,45 @@
# inputs(text) -> set_timesteps -> prepare_latents -> denoise
+# auto_docstring
class WanCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded conditions and runs the denoising process.
+
+ Components:
+ transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ num_frames (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
model_name = "wan"
block_classes = [
WanTextInputStep,
@@ -49,14 +88,11 @@ class WanCoreDenoiseStep(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "denoise block that takes encoded conditions and runs the denoising process.\n"
- + "This is a sequential pipeline blocks:\n"
- + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n"
- + " - `WanSetTimestepsStep` is used to set the timesteps\n"
- + " - `WanPrepareLatentsStep` is used to prepare the latents\n"
- + " - `WanDenoiseStep` is used to denoise the latents\n"
- )
+ return "denoise block that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
# ====================
@@ -64,7 +100,51 @@ def description(self):
# ====================
+# auto_docstring
class WanBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline blocks for Wan2.1.
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer
+ (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) vae (`AutoencoderKLWan`) video_processor
+ (`VideoProcessor`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ num_frames (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ output_type (`str`, *optional*, defaults to np):
+ The output type of the decoded videos
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
model_name = "wan"
block_classes = [
WanTextEncoderStep,
@@ -75,9 +155,8 @@ class WanBlocks(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "Modular pipeline blocks for Wan2.1.\n"
- + "- `WanTextEncoderStep` is used to encode the text\n"
- + "- `WanCoreDenoiseStep` is used to denoise the latents\n"
- + "- `WanVaeDecoderStep` is used to decode the latents to images"
- )
+ return "Modular pipeline blocks for Wan2.1."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py
index 21164422f3d9..9f602c24713b 100644
--- a/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py
+++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py
@@ -14,6 +14,7 @@
from ...utils import logging
from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
from .before_denoise import (
WanPrepareLatentsStep,
WanSetTimestepsStep,
@@ -38,7 +39,50 @@
# inputs(text) -> set_timesteps -> prepare_latents -> denoise
+# auto_docstring
class Wan22CoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded conditions and runs the denoising process.
+
+ Components:
+ transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`)
+ guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`)
+
+ Configs:
+ boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low
+ noise stages.
+
+ Inputs:
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ num_frames (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
model_name = "wan"
block_classes = [
WanTextInputStep,
@@ -50,14 +94,11 @@ class Wan22CoreDenoiseStep(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "denoise block that takes encoded conditions and runs the denoising process.\n"
- + "This is a sequential pipeline blocks:\n"
- + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n"
- + " - `WanSetTimestepsStep` is used to set the timesteps\n"
- + " - `WanPrepareLatentsStep` is used to prepare the latents\n"
- + " - `Wan22DenoiseStep` is used to denoise the latents in wan2.2\n"
- )
+ return "denoise block that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
# ====================
@@ -65,7 +106,55 @@ def description(self):
# ====================
+# auto_docstring
class Wan22Blocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline for text-to-video using Wan2.2.
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer
+ (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider_2 (`ClassifierFreeGuidance`)
+ transformer_2 (`WanTransformer3DModel`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Configs:
+ boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low
+ noise stages.
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ num_frames (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ output_type (`str`, *optional*, defaults to np):
+ The output type of the decoded videos
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
model_name = "wan"
block_classes = [
WanTextEncoderStep,
@@ -80,9 +169,8 @@ class Wan22Blocks(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "Modular pipeline for text-to-video using Wan2.2.\n"
- + " - `WanTextEncoderStep` encodes the text\n"
- + " - `Wan22CoreDenoiseStep` denoes the latents\n"
- + " - `WanVaeDecoderStep` decodes the latents to video frames\n"
- )
+ return "Modular pipeline for text-to-video using Wan2.2."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py
index 3db1c8fa837b..8e55b7a50f08 100644
--- a/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py
+++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py
@@ -14,6 +14,7 @@
from ...utils import logging
from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
from .before_denoise import (
WanAdditionalInputsStep,
WanPrepareLatentsStep,
@@ -40,7 +41,36 @@
# ====================
+# auto_docstring
class WanImage2VideoVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent
+ representation
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ first_frame_latents (`Tensor`):
+ video latent representation with the first frame image condition
+ image_condition_latents (`Tensor | NoneType`):
+ TODO: Add description.
+ """
+
model_name = "wan-i2v"
block_classes = [WanImageResizeStep, WanVaeEncoderStep, WanPrepareFirstFrameLatentsStep]
block_names = ["image_resize", "vae_encoder", "prepare_first_frame_latents"]
@@ -56,7 +86,52 @@ def description(self):
# inputs (text + image_condition_latents) -> set_timesteps -> prepare_latents -> denoise (latents)
+# auto_docstring
class Wan22Image2VideoCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded text and image latent conditions and runs the denoising process.
+
+ Components:
+ transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`)
+ guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`)
+
+ Configs:
+ boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low
+ noise stages.
+
+ Inputs:
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ num_frames (`None`, *optional*):
+ TODO: Add description.
+ image_condition_latents (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
model_name = "wan-i2v"
block_classes = [
WanTextInputStep,
@@ -75,15 +150,11 @@ class Wan22Image2VideoCoreDenoiseStep(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "denoise block that takes encoded text and image latent conditions and runs the denoising process.\n"
- + "This is a sequential pipeline blocks:\n"
- + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n"
- + " - `WanAdditionalInputsStep` is used to adjust the batch size of the latent conditions\n"
- + " - `WanSetTimestepsStep` is used to set the timesteps\n"
- + " - `WanPrepareLatentsStep` is used to prepare the latents\n"
- + " - `Wan22Image2VideoDenoiseStep` is used to denoise the latents in wan2.2\n"
- )
+ return "denoise block that takes encoded text and image latent conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
# ====================
@@ -91,7 +162,57 @@ def description(self):
# ====================
+# auto_docstring
class Wan22Image2VideoBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline for image-to-video using Wan2.2.
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`WanTransformer3DModel`) scheduler
+ (`UniPCMultistepScheduler`) guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`)
+
+ Configs:
+ boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low
+ noise stages.
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ output_type (`str`, *optional*, defaults to np):
+ The output type of the decoded videos
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
model_name = "wan-i2v"
block_classes = [
WanTextEncoderStep,
@@ -108,10 +229,8 @@ class Wan22Image2VideoBlocks(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "Modular pipeline for image-to-video using Wan2.2.\n"
- + " - `WanTextEncoderStep` encodes the text\n"
- + " - `WanImage2VideoVaeEncoderStep` encodes the image\n"
- + " - `Wan22Image2VideoCoreDenoiseStep` denoes the latents\n"
- + " - `WanVaeDecoderStep` decodes the latents to video frames\n"
- )
+ return "Modular pipeline for image-to-video using Wan2.2."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py
index d07ab8ecf473..c08db62c469a 100644
--- a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py
+++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py
@@ -14,6 +14,7 @@
from ...utils import logging
from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
from .before_denoise import (
WanAdditionalInputsStep,
WanPrepareLatentsStep,
@@ -45,7 +46,29 @@
# wan2.1 I2V (first frame only)
+# auto_docstring
class WanImage2VideoImageEncoderStep(SequentialPipelineBlocks):
+ """
+ Image2Video Image Encoder step that resize the image and encode the image to generate the image embeddings
+
+ Components:
+ image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ The image embeddings
+ """
+
model_name = "wan-i2v"
block_classes = [WanImageResizeStep, WanImageEncoderStep]
block_names = ["image_resize", "image_encoder"]
@@ -56,7 +79,34 @@ def description(self):
# wan2.1 FLF2V (first and last frame)
+# auto_docstring
class WanFLF2VImageEncoderStep(SequentialPipelineBlocks):
+ """
+ FLF2V Image Encoder step that resize and encode and encode the first and last frame images to generate the image
+ embeddings
+
+ Components:
+ image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`):
+ The last frameimage
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ resized_last_image (`Image`):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ The image embeddings
+ """
+
model_name = "wan-i2v"
block_classes = [WanImageResizeStep, WanImageCropResizeStep, WanFirstLastFrameImageEncoderStep]
block_names = ["image_resize", "last_image_resize", "image_encoder"]
@@ -67,7 +117,36 @@ def description(self):
# wan2.1 Auto Image Encoder
+# auto_docstring
class WanAutoImageEncoderStep(AutoPipelineBlocks):
+ """
+ Image Encoder step that encode the image to generate the image embeddingsThis is an auto pipeline block that works
+ for image2video tasks. - `WanFLF2VImageEncoderStep` (flf2v) is used when `last_image` is provided. -
+ `WanImage2VideoImageEncoderStep` (image2video) is used when `image` is provided. - if `last_image` or `image` is
+ not provided, step will be skipped.
+
+ Components:
+ image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`)
+
+ Inputs:
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`, *optional*):
+ The last frameimage
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ resized_last_image (`Image`):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ The image embeddings
+ """
+
block_classes = [WanFLF2VImageEncoderStep, WanImage2VideoImageEncoderStep]
block_names = ["flf2v_image_encoder", "image2video_image_encoder"]
block_trigger_inputs = ["last_image", "image"]
@@ -90,7 +169,36 @@ def description(self):
# wan2.1 I2V (first frame only)
+# auto_docstring
class WanImage2VideoVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent
+ representation
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ first_frame_latents (`Tensor`):
+ video latent representation with the first frame image condition
+ image_condition_latents (`Tensor | NoneType`):
+ TODO: Add description.
+ """
+
model_name = "wan-i2v"
block_classes = [WanImageResizeStep, WanVaeEncoderStep, WanPrepareFirstFrameLatentsStep]
block_names = ["image_resize", "vae_encoder", "prepare_first_frame_latents"]
@@ -101,7 +209,40 @@ def description(self):
# wan2.1 FLF2V (first and last frame)
+# auto_docstring
class WanFLF2VVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ FLF2V Vae Image Encoder step that resize and encode and encode the first and last frame images to generate the
+ latent conditions
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`):
+ The last frameimage
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ resized_last_image (`Image`):
+ TODO: Add description.
+ first_last_frame_latents (`Tensor`):
+ video latent representation with the first and last frame images condition
+ image_condition_latents (`Tensor | NoneType`):
+ TODO: Add description.
+ """
+
model_name = "wan-i2v"
block_classes = [
WanImageResizeStep,
@@ -117,7 +258,44 @@ def description(self):
# wan2.1 Auto Vae Encoder
+# auto_docstring
class WanAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Vae Image Encoder step that encode the image to generate the image latentsThis is an auto pipeline block that works
+ for image2video tasks. - `WanFLF2VVaeEncoderStep` (flf2v) is used when `last_image` is provided. -
+ `WanImage2VideoVaeEncoderStep` (image2video) is used when `image` is provided. - if `last_image` or `image` is not
+ provided, step will be skipped.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`, *optional*):
+ The last frameimage
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ resized_last_image (`Image`):
+ TODO: Add description.
+ first_last_frame_latents (`Tensor`):
+ video latent representation with the first and last frame images condition
+ image_condition_latents (`Tensor | NoneType`):
+ TODO: Add description.
+ first_frame_latents (`Tensor`):
+ video latent representation with the first frame image condition
+ """
+
model_name = "wan-i2v"
block_classes = [WanFLF2VVaeEncoderStep, WanImage2VideoVaeEncoderStep]
block_names = ["flf2v_vae_encoder", "image2video_vae_encoder"]
@@ -141,7 +319,53 @@ def description(self):
# wan2.1 I2V core denoise (support both I2V and FLF2V)
# inputs (text + image_condition_latents) -> set_timesteps -> prepare_latents -> denoise (latents)
+# auto_docstring
class WanImage2VideoCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded text and image latent conditions and runs the denoising process.
+
+ Components:
+ transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ num_frames (`None`, *optional*):
+ TODO: Add description.
+ image_condition_latents (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ TODO: Add description.
+
+ Outputs:
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_videos_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `transformer.dtype`)
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ """
+
model_name = "wan-i2v"
block_classes = [
WanTextInputStep,
@@ -160,15 +384,7 @@ class WanImage2VideoCoreDenoiseStep(SequentialPipelineBlocks):
@property
def description(self):
- return (
- "denoise block that takes encoded text and image latent conditions and runs the denoising process.\n"
- + "This is a sequential pipeline blocks:\n"
- + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n"
- + " - `WanAdditionalInputsStep` is used to adjust the batch size of the latent conditions\n"
- + " - `WanSetTimestepsStep` is used to set the timesteps\n"
- + " - `WanPrepareLatentsStep` is used to prepare the latents\n"
- + " - `WanImage2VideoDenoiseStep` is used to denoise the latents\n"
- )
+ return "denoise block that takes encoded text and image latent conditions and runs the denoising process."
# ====================
@@ -177,7 +393,64 @@ def description(self):
# wan2.1 Image2Video Auto Blocks
+# auto_docstring
class WanImage2VideoAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for image-to-video using Wan.
+
+ Supported workflows:
+ - `image2video`: requires `image`, `prompt`
+ - `flf2v`: requires `last_image`, `image`, `prompt`
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`)
+ image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`) vae (`AutoencoderKLWan`)
+ video_processor (`VideoProcessor`) transformer (`WanTransformer3DModel`) scheduler
+ (`UniPCMultistepScheduler`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`, *optional*):
+ The last frameimage
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_condition_latents (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ TODO: Add description.
+ output_type (`str`, *optional*, defaults to np):
+ The output type of the decoded videos
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
model_name = "wan-i2v"
block_classes = [
WanTextEncoderStep,
@@ -194,10 +467,15 @@ class WanImage2VideoAutoBlocks(SequentialPipelineBlocks):
"decode",
]
+ _workflow_map = {
+ "image2video": {"image": True, "prompt": True},
+ "flf2v": {"last_image": True, "image": True, "prompt": True},
+ }
+
@property
def description(self):
- return (
- "Auto Modular pipeline for image-to-video using Wan.\n"
- + "- for I2V workflow, all you need to provide is `image`"
- + "- for FLF2V workflow, all you need to provide is `last_image` and `image`"
- )
+ return "Auto Modular pipeline for image-to-video using Wan."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/src/diffusers/modular_pipelines/z_image/__init__.py b/src/diffusers/modular_pipelines/z_image/__init__.py
index c8a8c14396c0..5c04008d3305 100644
--- a/src/diffusers/modular_pipelines/z_image/__init__.py
+++ b/src/diffusers/modular_pipelines/z_image/__init__.py
@@ -21,12 +21,7 @@
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
- _import_structure["decoders"] = ["ZImageVaeDecoderStep"]
- _import_structure["encoders"] = ["ZImageTextEncoderStep", "ZImageVaeImageEncoderStep"]
- _import_structure["modular_blocks"] = [
- "ALL_BLOCKS",
- "ZImageAutoBlocks",
- ]
+ _import_structure["modular_blocks_z_image"] = ["ZImageAutoBlocks"]
_import_structure["modular_pipeline"] = ["ZImageModularPipeline"]
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
@@ -36,12 +31,7 @@
except OptionalDependencyNotAvailable:
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
else:
- from .decoders import ZImageVaeDecoderStep
- from .encoders import ZImageTextEncoderStep
- from .modular_blocks import (
- ALL_BLOCKS,
- ZImageAutoBlocks,
- )
+ from .modular_blocks_z_image import ZImageAutoBlocks
from .modular_pipeline import ZImageModularPipeline
else:
import sys
diff --git a/src/diffusers/modular_pipelines/z_image/before_denoise.py b/src/diffusers/modular_pipelines/z_image/before_denoise.py
index 35ea768f12c3..8558f2c67f65 100644
--- a/src/diffusers/modular_pipelines/z_image/before_denoise.py
+++ b/src/diffusers/modular_pipelines/z_image/before_denoise.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Tuple, Union
import torch
@@ -91,7 +90,7 @@ def repeat_tensor_to_batch_size(
return input_tensor
-def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor_spatial: int) -> Tuple[int, int]:
+def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor_spatial: int) -> tuple[int, int]:
"""Calculate image dimensions from latent tensor dimensions.
This function converts latent spatial dimensions to image spatial dimensions by multiplying the latent height/width
@@ -103,7 +102,7 @@ def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor_spa
vae_scale_factor (int): The scale factor used by the VAE to compress image spatial dimension.
By default, it is 16
Returns:
- Tuple[int, int]: The calculated image dimensions as (height, width)
+ tuple[int, int]: The calculated image dimensions as (height, width)
"""
latent_height, latent_width = latents.shape[2:]
height = latent_height * vae_scale_factor_spatial // 2
@@ -129,10 +128,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -147,15 +146,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -201,30 +200,30 @@ def description(self) -> str:
)
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("transformer", ZImageTransformer2DModel),
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("num_images_per_prompt", default=1),
InputParam(
"prompt_embeds",
required=True,
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
description="Pre-generated text embeddings. Can be generated from text_encoder step.",
),
InputParam(
"negative_prompt_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
description="Pre-generated negative text embeddings. Can be generated from text_encoder step.",
),
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"batch_size",
@@ -283,8 +282,8 @@ class ZImageAdditionalInputsStep(ModularPipelineBlocks):
def __init__(
self,
- image_latent_inputs: List[str] = ["image_latents"],
- additional_batch_inputs: List[str] = [],
+ image_latent_inputs: list[str] = ["image_latents"],
+ additional_batch_inputs: list[str] = [],
):
"""Initialize a configurable step that standardizes the inputs for the denoising step. It:\n"
@@ -295,10 +294,10 @@ def __init__(
This is a dynamic block that allows you to configure which inputs to process.
Args:
- image_latent_inputs (List[str], optional): Names of image latent tensors to process.
+ image_latent_inputs (list[str], optional): Names of image latent tensors to process.
In additional to adjust batch size of these inputs, they will be used to determine height/width. Can be
a single string or list of strings. Defaults to ["image_latents"].
- additional_batch_inputs (List[str], optional):
+ additional_batch_inputs (list[str], optional):
Names of additional conditional input tensors to expand batch size. These tensors will only have their
batch dimensions adjusted to match the final batch size. Can be a single string or list of strings.
Defaults to [].
@@ -346,7 +345,7 @@ def description(self) -> str:
return summary_section + inputs_info + placement_section
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
inputs = [
InputParam(name="num_images_per_prompt", default=1),
InputParam(name="batch_size", required=True),
@@ -406,11 +405,11 @@ def description(self) -> str:
return "Prepare latents step that prepares the latents for the text-to-video generation process"
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("height", type_hint=int),
InputParam("width", type_hint=int),
- InputParam("latents", type_hint=Optional[torch.Tensor]),
+ InputParam("latents", type_hint=torch.Tensor | None),
InputParam("num_images_per_prompt", type_hint=int, default=1),
InputParam("generator"),
InputParam(
@@ -423,7 +422,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
@@ -496,7 +495,7 @@ class ZImageSetTimestepsStep(ModularPipelineBlocks):
model_name = "z-image"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@@ -506,7 +505,7 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for inference. Need to run after prepare latents step."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("latents", required=True),
InputParam("num_inference_steps", default=9),
@@ -514,7 +513,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process"
@@ -554,7 +553,7 @@ class ZImageSetTimestepsWithStrengthStep(ModularPipelineBlocks):
model_name = "z-image"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@@ -564,7 +563,7 @@ def description(self) -> str:
return "Step that sets the scheduler's timesteps for inference with strength. Need to run after set timesteps step."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("timesteps", required=True),
InputParam("num_inference_steps", required=True),
@@ -602,7 +601,7 @@ def description(self) -> str:
return "step that prepares the latents with image condition, need to run after set timesteps and prepare latents step."
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("latents", required=True),
InputParam("image_latents", required=True),
diff --git a/src/diffusers/modular_pipelines/z_image/decoders.py b/src/diffusers/modular_pipelines/z_image/decoders.py
index cdb6a2e5eac1..353253102376 100644
--- a/src/diffusers/modular_pipelines/z_image/decoders.py
+++ b/src/diffusers/modular_pipelines/z_image/decoders.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, List, Tuple, Union
+from typing import Any
import numpy as np
import PIL
@@ -33,7 +33,7 @@ class ZImageVaeDecoderStep(ModularPipelineBlocks):
model_name = "z-image"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKL),
ComponentSpec(
@@ -49,7 +49,7 @@ def description(self) -> str:
return "Step that decodes the denoised latents into images"
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam(
"latents",
@@ -64,11 +64,11 @@ def inputs(self) -> List[Tuple[str, Any]]:
]
@property
- def intermediate_outputs(self) -> List[str]:
+ def intermediate_outputs(self) -> list[str]:
return [
OutputParam(
"images",
- type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]],
+ type_hint=list[PIL.Image.Image, list[torch.Tensor], list[np.ndarray]],
description="The generated images, can be a PIL.Image.Image, torch.Tensor or a numpy array",
)
]
diff --git a/src/diffusers/modular_pipelines/z_image/denoise.py b/src/diffusers/modular_pipelines/z_image/denoise.py
index 5f76a8459fde..863df312389a 100644
--- a/src/diffusers/modular_pipelines/z_image/denoise.py
+++ b/src/diffusers/modular_pipelines/z_image/denoise.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict, List, Tuple
+from typing import Any
import torch
@@ -46,7 +46,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam(
"latents",
@@ -80,7 +80,7 @@ class ZImageLoopDenoiser(ModularPipelineBlocks):
def __init__(
self,
- guider_input_fields: Dict[str, Any] = {"cap_feats": ("prompt_embeds", "negative_prompt_embeds")},
+ guider_input_fields: dict[str, Any] = {"cap_feats": ("prompt_embeds", "negative_prompt_embeds")},
):
"""Initialize a denoiser block that calls the denoiser model. This block is used in Z-Image.
@@ -101,7 +101,7 @@ def __init__(
super().__init__()
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
@@ -121,7 +121,7 @@ def description(self) -> str:
)
@property
- def inputs(self) -> List[Tuple[str, Any]]:
+ def inputs(self) -> list[tuple[str, Any]]:
inputs = [
InputParam(
"num_inference_steps",
@@ -205,7 +205,7 @@ class ZImageLoopAfterDenoiser(ModularPipelineBlocks):
model_name = "z-image"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@@ -246,13 +246,13 @@ def description(self) -> str:
)
@property
- def loop_expected_components(self) -> List[ComponentSpec]:
+ def loop_expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
]
@property
- def loop_inputs(self) -> List[InputParam]:
+ def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
"timesteps",
diff --git a/src/diffusers/modular_pipelines/z_image/encoders.py b/src/diffusers/modular_pipelines/z_image/encoders.py
index f5769fe2deec..06deb8236893 100644
--- a/src/diffusers/modular_pipelines/z_image/encoders.py
+++ b/src/diffusers/modular_pipelines/z_image/encoders.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Union
import PIL
import torch
@@ -37,10 +36,10 @@
def get_qwen_prompt_embeds(
text_encoder: Qwen3Model,
tokenizer: Qwen2Tokenizer,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
max_sequence_length: int = 512,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
prompt = [prompt] if isinstance(prompt, str) else prompt
for i, prompt_item in enumerate(prompt):
@@ -82,7 +81,7 @@ def get_qwen_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -134,7 +133,7 @@ def description(self) -> str:
return "Text Encoder step that generate text_embeddings to guide the video generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Qwen3Model),
ComponentSpec("tokenizer", Qwen2Tokenizer),
@@ -147,7 +146,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("prompt"),
InputParam("negative_prompt"),
@@ -155,17 +154,17 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
kwargs_type="denoiser_input_fields",
description="text embeddings used to guide the image generation",
),
OutputParam(
"negative_prompt_embeds",
- type_hint=List[torch.Tensor],
+ type_hint=list[torch.Tensor],
kwargs_type="denoiser_input_fields",
description="negative text embeddings used to guide the image generation",
),
@@ -182,22 +181,22 @@ def check_inputs(block_state):
def encode_prompt(
components,
prompt: str,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
prepare_unconditional_embeds: bool = True,
- negative_prompt: Optional[str] = None,
+ negative_prompt: str | None = None,
max_sequence_length: int = 512,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
prepare_unconditional_embeds (`bool`):
whether to use prepare unconditional embeddings or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -278,7 +277,7 @@ def description(self) -> str:
return "Vae Image Encoder step that generate condition_latents based on image to guide the image generation"
@property
- def expected_components(self) -> List[ComponentSpec]:
+ def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKL),
ComponentSpec(
@@ -290,7 +289,7 @@ def expected_components(self) -> List[ComponentSpec]:
]
@property
- def inputs(self) -> List[InputParam]:
+ def inputs(self) -> list[InputParam]:
return [
InputParam("image", type_hint=PIL.Image.Image, required=True),
InputParam("height"),
@@ -299,7 +298,7 @@ def inputs(self) -> List[InputParam]:
]
@property
- def intermediate_outputs(self) -> List[OutputParam]:
+ def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"image_latents",
diff --git a/src/diffusers/modular_pipelines/z_image/modular_blocks.py b/src/diffusers/modular_pipelines/z_image/modular_blocks.py
deleted file mode 100644
index a54baeccaf0c..000000000000
--- a/src/diffusers/modular_pipelines/z_image/modular_blocks.py
+++ /dev/null
@@ -1,191 +0,0 @@
-# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from ...utils import logging
-from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
-from ..modular_pipeline_utils import InsertableDict
-from .before_denoise import (
- ZImageAdditionalInputsStep,
- ZImagePrepareLatentsStep,
- ZImagePrepareLatentswithImageStep,
- ZImageSetTimestepsStep,
- ZImageSetTimestepsWithStrengthStep,
- ZImageTextInputStep,
-)
-from .decoders import ZImageVaeDecoderStep
-from .denoise import (
- ZImageDenoiseStep,
-)
-from .encoders import (
- ZImageTextEncoderStep,
- ZImageVaeImageEncoderStep,
-)
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-
-# z-image
-# text2image
-class ZImageCoreDenoiseStep(SequentialPipelineBlocks):
- block_classes = [
- ZImageTextInputStep,
- ZImagePrepareLatentsStep,
- ZImageSetTimestepsStep,
- ZImageDenoiseStep,
- ]
- block_names = ["input", "prepare_latents", "set_timesteps", "denoise"]
-
- @property
- def description(self):
- return (
- "denoise block that takes encoded conditions and runs the denoising process.\n"
- + "This is a sequential pipeline blocks:\n"
- + " - `ZImageTextInputStep` is used to adjust the batch size of the model inputs\n"
- + " - `ZImagePrepareLatentsStep` is used to prepare the latents\n"
- + " - `ZImageSetTimestepsStep` is used to set the timesteps\n"
- + " - `ZImageDenoiseStep` is used to denoise the latents\n"
- )
-
-
-# z-image: image2image
-## denoise
-class ZImageImage2ImageCoreDenoiseStep(SequentialPipelineBlocks):
- block_classes = [
- ZImageTextInputStep,
- ZImageAdditionalInputsStep(image_latent_inputs=["image_latents"]),
- ZImagePrepareLatentsStep,
- ZImageSetTimestepsStep,
- ZImageSetTimestepsWithStrengthStep,
- ZImagePrepareLatentswithImageStep,
- ZImageDenoiseStep,
- ]
- block_names = [
- "input",
- "additional_inputs",
- "prepare_latents",
- "set_timesteps",
- "set_timesteps_with_strength",
- "prepare_latents_with_image",
- "denoise",
- ]
-
- @property
- def description(self):
- return (
- "denoise block that takes encoded text and image latent conditions and runs the denoising process.\n"
- + "This is a sequential pipeline blocks:\n"
- + " - `ZImageTextInputStep` is used to adjust the batch size of the model inputs\n"
- + " - `ZImageAdditionalInputsStep` is used to adjust the batch size of the latent conditions\n"
- + " - `ZImagePrepareLatentsStep` is used to prepare the latents\n"
- + " - `ZImageSetTimestepsStep` is used to set the timesteps\n"
- + " - `ZImageSetTimestepsWithStrengthStep` is used to set the timesteps with strength\n"
- + " - `ZImagePrepareLatentswithImageStep` is used to prepare the latents with image\n"
- + " - `ZImageDenoiseStep` is used to denoise the latents\n"
- )
-
-
-## auto blocks
-class ZImageAutoDenoiseStep(AutoPipelineBlocks):
- block_classes = [
- ZImageImage2ImageCoreDenoiseStep,
- ZImageCoreDenoiseStep,
- ]
- block_names = ["image2image", "text2image"]
- block_trigger_inputs = ["image_latents", None]
-
- @property
- def description(self) -> str:
- return (
- "Denoise step that iteratively denoise the latents. "
- "This is a auto pipeline block that works for text2image and image2image tasks."
- " - `ZImageCoreDenoiseStep` (text2image) for text2image tasks."
- " - `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks."
- + " - if `image_latents` is provided, `ZImageImage2ImageCoreDenoiseStep` will be used.\n"
- + " - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used.\n"
- )
-
-
-class ZImageAutoVaeImageEncoderStep(AutoPipelineBlocks):
- block_classes = [ZImageVaeImageEncoderStep]
- block_names = ["vae_encoder"]
- block_trigger_inputs = ["image"]
-
- @property
- def description(self) -> str:
- return "Vae Image Encoder step that encode the image to generate the image latents"
- +"This is an auto pipeline block that works for image2image tasks."
- +" - `ZImageVaeImageEncoderStep` is used when `image` is provided."
- +" - if `image` is not provided, step will be skipped."
-
-
-class ZImageAutoBlocks(SequentialPipelineBlocks):
- block_classes = [
- ZImageTextEncoderStep,
- ZImageAutoVaeImageEncoderStep,
- ZImageAutoDenoiseStep,
- ZImageVaeDecoderStep,
- ]
- block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
-
- @property
- def description(self) -> str:
- return "Auto Modular pipeline for text-to-image and image-to-image using ZImage.\n"
- +" - for text-to-image generation, all you need to provide is `prompt`\n"
- +" - for image-to-image generation, you need to provide `image`\n"
- +" - if `image` is not provided, step will be skipped."
-
-
-# presets
-TEXT2IMAGE_BLOCKS = InsertableDict(
- [
- ("text_encoder", ZImageTextEncoderStep),
- ("input", ZImageTextInputStep),
- ("prepare_latents", ZImagePrepareLatentsStep),
- ("set_timesteps", ZImageSetTimestepsStep),
- ("denoise", ZImageDenoiseStep),
- ("decode", ZImageVaeDecoderStep),
- ]
-)
-
-IMAGE2IMAGE_BLOCKS = InsertableDict(
- [
- ("text_encoder", ZImageTextEncoderStep),
- ("vae_encoder", ZImageVaeImageEncoderStep),
- ("input", ZImageTextInputStep),
- ("additional_inputs", ZImageAdditionalInputsStep(image_latent_inputs=["image_latents"])),
- ("prepare_latents", ZImagePrepareLatentsStep),
- ("set_timesteps", ZImageSetTimestepsStep),
- ("set_timesteps_with_strength", ZImageSetTimestepsWithStrengthStep),
- ("prepare_latents_with_image", ZImagePrepareLatentswithImageStep),
- ("denoise", ZImageDenoiseStep),
- ("decode", ZImageVaeDecoderStep),
- ]
-)
-
-
-AUTO_BLOCKS = InsertableDict(
- [
- ("text_encoder", ZImageTextEncoderStep),
- ("vae_encoder", ZImageAutoVaeImageEncoderStep),
- ("denoise", ZImageAutoDenoiseStep),
- ("decode", ZImageVaeDecoderStep),
- ]
-)
-
-ALL_BLOCKS = {
- "text2image": TEXT2IMAGE_BLOCKS,
- "image2image": IMAGE2IMAGE_BLOCKS,
- "auto": AUTO_BLOCKS,
-}
diff --git a/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py b/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py
new file mode 100644
index 000000000000..23e20d55fb1e
--- /dev/null
+++ b/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py
@@ -0,0 +1,334 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ ZImageAdditionalInputsStep,
+ ZImagePrepareLatentsStep,
+ ZImagePrepareLatentswithImageStep,
+ ZImageSetTimestepsStep,
+ ZImageSetTimestepsWithStrengthStep,
+ ZImageTextInputStep,
+)
+from .decoders import ZImageVaeDecoderStep
+from .denoise import (
+ ZImageDenoiseStep,
+)
+from .encoders import (
+ ZImageTextEncoderStep,
+ ZImageVaeImageEncoderStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# ====================
+# 1. DENOISE
+# ====================
+
+
+# text2image: inputs(text) -> set_timesteps -> prepare_latents -> denoise
+# auto_docstring
+class ZImageCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded conditions and runs the denoising process.
+
+ Components:
+ transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`list`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`list`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 9):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [
+ ZImageTextInputStep,
+ ZImagePrepareLatentsStep,
+ ZImageSetTimestepsStep,
+ ZImageDenoiseStep,
+ ]
+ block_names = ["input", "prepare_latents", "set_timesteps", "denoise"]
+
+ @property
+ def description(self):
+ return "denoise block that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# image2image: inputs(text + image_latents) -> prepare_latents -> set_timesteps -> set_timesteps_with_strength -> prepare_latents_with_image -> denoise
+# auto_docstring
+class ZImageImage2ImageCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded text and image latent conditions and runs the denoising process.
+
+ Components:
+ transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`list`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`list`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 9):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [
+ ZImageTextInputStep,
+ ZImageAdditionalInputsStep(image_latent_inputs=["image_latents"]),
+ ZImagePrepareLatentsStep,
+ ZImageSetTimestepsStep,
+ ZImageSetTimestepsWithStrengthStep,
+ ZImagePrepareLatentswithImageStep,
+ ZImageDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "prepare_latents",
+ "set_timesteps",
+ "set_timesteps_with_strength",
+ "prepare_latents_with_image",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "denoise block that takes encoded text and image latent conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class ZImageAutoDenoiseStep(AutoPipelineBlocks):
+ """
+ Denoise step that iteratively denoise the latents. This is a auto pipeline block that works for text2image and
+ image2image tasks. - `ZImageCoreDenoiseStep` (text2image) for text2image tasks. -
+ `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks. - if `image_latents` is provided,
+ `ZImageImage2ImageCoreDenoiseStep` will be used.
+ - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used.
+
+ Components:
+ transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`list`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`list`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [
+ ZImageImage2ImageCoreDenoiseStep,
+ ZImageCoreDenoiseStep,
+ ]
+ block_names = ["image2image", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. "
+ "This is a auto pipeline block that works for text2image and image2image tasks."
+ " - `ZImageCoreDenoiseStep` (text2image) for text2image tasks."
+ " - `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks."
+ + " - if `image_latents` is provided, `ZImageImage2ImageCoreDenoiseStep` will be used.\n"
+ + " - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used.\n"
+ )
+
+
+# auto_docstring
+class ZImageAutoVaeImageEncoderStep(AutoPipelineBlocks):
+ """
+ Vae Image Encoder step that encode the image to generate the image latents
+
+ Components:
+ vae (`AutoencoderKL`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ image_latents (`Tensor`):
+ video latent representation with the first frame image condition
+ """
+
+ block_classes = [ZImageVaeImageEncoderStep]
+ block_names = ["vae_encoder"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self) -> str:
+ return "Vae Image Encoder step that encode the image to generate the image latents"
+ +"This is an auto pipeline block that works for image2image tasks."
+ +" - `ZImageVaeImageEncoderStep` is used when `image` is provided."
+ +" - if `image` is not provided, step will be skipped."
+
+
+# auto_docstring
+class ZImageAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image and image-to-image using ZImage.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Qwen3Model`) tokenizer (`Qwen2Tokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKL`) image_processor (`VaeImageProcessor`) transformer (`ZImageTransformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ output_type (`str`, *optional*, defaults to pil):
+ The type of the output images, can be 'pil', 'np', 'pt'
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ block_classes = [
+ ZImageTextEncoderStep,
+ ZImageAutoVaeImageEncoderStep,
+ ZImageAutoDenoiseStep,
+ ZImageVaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self) -> str:
+ return "Auto Modular pipeline for text-to-image and image-to-image using ZImage."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/src/diffusers/optimization.py b/src/diffusers/optimization.py
index e0b3576e4426..044bb0db1908 100644
--- a/src/diffusers/optimization.py
+++ b/src/diffusers/optimization.py
@@ -16,7 +16,6 @@
import math
from enum import Enum
-from typing import Optional, Union
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR
@@ -287,11 +286,11 @@ def lr_lambda(current_step: int):
def get_scheduler(
- name: Union[str, SchedulerType],
+ name: str | SchedulerType,
optimizer: Optimizer,
- step_rules: Optional[str] = None,
- num_warmup_steps: Optional[int] = None,
- num_training_steps: Optional[int] = None,
+ step_rules: str | None = None,
+ num_warmup_steps: int | None = None,
+ num_training_steps: int | None = None,
num_cycles: int = 1,
power: float = 1.0,
last_epoch: int = -1,
diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py
index 72923cbb5c18..638598051d64 100644
--- a/src/diffusers/pipelines/__init__.py
+++ b/src/diffusers/pipelines/__init__.py
@@ -6,7 +6,6 @@
_LazyModule,
get_objects_from_module,
is_flax_available,
- is_k_diffusion_available,
is_librosa_available,
is_note_seq_available,
is_onnx_available,
@@ -167,6 +166,7 @@
_import_structure["consisid"] = ["ConsisIDPipeline"]
_import_structure["cosmos"] = [
"Cosmos2_5_PredictBasePipeline",
+ "Cosmos2_5_TransferPipeline",
"Cosmos2TextToImagePipeline",
"CosmosTextToWorldPipeline",
"CosmosVideoToWorldPipeline",
@@ -465,21 +465,6 @@
]
)
-try:
- if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- from ..utils import (
- dummy_torch_and_transformers_and_k_diffusion_objects,
- )
-
- _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_k_diffusion_objects))
-else:
- _import_structure["stable_diffusion_k_diffusion"] = [
- "StableDiffusionKDiffusionPipeline",
- "StableDiffusionXLKDiffusionPipeline",
- ]
-
try:
if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
raise OptionalDependencyNotAvailable()
@@ -631,6 +616,7 @@
)
from .cosmos import (
Cosmos2_5_PredictBasePipeline,
+ Cosmos2_5_TransferPipeline,
Cosmos2TextToImagePipeline,
Cosmos2VideoToWorldPipeline,
CosmosTextToWorldPipeline,
@@ -899,17 +885,6 @@
StableDiffusionOnnxPipeline,
)
- try:
- if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- from ..utils.dummy_torch_and_transformers_and_k_diffusion_objects import *
- else:
- from .stable_diffusion_k_diffusion import (
- StableDiffusionKDiffusionPipeline,
- StableDiffusionXLKDiffusionPipeline,
- )
-
try:
if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
raise OptionalDependencyNotAvailable()
diff --git a/src/diffusers/pipelines/allegro/pipeline_allegro.py b/src/diffusers/pipelines/allegro/pipeline_allegro.py
index 42083378d465..e54e9ed20739 100644
--- a/src/diffusers/pipelines/allegro/pipeline_allegro.py
+++ b/src/diffusers/pipelines/allegro/pipeline_allegro.py
@@ -18,7 +18,7 @@
import math
import re
import urllib.parse as ul
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import T5EncoderModel, T5Tokenizer
@@ -84,10 +84,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -102,15 +102,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -214,15 +214,15 @@ def __init__(
# Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->512, num_images_per_prompt->num_videos_per_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 512,
**kwargs,
@@ -231,9 +231,9 @@ def encode_prompt(
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -719,46 +719,44 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 100,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 7.5,
- num_frames: Optional[int] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ num_frames: int | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_videos_per_prompt: int = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
clean_caption: bool = True,
max_sequence_length: int = 512,
- ) -> Union[AllegroPipelineOutput, Tuple]:
+ ) -> AllegroPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the video generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality video at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 7.5):
@@ -778,7 +776,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/allegro/pipeline_output.py b/src/diffusers/pipelines/allegro/pipeline_output.py
index 6a721783ca86..bf85a4954ce9 100644
--- a/src/diffusers/pipelines/allegro/pipeline_output.py
+++ b/src/diffusers/pipelines/allegro/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL
@@ -14,10 +13,10 @@ class AllegroPipelineOutput(BaseOutput):
Output class for Allegro pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
- frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]]
+ frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]
diff --git a/src/diffusers/pipelines/amused/pipeline_amused.py b/src/diffusers/pipelines/amused/pipeline_amused.py
index 131e34d1a4a1..b23adf0d2152 100644
--- a/src/diffusers/pipelines/amused/pipeline_amused.py
+++ b/src/diffusers/pipelines/amused/pipeline_amused.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
@@ -84,33 +84,33 @@ def __init__(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[List[str], str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: list[str] | str | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 12,
guidance_scale: float = 10.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.IntTensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_encoder_hidden_states: Optional[torch.Tensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | None = None,
+ latents: torch.IntTensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_encoder_hidden_states: torch.Tensor | None = None,
output_type="pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
micro_conditioning_aesthetic_score: int = 6,
- micro_conditioning_crop_coord: Tuple[int, int] = (0, 0),
- temperature: Union[int, Tuple[int, int], List[int]] = (2, 0),
+ micro_conditioning_crop_coord: tuple[int, int] = (0, 0),
+ temperature: int | tuple[int, int] | list[int] = (2, 0),
):
"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.transformer.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -122,7 +122,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 10.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -162,10 +162,10 @@ def __call__(
The targeted aesthetic score according to the laion aesthetic classifier. See
https://laion.ai/blog/laion-aesthetics/ and the micro-conditioning section of
https://huggingface.co/papers/2307.01952.
- micro_conditioning_crop_coord (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ micro_conditioning_crop_coord (`tuple[int]`, *optional*, defaults to (0, 0)):
The targeted height, width crop coordinates. See the micro-conditioning section of
https://huggingface.co/papers/2307.01952.
- temperature (`Union[int, Tuple[int, int], List[int]]`, *optional*, defaults to (2, 0)):
+ temperature (`int | tuple[int, int, list[int]]`, *optional*, defaults to (2, 0)):
Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`.
Examples:
diff --git a/src/diffusers/pipelines/amused/pipeline_amused_img2img.py b/src/diffusers/pipelines/amused/pipeline_amused_img2img.py
index a122c12236dd..79ebd96dedeb 100644
--- a/src/diffusers/pipelines/amused/pipeline_amused_img2img.py
+++ b/src/diffusers/pipelines/amused/pipeline_amused_img2img.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
@@ -99,34 +99,34 @@ def __init__(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[List[str], str]] = None,
+ prompt: list[str] | str | None = None,
image: PipelineImageInput = None,
strength: float = 0.5,
num_inference_steps: int = 12,
guidance_scale: float = 10.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[torch.Generator] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_encoder_hidden_states: Optional[torch.Tensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_encoder_hidden_states: torch.Tensor | None = None,
output_type="pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
micro_conditioning_aesthetic_score: int = 6,
- micro_conditioning_crop_coord: Tuple[int, int] = (0, 0),
- temperature: Union[int, Tuple[int, int], List[int]] = (2, 0),
+ micro_conditioning_crop_coord: tuple[int, int] = (0, 0),
+ temperature: int | tuple[int, int] | list[int] = (2, 0),
):
"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -144,7 +144,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 10.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -181,10 +181,10 @@ def __call__(
The targeted aesthetic score according to the laion aesthetic classifier. See
https://laion.ai/blog/laion-aesthetics/ and the micro-conditioning section of
https://huggingface.co/papers/2307.01952.
- micro_conditioning_crop_coord (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ micro_conditioning_crop_coord (`tuple[int]`, *optional*, defaults to (0, 0)):
The targeted height, width crop coordinates. See the micro-conditioning section of
https://huggingface.co/papers/2307.01952.
- temperature (`Union[int, Tuple[int, int], List[int]]`, *optional*, defaults to (2, 0)):
+ temperature (`int | tuple[int, int, list[int]]`, *optional*, defaults to (2, 0)):
Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`.
Examples:
diff --git a/src/diffusers/pipelines/amused/pipeline_amused_inpaint.py b/src/diffusers/pipelines/amused/pipeline_amused_inpaint.py
index f4bd4944ff9a..55302401832c 100644
--- a/src/diffusers/pipelines/amused/pipeline_amused_inpaint.py
+++ b/src/diffusers/pipelines/amused/pipeline_amused_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
@@ -115,41 +115,41 @@ def __init__(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[List[str], str]] = None,
+ prompt: list[str] | str | None = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
strength: float = 1.0,
num_inference_steps: int = 12,
guidance_scale: float = 10.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[torch.Generator] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_encoder_hidden_states: Optional[torch.Tensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_encoder_hidden_states: torch.Tensor | None = None,
output_type="pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
micro_conditioning_aesthetic_score: int = 6,
- micro_conditioning_crop_coord: Tuple[int, int] = (0, 0),
- temperature: Union[int, Tuple[int, int], List[int]] = (2, 0),
+ micro_conditioning_crop_coord: tuple[int, int] = (0, 0),
+ temperature: int | tuple[int, int] | list[int] = (2, 0),
):
"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
@@ -168,7 +168,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 10.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -205,10 +205,10 @@ def __call__(
The targeted aesthetic score according to the laion aesthetic classifier. See
https://laion.ai/blog/laion-aesthetics/ and the micro-conditioning section of
https://huggingface.co/papers/2307.01952.
- micro_conditioning_crop_coord (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ micro_conditioning_crop_coord (`tuple[int]`, *optional*, defaults to (0, 0)):
The targeted height, width crop coordinates. See the micro-conditioning section of
https://huggingface.co/papers/2307.01952.
- temperature (`Union[int, Tuple[int, int], List[int]]`, *optional*, defaults to (2, 0)):
+ temperature (`int | tuple[int, int, list[int]]`, *optional*, defaults to (2, 0)):
Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`.
Examples:
diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff.py
index 091b6db713ba..4d7477bc8754 100644
--- a/src/diffusers/pipelines/animatediff/pipeline_animatediff.py
+++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
@@ -122,16 +122,14 @@ def __init__(
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetMotionModel],
+ unet: UNet2DConditionModel | UNetMotionModel,
motion_adapter: MotionAdapter,
- scheduler: Union[
- DDIMScheduler,
- PNDMScheduler,
- LMSDiscreteScheduler,
- EulerDiscreteScheduler,
- EulerAncestralDiscreteScheduler,
- DPMSolverMultistepScheduler,
- ],
+ scheduler: DDIMScheduler
+ | PNDMScheduler
+ | LMSDiscreteScheduler
+ | EulerDiscreteScheduler
+ | EulerAncestralDiscreteScheduler
+ | DPMSolverMultistepScheduler,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
):
@@ -160,16 +158,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -177,7 +175,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -276,7 +274,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -574,27 +572,27 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- num_frames: Optional[int] = 16,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ num_frames: int | None = 16,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
decode_chunk_size: int = 16,
**kwargs,
):
@@ -602,7 +600,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated video.
@@ -617,13 +615,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -639,7 +637,7 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -660,7 +658,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py
index 70180ccf0650..eb511129cc6f 100644
--- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py
+++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
import torch.nn.functional as F
@@ -164,12 +164,12 @@ def __init__(
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetMotionModel],
+ unet: UNet2DConditionModel | UNetMotionModel,
motion_adapter: MotionAdapter,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
- feature_extractor: Optional[CLIPImageProcessor] = None,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ feature_extractor: CLIPImageProcessor | None = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
):
super().__init__()
if isinstance(unet, UNet2DConditionModel):
@@ -203,16 +203,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -220,7 +220,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -319,7 +319,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -721,39 +721,39 @@ def interrupt(self):
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- num_frames: Optional[int] = 16,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ num_frames: int | None = 16,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[PipelineImageInput] = None,
- conditioning_frames: Optional[List[PipelineImageInput]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: PipelineImageInput | None = None,
+ conditioning_frames: list[PipelineImageInput] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
decode_chunk_size: int = 16,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated video.
@@ -768,13 +768,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -790,12 +790,12 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
- conditioning_frames (`List[PipelineImageInput]`, *optional*):
+ conditioning_frames (`list[PipelineImageInput]`, *optional*):
The ControlNet input condition to provide guidance to the `unet` for generation. If multiple
ControlNets are specified, images must be passed as a list such that each element of the list can be
correctly batched for input to a single ControlNet.
@@ -807,16 +807,16 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
@@ -826,7 +826,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py
index 56d319027595..68ce7c92896a 100644
--- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py
+++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import (
@@ -150,10 +150,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -168,15 +168,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -284,16 +284,14 @@ def __init__(
text_encoder_2: CLIPTextModelWithProjection,
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetMotionModel],
+ unet: UNet2DConditionModel | UNetMotionModel,
motion_adapter: MotionAdapter,
- scheduler: Union[
- DDIMScheduler,
- PNDMScheduler,
- LMSDiscreteScheduler,
- EulerDiscreteScheduler,
- EulerAncestralDiscreteScheduler,
- DPMSolverMultistepScheduler,
- ],
+ scheduler: DDIMScheduler
+ | PNDMScheduler
+ | LMSDiscreteScheduler
+ | EulerDiscreteScheduler
+ | EulerAncestralDiscreteScheduler
+ | DPMSolverMultistepScheduler,
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
force_zeros_for_empty_prompt: bool = True,
@@ -329,26 +327,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_videos_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -357,11 +355,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -479,7 +477,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -870,50 +868,50 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
num_frames: int = 16,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
num_frames:
@@ -932,11 +930,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality video at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -953,11 +951,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower video quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the video generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the video generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_videos_per_prompt (`int`, *optional*, defaults to 1):
@@ -965,7 +963,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -988,7 +986,7 @@ def __call__(
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. If not provided, embeddings are computed from the
`ip_adapter_image` input argument.
output_type (`str`, *optional*, defaults to `"pil"`):
@@ -1007,31 +1005,31 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1041,7 +1039,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py
index 46d650efe8b6..8b3eb8fc3c03 100644
--- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py
+++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -117,7 +117,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -176,7 +176,7 @@ def __init__(
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetMotionModel],
+ unet: UNet2DConditionModel | UNetMotionModel,
motion_adapter: MotionAdapter,
controlnet: SparseControlNetModel,
scheduler: KarrasDiffusionSchedulers,
@@ -212,16 +212,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -229,7 +229,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -328,7 +328,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -675,7 +675,7 @@ def prepare_sparse_control_conditioning(
controlnet_frame_indices: int,
device: torch.device,
dtype: torch.dtype,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
assert conditioning_frames.shape[2] >= len(controlnet_frame_indices)
batch_size, channels, _, height, width = conditioning_frames.shape
@@ -713,37 +713,37 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_frames: int = 16,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
num_videos_per_prompt: int = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- conditioning_frames: Optional[List[PipelineImageInput]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ conditioning_frames: list[PipelineImageInput] | None = None,
output_type: str = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- controlnet_frame_indices: List[int] = [0],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ controlnet_frame_indices: list[int] = [0],
guess_mode: bool = False,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated video.
@@ -758,13 +758,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -780,12 +780,12 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
- conditioning_frames (`List[PipelineImageInput]`, *optional*):
+ conditioning_frames (`list[PipelineImageInput]`, *optional*):
The SparseControlNet input to provide guidance to the `unet` for generation.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generated video. Choose between `torch.Tensor`, `PIL.Image` or `np.array`.
@@ -795,11 +795,11 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
- controlnet_frame_indices (`List[int]`):
+ controlnet_frame_indices (`list[int]`):
The indices where the conditioning frames must be applied for generation. Multiple frames can be
provided to guide the model to generate similar structure outputs, where the `unet` can
"fill-in-the-gaps" for interpolation videos, or a single frame could be provided for general expected
@@ -812,7 +812,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py
index 51a9a31c4259..4e7cd21fc25d 100644
--- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py
+++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
@@ -106,7 +106,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -121,10 +121,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -139,15 +139,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -225,16 +225,14 @@ def __init__(
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetMotionModel],
+ unet: UNet2DConditionModel | UNetMotionModel,
motion_adapter: MotionAdapter,
- scheduler: Union[
- DDIMScheduler,
- PNDMScheduler,
- LMSDiscreteScheduler,
- EulerDiscreteScheduler,
- EulerAncestralDiscreteScheduler,
- DPMSolverMultistepScheduler,
- ],
+ scheduler: DDIMScheduler
+ | PNDMScheduler
+ | LMSDiscreteScheduler
+ | EulerDiscreteScheduler
+ | EulerAncestralDiscreteScheduler
+ | DPMSolverMultistepScheduler,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
):
@@ -262,16 +260,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -279,7 +277,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -378,7 +376,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -636,16 +634,16 @@ def get_timesteps(self, num_inference_steps, timesteps, strength, device):
def prepare_latents(
self,
- video: Optional[torch.Tensor] = None,
+ video: torch.Tensor | None = None,
height: int = 64,
width: int = 64,
num_channels_latents: int = 4,
batch_size: int = 1,
- timestep: Optional[int] = None,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ timestep: int | None = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
decode_chunk_size: int = 16,
add_noise: bool = False,
) -> torch.Tensor:
@@ -746,40 +744,40 @@ def interrupt(self):
@torch.no_grad()
def __call__(
self,
- video: List[List[PipelineImageInput]] = None,
- prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ video: list[list[PipelineImageInput]] = None,
+ prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
enforce_inference_steps: bool = False,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.5,
strength: float = 0.8,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
decode_chunk_size: int = 16,
):
r"""
The call function to the pipeline for generation.
Args:
- video (`List[PipelineImageInput]`):
+ video (`list[PipelineImageInput]`):
The input video to condition the generation on. Must be a list of images/frames of the video.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated video.
@@ -788,11 +786,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality videos at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -801,13 +799,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -823,7 +821,7 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -843,7 +841,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py
index c3ac7df2cc8c..56ed5e23c1db 100644
--- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py
+++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
import torch.nn.functional as F
@@ -124,7 +124,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -139,10 +139,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -157,15 +157,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -229,7 +229,7 @@ class AnimateDiffVideoToVideoControlNetPipeline(
A [`UNet2DConditionModel`] used to create a UNetMotionModel to denoise the encoded video latents.
motion_adapter ([`MotionAdapter`]):
A [`MotionAdapter`] to be used in combination with `unet` to denoise the encoded video latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]` or `Tuple[ControlNetModel]` or `MultiControlNetModel`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]` or `tuple[ControlNetModel]` or `MultiControlNetModel`):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -247,17 +247,15 @@ def __init__(
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetMotionModel],
+ unet: UNet2DConditionModel | UNetMotionModel,
motion_adapter: MotionAdapter,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
- scheduler: Union[
- DDIMScheduler,
- PNDMScheduler,
- LMSDiscreteScheduler,
- EulerDiscreteScheduler,
- EulerAncestralDiscreteScheduler,
- DPMSolverMultistepScheduler,
- ],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
+ scheduler: DDIMScheduler
+ | PNDMScheduler
+ | LMSDiscreteScheduler
+ | EulerDiscreteScheduler
+ | EulerAncestralDiscreteScheduler
+ | DPMSolverMultistepScheduler,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
):
@@ -293,16 +291,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -310,7 +308,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -409,7 +407,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -768,16 +766,16 @@ def get_timesteps(self, num_inference_steps, timesteps, strength, device):
# Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.prepare_latents
def prepare_latents(
self,
- video: Optional[torch.Tensor] = None,
+ video: torch.Tensor | None = None,
height: int = 64,
width: int = 64,
num_channels_latents: int = 4,
batch_size: int = 1,
- timestep: Optional[int] = None,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ timestep: int | None = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
decode_chunk_size: int = 16,
add_noise: bool = False,
) -> torch.Tensor:
@@ -911,45 +909,45 @@ def interrupt(self):
@torch.no_grad()
def __call__(
self,
- video: List[List[PipelineImageInput]] = None,
- prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ video: list[list[PipelineImageInput]] = None,
+ prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
enforce_inference_steps: bool = False,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.5,
strength: float = 0.8,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- conditioning_frames: Optional[List[PipelineImageInput]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ conditioning_frames: list[PipelineImageInput] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
decode_chunk_size: int = 16,
):
r"""
The call function to the pipeline for generation.
Args:
- video (`List[PipelineImageInput]`):
+ video (`list[PipelineImageInput]`):
The input video to condition the generation on. Must be a list of images/frames of the video.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated video.
@@ -958,11 +956,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality videos at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -971,13 +969,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -993,12 +991,12 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
- conditioning_frames (`List[PipelineImageInput]`, *optional*):
+ conditioning_frames (`list[PipelineImageInput]`, *optional*):
The ControlNet input condition to provide guidance to the `unet` for generation. If multiple
ControlNets are specified, images must be passed as a list such that each element of the list can be
correctly batched for input to a single ControlNet.
@@ -1009,16 +1007,16 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
@@ -1028,7 +1026,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/animatediff/pipeline_output.py b/src/diffusers/pipelines/animatediff/pipeline_output.py
index 2417223cf95e..436a20f455bf 100644
--- a/src/diffusers/pipelines/animatediff/pipeline_output.py
+++ b/src/diffusers/pipelines/animatediff/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -14,11 +13,11 @@ class AnimateDiffPipelineOutput(BaseOutput):
Output class for AnimateDiff pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised
PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`
"""
- frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]]
+ frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]
diff --git a/src/diffusers/pipelines/audioldm/pipeline_audioldm.py b/src/diffusers/pipelines/audioldm/pipeline_audioldm.py
index 6a70f00c76c7..357c3582b21c 100644
--- a/src/diffusers/pipelines/audioldm/pipeline_audioldm.py
+++ b/src/diffusers/pipelines/audioldm/pipeline_audioldm.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -88,7 +88,7 @@ def __init__(
self,
vae: AutoencoderKL,
text_encoder: ClapTextModelWithProjection,
- tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast],
+ tokenizer: RobertaTokenizer | RobertaTokenizerFast,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
vocoder: SpeechT5HifiGan,
@@ -112,14 +112,14 @@ def _encode_prompt(
num_waveforms_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device (`torch.device`):
torch device
@@ -127,7 +127,7 @@ def _encode_prompt(
number of waveforms that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the audio generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -189,7 +189,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -361,28 +361,28 @@ def prepare_latents(self, batch_size, num_channels_latents, height, dtype, devic
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- audio_length_in_s: Optional[float] = None,
+ prompt: str | list[str] = None,
+ audio_length_in_s: float | None = None,
num_inference_steps: int = 10,
guidance_scale: float = 2.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_waveforms_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_waveforms_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- output_type: Optional[str] = "np",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ output_type: str | None = "np",
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide audio generation. If not defined, you need to pass `prompt_embeds`.
audio_length_in_s (`int`, *optional*, defaults to 5.12):
The length of the generated audio sample in seconds.
@@ -392,7 +392,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 2.5):
A higher guidance scale value encourages the model to generate audio that is closely linked to the text
`prompt` at the expense of lower sound quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in audio generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_waveforms_per_prompt (`int`, *optional*, defaults to 1):
@@ -400,7 +400,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py
index 878f6f08db42..09aa0ad17003 100644
--- a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py
+++ b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py
@@ -13,7 +13,7 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import torch
import torch.nn as nn
@@ -28,10 +28,7 @@
AttnAddedKVProcessor,
AttnProcessor,
)
-from ...models.embeddings import (
- TimestepEmbedding,
- Timesteps,
-)
+from ...models.embeddings import TimestepEmbedding, Timesteps
from ...models.modeling_utils import ModelMixin
from ...models.resnet import Downsample2D, ResnetBlock2D, Upsample2D
from ...models.transformers.transformer_2d import Transformer2DModel
@@ -75,7 +72,7 @@ class AudioLDM2ProjectionModelOutput(BaseOutput):
"""
hidden_states: torch.Tensor
- attention_mask: Optional[torch.LongTensor] = None
+ attention_mask: torch.LongTensor | None = None
class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin):
@@ -124,10 +121,10 @@ def __init__(
def forward(
self,
- hidden_states: Optional[torch.Tensor] = None,
- hidden_states_1: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.LongTensor] = None,
- attention_mask_1: Optional[torch.LongTensor] = None,
+ hidden_states: torch.Tensor | None = None,
+ hidden_states_1: torch.Tensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ attention_mask_1: torch.LongTensor | None = None,
):
hidden_states = self.projection(hidden_states)
hidden_states, attention_mask = add_special_tokens(
@@ -174,23 +171,23 @@ class AudioLDM2UNet2DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNe
for all models (such as downloading or saving).
Parameters:
- sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+ sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`):
Height and width of input/output sample.
in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
flip_sin_to_cos (`bool`, *optional*, defaults to `False`):
Whether to flip the sin to cos in the time embedding.
freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
The tuple of downsample blocks to use.
mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
Block type for middle of UNet, it can only be `UNetMidBlock2DCrossAttn` for AudioLDM2.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
The tuple of upsample blocks to use.
- only_cross_attention (`bool` or `Tuple[bool]`, *optional*, default to `False`):
+ only_cross_attention (`bool` or `tuple[bool]`, *optional*, default to `False`):
Whether to include self-attention in the basic transformer blocks, see
[`~models.attention.BasicTransformerBlock`].
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
@@ -199,9 +196,9 @@ class AudioLDM2UNet2DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNe
norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
If `None`, normalization and activation layers is skipped in post-processing.
norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
- cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+ cross_attention_dim (`int` or `tuple[int]`, *optional*, defaults to 1280):
The dimension of the cross attention features.
- transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
+ transformer_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 1):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
[`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
[`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
@@ -240,49 +237,44 @@ class conditioning with `class_embed_type` equal to `None`.
@register_to_config
def __init__(
self,
- sample_size: Optional[int] = None,
+ sample_size: int | None = None,
in_channels: int = 4,
out_channels: int = 4,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
),
- mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
- up_block_types: Tuple[str, ...] = (
- "UpBlock2D",
- "CrossAttnUpBlock2D",
- "CrossAttnUpBlock2D",
- "CrossAttnUpBlock2D",
- ),
- only_cross_attention: Union[bool, Tuple[bool]] = False,
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
- layers_per_block: Union[int, Tuple[int]] = 2,
+ mid_block_type: str = "UNetMidBlock2DCrossAttn",
+ up_block_types: tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
+ only_cross_attention: bool | tuple[bool] = False,
+ block_out_channels: tuple[int] = (320, 640, 1280, 1280),
+ layers_per_block: int | tuple[int] = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
act_fn: str = "silu",
- norm_num_groups: Optional[int] = 32,
+ norm_num_groups: int | None = 32,
norm_eps: float = 1e-5,
- cross_attention_dim: Union[int, Tuple[int]] = 1280,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
- attention_head_dim: Union[int, Tuple[int]] = 8,
- num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+ cross_attention_dim: int | tuple[int] = 1280,
+ transformer_layers_per_block: int | tuple[int] = 1,
+ attention_head_dim: int | tuple[int] = 8,
+ num_attention_heads: int | tuple[int] | None = None,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- num_class_embeds: Optional[int] = None,
+ class_embed_type: str | None = None,
+ num_class_embeds: int | None = None,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
time_embedding_type: str = "positional",
- time_embedding_dim: Optional[int] = None,
- time_embedding_act_fn: Optional[str] = None,
- timestep_post_act: Optional[str] = None,
- time_cond_proj_dim: Optional[int] = None,
+ time_embedding_dim: int | None = None,
+ time_embedding_act_fn: str | None = None,
+ timestep_post_act: str | None = None,
+ time_cond_proj_dim: int | None = None,
conv_in_kernel: int = 3,
conv_out_kernel: int = 3,
- projection_class_embeddings_input_dim: Optional[int] = None,
+ projection_class_embeddings_input_dim: int | None = None,
class_embeddings_concat: bool = False,
):
super().__init__()
@@ -606,7 +598,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
# Recursively walk through all the children.
# Any children which exposes the set_attention_slice method
# gets the message
- def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]):
if hasattr(module, "set_attention_slice"):
module.set_attention_slice(slice_size.pop())
@@ -620,17 +612,17 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
return_dict: bool = True,
- encoder_hidden_states_1: Optional[torch.Tensor] = None,
- encoder_attention_mask_1: Optional[torch.Tensor] = None,
- ) -> Union[UNet2DConditionOutput, Tuple]:
+ encoder_hidden_states_1: torch.Tensor | None = None,
+ encoder_attention_mask_1: torch.Tensor | None = None,
+ ) -> UNet2DConditionOutput | tuple:
r"""
The [`AudioLDM2UNet2DConditionModel`] forward method.
@@ -1032,13 +1024,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states_1: Optional[torch.Tensor] = None,
- encoder_attention_mask_1: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states_1: torch.Tensor | None = None,
+ encoder_attention_mask_1: torch.Tensor | None = None,
):
output_states = ()
num_layers = len(self.resnets)
@@ -1194,13 +1186,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states_1: Optional[torch.Tensor] = None,
- encoder_attention_mask_1: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states_1: torch.Tensor | None = None,
+ encoder_attention_mask_1: torch.Tensor | None = None,
) -> torch.Tensor:
hidden_states = self.resnets[0](hidden_states, temb)
num_attention_per_layer = len(self.attentions) // (len(self.resnets) - 1)
@@ -1344,15 +1336,15 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states_1: Optional[torch.Tensor] = None,
- encoder_attention_mask_1: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ upsample_size: int | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states_1: torch.Tensor | None = None,
+ encoder_attention_mask_1: torch.Tensor | None = None,
):
num_layers = len(self.resnets)
num_attention_per_layer = len(self.attentions) // num_layers
diff --git a/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py b/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py
index 452fc3c01b27..b023974a33dd 100644
--- a/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py
+++ b/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -196,11 +196,11 @@ def __init__(
self,
vae: AutoencoderKL,
text_encoder: ClapModel,
- text_encoder_2: Union[T5EncoderModel, VitsModel],
+ text_encoder_2: T5EncoderModel | VitsModel,
projection_model: AudioLDM2ProjectionModel,
language_model: GPT2LMHeadModel,
- tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast],
- tokenizer_2: Union[T5Tokenizer, T5TokenizerFast, VitsTokenizer],
+ tokenizer: RobertaTokenizer | RobertaTokenizerFast,
+ tokenizer_2: T5Tokenizer | T5TokenizerFast | VitsTokenizer,
feature_extractor: ClapFeatureExtractor,
unet: AudioLDM2UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
@@ -251,7 +251,7 @@ def disable_vae_slicing(self):
)
self.vae.disable_slicing()
- def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
+ def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = "cuda"):
r"""
Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
@@ -316,7 +316,7 @@ def generate_language_model(
The sequence used as a prompt for the generation.
max_new_tokens (`int`):
Number of new tokens to generate.
- model_kwargs (`Dict[str, Any]`, *optional*):
+ model_kwargs (`dict[str, Any]`, *optional*):
Ad hoc parametrization of additional model-specific kwargs that will be forwarded to the `forward`
function of the model.
@@ -361,21 +361,21 @@ def encode_prompt(
do_classifier_free_guidance,
transcription=None,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- generated_prompt_embeds: Optional[torch.Tensor] = None,
- negative_generated_prompt_embeds: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.LongTensor] = None,
- negative_attention_mask: Optional[torch.LongTensor] = None,
- max_new_tokens: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ generated_prompt_embeds: torch.Tensor | None = None,
+ negative_generated_prompt_embeds: torch.Tensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ negative_attention_mask: torch.LongTensor | None = None,
+ max_new_tokens: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- transcription (`str` or `List[str]`):
+ transcription (`str` or `list[str]`):
transcription of text to speech
device (`torch.device`):
torch device
@@ -383,7 +383,7 @@ def encode_prompt(
number of waveforms that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the audio generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -502,6 +502,10 @@ def encode_prompt(
text_input_ids,
attention_mask=attention_mask,
)
+ # Extract the pooler output if it's a BaseModelOutputWithPooling (Transformers v5+)
+ # otherwise use it directly (Transformers v4)
+ if hasattr(prompt_embeds, "pooler_output"):
+ prompt_embeds = prompt_embeds.pooler_output
# append the seq-len dim: (bs, hidden_size) -> (bs, seq_len, hidden_size)
prompt_embeds = prompt_embeds[:, None, :]
# make sure that we attend to this single hidden-state
@@ -569,7 +573,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -610,6 +614,10 @@ def encode_prompt(
uncond_input_ids,
attention_mask=negative_attention_mask,
)
+ # Extract the pooler output if it's a BaseModelOutputWithPooling (Transformers v5+)
+ # otherwise use it directly (Transformers v4)
+ if hasattr(negative_prompt_embeds, "pooler_output"):
+ negative_prompt_embeds = negative_prompt_embeds.pooler_output
# append the seq-len dim: (bs, hidden_size) -> (bs, seq_len, hidden_size)
negative_prompt_embeds = negative_prompt_embeds[:, None, :]
# make sure that we attend to this single hidden-state
@@ -862,36 +870,36 @@ def prepare_latents(self, batch_size, num_channels_latents, height, dtype, devic
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- transcription: Union[str, List[str]] = None,
- audio_length_in_s: Optional[float] = None,
+ prompt: str | list[str] = None,
+ transcription: str | list[str] = None,
+ audio_length_in_s: float | None = None,
num_inference_steps: int = 200,
guidance_scale: float = 3.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_waveforms_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_waveforms_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- generated_prompt_embeds: Optional[torch.Tensor] = None,
- negative_generated_prompt_embeds: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.LongTensor] = None,
- negative_attention_mask: Optional[torch.LongTensor] = None,
- max_new_tokens: Optional[int] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ generated_prompt_embeds: torch.Tensor | None = None,
+ negative_generated_prompt_embeds: torch.Tensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ negative_attention_mask: torch.LongTensor | None = None,
+ max_new_tokens: int | None = None,
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- output_type: Optional[str] = "np",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ output_type: str | None = "np",
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide audio generation. If not defined, you need to pass `prompt_embeds`.
- transcription (`str` or `List[str]`, *optional*):\
+ transcription (`str` or `list[str]`, *optional*):\
The transcript for text to speech.
audio_length_in_s (`int`, *optional*, defaults to 10.24):
The length of the generated audio sample in seconds.
@@ -901,7 +909,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 3.5):
A higher guidance scale value encourages the model to generate audio that is closely linked to the text
`prompt` at the expense of lower sound quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in audio generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_waveforms_per_prompt (`int`, *optional*, defaults to 1):
@@ -912,7 +920,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py b/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py
index 1d75e4bef31e..c609bc8b29f8 100644
--- a/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py
+++ b/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import T5Tokenizer, UMT5EncoderModel
@@ -63,10 +63,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -81,15 +81,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -231,25 +231,25 @@ def check_inputs(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] = None,
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
@@ -429,37 +429,35 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
num_inference_steps: int = 50,
- sigmas: List[float] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 3.5,
- num_images_per_prompt: Optional[int] = 1,
- height: Optional[int] = 1024,
- width: Optional[int] = 1024,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ num_images_per_prompt: int | None = 1,
+ height: int | None = 1024,
+ width: int | None = 1024,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 256,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- ) -> Union[ImagePipelineOutput, Tuple]:
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ ) -> ImagePipelineOutput | tuple:
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -470,7 +468,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
guidance_scale (`float`, *optional*, defaults to 5.0):
@@ -481,7 +479,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -514,7 +512,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py
index be1d7ea5a54e..7247ca0d161c 100644
--- a/src/diffusers/pipelines/auto_pipeline.py
+++ b/src/diffusers/pipelines/auto_pipeline.py
@@ -396,11 +396,11 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
@@ -422,7 +422,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not
guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
information.
- device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+ device_map (`str` or `dict[str, int | str | torch.device]`, *optional*):
A map that specifies where each submodule should go. It doesn’t need to be defined for each
parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the
same device.
@@ -687,11 +687,11 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
@@ -713,7 +713,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not
guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
information.
- device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+ device_map (`str` or `dict[str, int | str | torch.device]`, *optional*):
A map that specifies where each submodule should go. It doesn’t need to be defined for each
parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the
same device.
@@ -993,11 +993,11 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
@@ -1019,7 +1019,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not
guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
information.
- device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+ device_map (`str` or `dict[str, int | str | torch.device]`, *optional*):
A map that specifies where each submodule should go. It doesn’t need to be defined for each
parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the
same device.
diff --git a/src/diffusers/pipelines/blip_diffusion/__init__.py b/src/diffusers/pipelines/blip_diffusion/__init__.py
index af6c879d5ce8..c245313e2f8a 100644
--- a/src/diffusers/pipelines/blip_diffusion/__init__.py
+++ b/src/diffusers/pipelines/blip_diffusion/__init__.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL
diff --git a/src/diffusers/pipelines/blip_diffusion/blip_image_processing.py b/src/diffusers/pipelines/blip_diffusion/blip_image_processing.py
index e45f431d0b9d..5d2a0186f041 100644
--- a/src/diffusers/pipelines/blip_diffusion/blip_image_processing.py
+++ b/src/diffusers/pipelines/blip_diffusion/blip_image_processing.py
@@ -14,8 +14,6 @@
# limitations under the License.
"""Image processor class for BLIP."""
-from typing import Dict, List, Optional, Union
-
import numpy as np
import torch
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
@@ -69,11 +67,11 @@ class BlipImageProcessor(BaseImageProcessor):
do_normalize (`bool`, *optional*, defaults to `True`):
Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
method. Can be overridden by the `do_normalize` parameter in the `preprocess` method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
+ image_mean (`float` or `list[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
Mean to use if normalizing the image. This is a float or list of floats the length of the number of
channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method. Can be
overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
+ image_std (`float` or `list[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
Can be overridden by the `image_std` parameter in the `preprocess` method.
@@ -86,13 +84,13 @@ class BlipImageProcessor(BaseImageProcessor):
def __init__(
self,
do_resize: bool = True,
- size: Dict[str, int] = None,
+ size: dict[str, int] = None,
resample: PILImageResampling = PILImageResampling.BICUBIC,
do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
+ rescale_factor: int | float = 1 / 255,
do_normalize: bool = True,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
+ image_mean: float | list[float] | None = None,
+ image_std: float | list[float] | None = None,
do_convert_rgb: bool = True,
do_center_crop: bool = True,
**kwargs,
@@ -116,10 +114,10 @@ def __init__(
def resize(
self,
image: np.ndarray,
- size: Dict[str, int],
+ size: dict[str, int],
resample: PILImageResampling = PILImageResampling.BICUBIC,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ data_format: str | ChannelDimension | None = None,
+ input_data_format: str | ChannelDimension | None = None,
**kwargs,
) -> np.ndarray:
"""
@@ -128,7 +126,7 @@ def resize(
Args:
image (`np.ndarray`):
Image to resize.
- size (`Dict[str, int]`):
+ size (`dict[str, int]`):
Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
`PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BICUBIC`.
@@ -164,19 +162,19 @@ def resize(
def preprocess(
self,
images: ImageInput,
- do_resize: Optional[bool] = None,
- size: Optional[Dict[str, int]] = None,
+ do_resize: bool | None = None,
+ size: dict[str, int] | None = None,
resample: PILImageResampling = None,
- do_rescale: Optional[bool] = None,
- do_center_crop: Optional[bool] = None,
- rescale_factor: Optional[float] = None,
- do_normalize: Optional[bool] = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
+ do_rescale: bool | None = None,
+ do_center_crop: bool | None = None,
+ rescale_factor: float | None = None,
+ do_normalize: bool | None = None,
+ image_mean: float | list[float] | None = None,
+ image_std: float | list[float] | None = None,
+ return_tensors: str | TensorType | None = None,
do_convert_rgb: bool = None,
data_format: ChannelDimension = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ input_data_format: str | ChannelDimension | None = None,
**kwargs,
) -> PIL.Image.Image:
"""
@@ -188,7 +186,7 @@ def preprocess(
passing in images with pixel values between 0 and 1, set `do_rescale=False`.
do_resize (`bool`, *optional*, defaults to `self.do_resize`):
Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
+ size (`dict[str, int]`, *optional*, defaults to `self.size`):
Controls the size of the image after `resize`. The shortest edge of the image is resized to
`size["shortest_edge"]` whilst preserving the aspect ratio. If the longest edge of this resized image
is > `int(size["shortest_edge"] * (1333 / 800))`, then the image is resized again to make the longest
@@ -201,9 +199,9 @@ def preprocess(
Rescale factor to rescale the image by if `do_rescale` is set to `True`.
do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+ image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):
Image mean to normalize the image by if `do_normalize` is set to `True`.
- image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+ image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):
Image standard deviation to normalize the image by if `do_normalize` is set to `True`.
do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
Whether to convert the image to RGB.
diff --git a/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py b/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py
index b061ac2636a5..c434ccdaccca 100644
--- a/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py
+++ b/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py
@@ -11,8 +11,6 @@
# 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.
-from typing import Optional, Tuple, Union
-
import torch
from torch import nn
from transformers import BertTokenizer
@@ -372,11 +370,11 @@ def __init__(self, config: Blip2VisionConfig):
@replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=Blip2VisionConfig)
def forward(
self,
- pixel_values: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPooling]:
+ pixel_values: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | BaseModelOutputWithPooling:
r"""
Returns:
@@ -464,7 +462,7 @@ class PreTrainedModel
def get_extended_attention_mask(
self,
attention_mask: torch.Tensor,
- input_shape: Tuple[int],
+ input_shape: tuple[int],
device: torch.device,
has_query: bool = False,
) -> torch.Tensor:
@@ -474,7 +472,7 @@ def get_extended_attention_mask(
Arguments:
attention_mask (`torch.Tensor`):
Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
- input_shape (`Tuple[int]`):
+ input_shape (`tuple[int]`):
The shape of the input to the model.
device (`torch.device`):
The device of the input to the model.
diff --git a/src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py b/src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py
index 1b0342ce7a56..c5364f8985aa 100644
--- a/src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py
+++ b/src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py
@@ -12,8 +12,6 @@
# 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.
-from typing import Optional, Tuple, Union
-
import torch
from torch import nn
from transformers import CLIPPreTrainedModel
@@ -22,7 +20,7 @@
from transformers.models.clip.modeling_clip import CLIPEncoder
-def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
"""
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
"""
@@ -54,13 +52,13 @@ def forward(
self,
ctx_embeddings: torch.Tensor = None,
ctx_begin_pos: list = None,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPooling]:
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | BaseModelOutputWithPooling:
return self.text_model(
ctx_embeddings=ctx_embeddings,
ctx_begin_pos=ctx_begin_pos,
@@ -86,13 +84,13 @@ def forward(
self,
ctx_embeddings: torch.Tensor,
ctx_begin_pos: list,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPooling]:
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | BaseModelOutputWithPooling:
r"""
Returns:
@@ -184,9 +182,9 @@ def forward(
self,
ctx_embeddings: torch.Tensor,
ctx_begin_pos: list,
- input_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
+ input_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
if ctx_embeddings is None:
ctx_len = 0
diff --git a/src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py b/src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py
index 705d930b59fe..aa3dbdae966b 100644
--- a/src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py
+++ b/src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py
@@ -11,8 +11,6 @@
# 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.
-from typing import List, Optional, Union
-
import PIL.Image
import torch
from transformers import CLIPTokenizer
@@ -116,8 +114,8 @@ def __init__(
qformer: Blip2QFormerModel,
image_processor: BlipImageProcessor,
ctx_begin_pos: int = 2,
- mean: List[float] = None,
- std: List[float] = None,
+ mean: list[float] = None,
+ std: list[float] = None,
):
super().__init__()
@@ -193,33 +191,33 @@ def encode_prompt(self, query_embeds, prompt, device=None):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: List[str],
+ prompt: list[str],
reference_image: PIL.Image.Image,
- source_subject_category: List[str],
- target_subject_category: List[str],
- latents: Optional[torch.Tensor] = None,
+ source_subject_category: list[str],
+ target_subject_category: list[str],
+ latents: torch.Tensor | None = None,
guidance_scale: float = 7.5,
height: int = 512,
width: int = 512,
num_inference_steps: int = 50,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- neg_prompt: Optional[str] = "",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ neg_prompt: str | None = "",
prompt_strength: float = 1.0,
prompt_reps: int = 20,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`List[str]`):
+ prompt (`list[str]`):
The prompt or prompts to guide the image generation.
reference_image (`PIL.Image.Image`):
The reference image to condition the generation on.
- source_subject_category (`List[str]`):
+ source_subject_category (`list[str]`):
The source subject category.
- target_subject_category (`List[str]`):
+ target_subject_category (`list[str]`):
The target subject category.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
@@ -238,7 +236,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
neg_prompt (`str`, *optional*, defaults to ""):
diff --git a/src/diffusers/pipelines/bria/pipeline_bria.py b/src/diffusers/pipelines/bria/pipeline_bria.py
index a22a756005ac..95ae9ce96e7e 100644
--- a/src/diffusers/pipelines/bria/pipeline_bria.py
+++ b/src/diffusers/pipelines/bria/pipeline_bria.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -116,7 +116,7 @@ class BriaPipeline(DiffusionPipeline):
def __init__(
self,
transformer: BriaTransformer2DModel,
- scheduler: Union[FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers],
+ scheduler: FlowMatchEulerDiscreteScheduler | KarrasDiffusionSchedulers,
vae: AutoencoderKL,
text_encoder: T5EncoderModel,
tokenizer: T5TokenizerFast,
@@ -145,20 +145,20 @@ def __init__(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 128,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -166,7 +166,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -320,10 +320,10 @@ def check_inputs(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
tokenizer = self.tokenizer
text_encoder = self.text_encoder
@@ -449,97 +449,109 @@ def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 30,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 128,
- clip_value: Union[None, float] = None,
+ clip_value: None | float = None,
normalize: bool = False,
):
r"""
- Function invoked when calling the pipeline for generation.
-
- Args:
- prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
- instead.
- height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
- The height in pixels of the generated image. This is set to 1024 by default for the best results.
- width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
- The width in pixels of the generated image. This is set to 1024 by default for the best results.
- num_inference_steps (`int`, *optional*, defaults to 50):
- The number of denoising steps. More denoising steps usually lead to a higher quality image at the
- expense of slower inference.
- timesteps (`List[int]`, *optional*):
- Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
- in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
- passed will be used. Must be in descending order.
- guidance_scale (`float`, *optional*, defaults to 5.0):
- Guidance scale as defined in [Classifier-Free Diffusion
- Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
- of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
- `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
- the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation. If not defined, one has to pass
- `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
- less than `1`).
- num_images_per_prompt (`int`, *optional*, defaults to 1):
- The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
- One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
- to make generation deterministic.
- latents (`torch.FloatTensor`, *optional*):
- Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
- generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
- tensor will be generated by sampling using the supplied random `generator`.
- prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
- provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
- argument.
- output_type (`str`, *optional*, defaults to `"pil"`):
- The output format of the generate image. Choose between
- [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
- return_dict (`bool`, *optional*, defaults to `True`):
- Whether or not to return a [`~pipelines.bria.BriaPipelineOutput`] instead of a plain tuple.
- attention_kwargs (`dict`, *optional*):
- A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
- `self.processor` in
- [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- callback_on_step_end (`Callable`, *optional*):
- A function that calls at the end of each denoising steps during the inference. The function is called
- with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
- callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
- `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
- The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
- will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
- `._callback_tensor_inputs` attribute of your pipeline class.
- max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`.
-
- Examples:
-
- Returns:
- [`~pipelines.bria.BriaPipelineOutput`] or `tuple`: [`~pipelines.bria.BriaPipelineOutput`] if `return_dict`
- is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
- images.
+ Function invoked when calling the pipeline for generation.
+
+ Args:
+ prompt (`str` or `list[str]`, *optional*):
+ The prompt or prompts to guide the image generation. If not defined, one has to pass
+ `prompt_embeds`. instead.
+ height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+ The height in pixels of the generated image. This is set to 1024 by default for the best
+ results.
+ width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+ The width in pixels of the generated image. This is set to 1024 by default for the best
+ results.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at
+ the expense of slower inference.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps to use for the denoising process with schedulers which support a `timesteps`
+ argument in their `set_timesteps` method. If not defined, the default behavior when
+ `num_inference_steps` is passed will be used. Must be in descending order.
+ guidance_scale (`float`, *optional*, defaults to 5.0):
+ <<<<<<< HEAD
+ Guidance scale as defined in [Classifier-Free Diffusion
+ Guidance](https://arxiv.org/abs/2207.12598). `guidance_scale` is defined as `w` of equation 2.
+ of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting
+ `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely
+ linked to the text `prompt`, usually at the expense of lower image quality.
+ negative_prompt (`str` or `list[str]`, *optional*):
+ =======
+ Guidance scale as defined in [Classifier-Free Diffusion
+ Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of
+ equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is
+ enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images
+ that are closely linked to the text `prompt`, usually at the expense of lower image quality.
+ negative_prompt (`str` or `list[str]`, *optional*):
+ >>>>>>> main
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if
+ `guidance_scale` is less than `1`).
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+ One or a list of [torch
+ generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+ generation deterministic.
+ latents (`torch.FloatTensor`, *optional*):
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for
+ image generation. Can be used to tweak the same generation with different prompts. If not
+ provided, a latents tensor will be generated by sampling using the supplied random `generator`.
+ prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt`
+ input argument.
+ output_type (`str`, *optional*, defaults to `"pil"`):
+ The output format of the generate image. Choose between
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~pipelines.bria.BriaPipelineOutput`] instead of a plain tuple.
+ attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined
+ under `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ callback_on_step_end (`Callable`, *optional*):
+ A function that calls at the end of each denoising steps during the inference. The function is
+ called with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int,
+ timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as
+ specified by `callback_on_step_end_tensor_inputs`.
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the
+ list will be passed as `callback_kwargs` argument. You will only be able to include variables
+ listed in the `._callback_tensor_inputs` attribute of your pipeline class.
+ max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`.
+
+ Examples:
+
+ Returns:
+ [`~pipelines.bria.BriaPipelineOutput`] or `tuple`: [`~pipelines.bria.BriaPipelineOutput`] if
+ `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list
+ with the generated images.
"""
height = height or self.default_sample_size * self.vae_scale_factor
diff --git a/src/diffusers/pipelines/bria/pipeline_output.py b/src/diffusers/pipelines/bria/pipeline_output.py
index 54eed0623371..4bf2ed949c26 100644
--- a/src/diffusers/pipelines/bria/pipeline_output.py
+++ b/src/diffusers/pipelines/bria/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class BriaPipelineOutput(BaseOutput):
Output class for Bria pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py
index 8fd29756b290..1f178066b17d 100644
--- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py
+++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py
@@ -8,7 +8,7 @@
#
# See the license for further details.
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -94,7 +94,7 @@ class BriaFiboPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
def __init__(
self,
transformer: BriaFiboTransformer2DModel,
- scheduler: Union[FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers],
+ scheduler: FlowMatchEulerDiscreteScheduler | KarrasDiffusionSchedulers,
vae: AutoencoderKLWan,
text_encoder: SmolLM3ForCausalLM,
tokenizer: AutoTokenizer,
@@ -113,11 +113,11 @@ def __init__(
def get_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
max_sequence_length: int = 2048,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -200,19 +200,19 @@ def pad_embedding(prompt_embeds, max_tokens, attention_mask=None):
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
guidance_scale: float = 5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 3000,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -220,7 +220,7 @@ def encode_prompt(
number of images that should be generated per prompt
guidance_scale (`float`):
Guidance scale for classifier free guidance.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -459,23 +459,23 @@ def _prepare_attention_mask(attention_mask):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 30,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 3000,
do_patching=False,
):
@@ -483,7 +483,7 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -493,7 +493,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -503,13 +503,13 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py
index aae8fc7367da..c2327bbce1c7 100644
--- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py
+++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py
@@ -274,8 +274,8 @@ def get_prompt_embeds(
prompt: Union[str, List[str]],
num_images_per_prompt: int = 1,
max_sequence_length: int = 2048,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -359,14 +359,14 @@ def pad_embedding(prompt_embeds, max_tokens, attention_mask=None):
def encode_prompt(
self,
prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
guidance_scale: float = 5,
negative_prompt: Optional[Union[str, List[str]]] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
max_sequence_length: int = 3000,
- lora_scale: Optional[float] = None,
+ lora_scale: bool | None = None,
):
r"""
Args:
@@ -620,19 +620,19 @@ def __call__(
prompt: Union[str, List[str]] = None,
image: Optional[PipelineImageInput] = None,
mask: Optional[PipelineMaskInput] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 30,
timesteps: List[int] = None,
- seed: Optional[int] = None,
+ seed: int | None = None,
guidance_scale: float = 5,
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
@@ -1019,7 +1019,7 @@ def prepare_image_latents(
width: int,
dtype: torch.dtype,
device: torch.device,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
):
image = image.to(device=device, dtype=dtype)
diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_output.py b/src/diffusers/pipelines/bria_fibo/pipeline_output.py
index f459185a2c7c..0c131db29d9f 100644
--- a/src/diffusers/pipelines/bria_fibo/pipeline_output.py
+++ b/src/diffusers/pipelines/bria_fibo/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class BriaFiboPipelineOutput(BaseOutput):
Output class for BriaFibo pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image, np.ndarray]
diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma.py b/src/diffusers/pipelines/chroma/pipeline_chroma.py
index ed6c2c2105b6..bc782107022d 100644
--- a/src/diffusers/pipelines/chroma/pipeline_chroma.py
+++ b/src/diffusers/pipelines/chroma/pipeline_chroma.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -91,10 +91,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -109,15 +109,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -208,11 +208,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -264,24 +264,24 @@ def _get_t5_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
do_classifier_free_guidance: bool = True,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
device: (`torch.device`):
@@ -642,39 +642,39 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 35,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_ip_adapter_image: Optional[PipelineImageInput] = None,
- negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_ip_adapter_image: PipelineImageInput | None = None,
+ negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
not greater than `1`).
@@ -685,7 +685,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -697,7 +697,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -708,13 +708,13 @@ def __call__(
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
negative_ip_adapter_image:
(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
@@ -744,7 +744,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py
index 470c746e4146..e1f6e2f8d8af 100644
--- a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py
+++ b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -88,7 +88,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -103,10 +103,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -121,15 +121,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -222,11 +222,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -291,24 +291,24 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
do_classifier_free_guidance: bool = True,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
device: (`torch.device`):
@@ -701,41 +701,41 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 35,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
strength: float = 0.9,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_ip_adapter_image: Optional[PipelineImageInput] = None,
- negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.tensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_ip_adapter_image: PipelineImageInput | None = None,
+ negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
not greater than `1`).
@@ -746,7 +746,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 35):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -764,7 +764,7 @@ def __call__(
A value of 1, therefore, essentially ignores image.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -775,13 +775,13 @@ def __call__(
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
negative_ip_adapter_image:
(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
@@ -811,7 +811,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py b/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py
index 3ea1ece36c87..52c2f7e51cf2 100644
--- a/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py
+++ b/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py
@@ -94,7 +94,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -109,10 +109,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -127,15 +127,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -241,8 +241,8 @@ def _get_t5_prompt_embeds(
prompt: Union[str, List[str], None] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -294,15 +294,15 @@ def encode_prompt(
self,
prompt: Union[str, List[str]],
negative_prompt: Union[str, List[str], None] = None,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
do_classifier_free_guidance: bool = True,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: bool | None = None,
):
r"""
@@ -771,15 +771,15 @@ def __call__(
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
@@ -787,14 +787,14 @@ def __call__(
negative_ip_adapter_image: Optional[PipelineImageInput] = None,
negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ output_type: str = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
max_sequence_length: int = 256,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
diff --git a/src/diffusers/pipelines/chroma/pipeline_output.py b/src/diffusers/pipelines/chroma/pipeline_output.py
index 951d132dba2e..229b0fe42b90 100644
--- a/src/diffusers/pipelines/chroma/pipeline_output.py
+++ b/src/diffusers/pipelines/chroma/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class ChromaPipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py b/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py
index 79f6580fbed6..1e0cc0ea5c2a 100644
--- a/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py
+++ b/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py
@@ -13,7 +13,7 @@
# limitations under the License.
import html
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL
import regex as re
@@ -113,7 +113,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -185,11 +185,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan_i2v.WanImageToVideoPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -228,7 +228,7 @@ def _get_t5_prompt_embeds(
def encode_image(
self,
image: PipelineImageInput,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
image = self.image_processor(images=image, return_tensors="pt").to(device)
@@ -238,23 +238,23 @@ def encode_image(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -381,11 +381,11 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
latent_width = width // self.vae_scale_factor_spatial
@@ -469,26 +469,24 @@ def attention_kwargs(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 480,
width: int = 832,
num_frames: int = 81,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ image_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
enable_temporal_reasoning: bool = False,
num_temporal_reasoning_steps: int = 0,
@@ -499,10 +497,10 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -523,7 +521,7 @@ def __call__(
usually at the expense of lower image quality.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/chronoedit/pipeline_output.py b/src/diffusers/pipelines/chronoedit/pipeline_output.py
index b1df5b9de35d..6247ce9f3a0c 100644
--- a/src/diffusers/pipelines/chronoedit/pipeline_output.py
+++ b/src/diffusers/pipelines/chronoedit/pipeline_output.py
@@ -11,7 +11,7 @@ class ChronoEditPipelineOutput(BaseOutput):
Output class for ChronoEdit pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py
index 245c794c9c93..b883e10a6732 100644
--- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py
+++ b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py
@@ -15,7 +15,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import T5EncoderModel, T5Tokenizer
@@ -87,10 +87,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -105,15 +105,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -182,7 +182,7 @@ def __init__(
text_encoder: T5EncoderModel,
vae: AutoencoderKLCogVideoX,
transformer: CogVideoXTransformer3DModel,
- scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+ scheduler: CogVideoXDDIMScheduler | CogVideoXDPMScheduler,
):
super().__init__()
@@ -201,11 +201,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -243,23 +243,23 @@ def _get_t5_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -444,7 +444,7 @@ def _prepare_rotary_positional_embeddings(
width: int,
num_frames: int,
device: torch.device,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
@@ -506,38 +506,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- num_frames: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ num_frames: int | None = None,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
guidance_scale: float = 6,
use_dynamic_cfg: bool = False,
num_videos_per_prompt: int = 1,
eta: float = 0.0,
- 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,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 226,
- ) -> Union[CogVideoXPipelineOutput, Tuple]:
+ ) -> CogVideoXPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -553,7 +551,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -565,7 +563,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -594,7 +592,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py
index 456f0bda1644..de5b969a9adc 100644
--- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py
+++ b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py
@@ -15,7 +15,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from PIL import Image
@@ -94,10 +94,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -112,15 +112,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -209,11 +209,11 @@ def __init__(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -252,23 +252,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -360,8 +360,8 @@ def prepare_latents(
# Adapted from https://github.com/aigc-apps/CogVideoX-Fun/blob/2a93e5c14e02b2b5921d533fd59fc8c0ed69fb24/cogvideox/pipeline/pipeline_cogvideox_control.py#L366
def prepare_control_latents(
- self, mask: Optional[torch.Tensor] = None, masked_image: Optional[torch.Tensor] = None
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ self, mask: torch.Tensor | None = None, masked_image: torch.Tensor | None = None
+ ) -> tuple[torch.Tensor, torch.Tensor]:
if mask is not None:
masks = []
for i in range(mask.size(0)):
@@ -490,7 +490,7 @@ def _prepare_rotary_positional_embeddings(
width: int,
num_frames: int,
device: torch.device,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
@@ -552,43 +552,41 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- control_video: Optional[List[Image.Image]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
+ control_video: list[Image.Image] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
guidance_scale: float = 6,
use_dynamic_cfg: bool = False,
num_videos_per_prompt: int = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- control_video_latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ control_video_latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 226,
- ) -> Union[CogVideoXPipelineOutput, Tuple]:
+ ) -> CogVideoXPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- control_video (`List[PIL.Image.Image]`):
+ control_video (`list[PIL.Image.Image]`):
The control video to condition the generation on. Must be a list of images/frames of the video. If not
provided, `control_video_latents` must be provided.
height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):
@@ -598,7 +596,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -610,7 +608,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -642,7 +640,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py
index 321f0f073fe7..9687d63bc7bf 100644
--- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py
+++ b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py
@@ -15,7 +15,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL
import torch
@@ -86,10 +86,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -104,15 +104,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -145,7 +145,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -195,7 +195,7 @@ def __init__(
text_encoder: T5EncoderModel,
vae: AutoencoderKLCogVideoX,
transformer: CogVideoXTransformer3DModel,
- scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+ scheduler: CogVideoXDDIMScheduler | CogVideoXDPMScheduler,
):
super().__init__()
@@ -219,11 +219,11 @@ def __init__(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -262,23 +262,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -349,10 +349,10 @@ def prepare_latents(
num_frames: int = 13,
height: int = 60,
width: int = 90,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
):
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -470,7 +470,7 @@ def check_inputs(
and not isinstance(image, list)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+ "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
f" {type(image)}"
)
@@ -537,7 +537,7 @@ def _prepare_rotary_positional_embeddings(
width: int,
num_frames: int,
device: torch.device,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
@@ -600,40 +600,38 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_frames: int = 49,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
guidance_scale: float = 6,
use_dynamic_cfg: bool = False,
num_videos_per_prompt: int = 1,
eta: float = 0.0,
- 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,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 226,
- ) -> Union[CogVideoXPipelineOutput, Tuple]:
+ ) -> CogVideoXPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -649,7 +647,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -661,7 +659,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -690,7 +688,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py
index e27c572020d6..e3ce8292fad6 100644
--- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py
+++ b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py
@@ -15,7 +15,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from PIL import Image
@@ -95,10 +95,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -113,15 +113,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -154,7 +154,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -204,7 +204,7 @@ def __init__(
text_encoder: T5EncoderModel,
vae: AutoencoderKLCogVideoX,
transformer: CogVideoXTransformer3DModel,
- scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+ scheduler: CogVideoXDDIMScheduler | CogVideoXDPMScheduler,
):
super().__init__()
@@ -225,11 +225,11 @@ def __init__(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -268,23 +268,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -349,16 +349,16 @@ def encode_prompt(
def prepare_latents(
self,
- video: Optional[torch.Tensor] = None,
+ video: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 16,
height: int = 60,
width: int = 90,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
- timestep: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
+ timestep: torch.Tensor | None = None,
):
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -514,7 +514,7 @@ def _prepare_rotary_positional_embeddings(
width: int,
num_frames: int,
device: torch.device,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
@@ -576,41 +576,39 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- video: List[Image.Image] = None,
- prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ video: list[Image.Image] = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
strength: float = 0.8,
guidance_scale: float = 6,
use_dynamic_cfg: bool = False,
num_videos_per_prompt: int = 1,
eta: float = 0.0,
- 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,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 226,
- ) -> Union[CogVideoXPipelineOutput, Tuple]:
+ ) -> CogVideoXPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- video (`List[PIL.Image.Image]`):
+ video (`list[PIL.Image.Image]`):
The input video to condition the generation on. Must be a list of images/frames of the video.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -621,7 +619,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -635,7 +633,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -664,7 +662,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cogvideo/pipeline_output.py b/src/diffusers/pipelines/cogvideo/pipeline_output.py
index 3de030dd6928..56ff50132231 100644
--- a/src/diffusers/pipelines/cogvideo/pipeline_output.py
+++ b/src/diffusers/pipelines/cogvideo/pipeline_output.py
@@ -11,8 +11,8 @@ class CogVideoXPipelineOutput(BaseOutput):
Output class for CogVideo pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
diff --git a/src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py b/src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py
index 46f60d24a467..8880e3a0d1e2 100644
--- a/src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py
+++ b/src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py
@@ -14,7 +14,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import T5EncoderModel, T5Tokenizer
@@ -58,10 +58,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -76,15 +76,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -153,7 +153,7 @@ def __init__(
text_encoder: T5EncoderModel,
vae: AutoencoderKL,
transformer: CogView3PlusTransformer2DModel,
- scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+ scheduler: CogVideoXDDIMScheduler | CogVideoXDPMScheduler,
):
super().__init__()
@@ -167,11 +167,11 @@ def __init__(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds with num_videos_per_prompt->num_images_per_prompt
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -209,23 +209,23 @@ def _get_t5_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 224,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -408,36 +408,34 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
guidance_scale: float = 5.0,
num_images_per_prompt: int = 1,
eta: float = 0.0,
- 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,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
output_type: str = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 224,
- ) -> Union[CogView3PipelineOutput, Tuple]:
+ ) -> CogView3PipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -448,7 +446,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -460,7 +458,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to `1`):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -474,12 +472,12 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
@@ -499,7 +497,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cogview3/pipeline_output.py b/src/diffusers/pipelines/cogview3/pipeline_output.py
index 3891dd51e691..6c89e117b74c 100644
--- a/src/diffusers/pipelines/cogview3/pipeline_output.py
+++ b/src/diffusers/pipelines/cogview3/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class CogView3PipelineOutput(BaseOutput):
Output class for CogView3 pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/cogview4/pipeline_cogview4.py b/src/diffusers/pipelines/cogview4/pipeline_cogview4.py
index 9a2d555538d5..329b76d11e0d 100644
--- a/src/diffusers/pipelines/cogview4/pipeline_cogview4.py
+++ b/src/diffusers/pipelines/cogview4/pipeline_cogview4.py
@@ -14,7 +14,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -69,10 +69,10 @@ def calculate_shift(
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -87,15 +87,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
@@ -177,10 +177,10 @@ def __init__(
def _get_glm_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 1024,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -220,23 +220,23 @@ def _get_glm_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 1024,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -403,37 +403,35 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
num_images_per_prompt: int = 1,
- 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,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 1024,
- ) -> Union[CogView4PipelineOutput, Tuple]:
+ ) -> CogView4PipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -444,11 +442,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -460,7 +458,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to `1`):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -474,12 +472,12 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
@@ -499,7 +497,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py b/src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py
index 2d6785f791db..6282bf4cd7a4 100644
--- a/src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py
+++ b/src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py
@@ -14,7 +14,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -71,10 +71,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.cogview4.pipeline_cogview4.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -89,15 +89,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
@@ -180,10 +180,10 @@ def __init__(
# Copied from diffusers.pipelines.cogview4.pipeline_cogview4.CogView4Pipeline._get_glm_embeds
def _get_glm_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 1024,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -224,23 +224,23 @@ def _get_glm_embeds(
# Copied from diffusers.pipelines.cogview4.pipeline_cogview4.CogView4Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 1024,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -435,38 +435,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
num_images_per_prompt: int = 1,
- 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,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 1024,
- ) -> Union[CogView4PipelineOutput, Tuple]:
+ ) -> CogView4PipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -477,11 +475,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -493,7 +491,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to `1`):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -507,12 +505,12 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
@@ -532,7 +530,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cogview4/pipeline_output.py b/src/diffusers/pipelines/cogview4/pipeline_output.py
index 4efec1310845..997444c6c009 100644
--- a/src/diffusers/pipelines/cogview4/pipeline_output.py
+++ b/src/diffusers/pipelines/cogview4/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class CogView4PipelineOutput(BaseOutput):
Output class for CogView3 pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/consisid/consisid_utils.py b/src/diffusers/pipelines/consisid/consisid_utils.py
index 521d4d787e54..c1646e15efbc 100644
--- a/src/diffusers/pipelines/consisid/consisid_utils.py
+++ b/src/diffusers/pipelines/consisid/consisid_utils.py
@@ -135,7 +135,7 @@ def process_face_embeddings(
is_align_face: Boolean flag indicating whether face alignment should be performed.
Returns:
- Tuple:
+ tuple:
- id_cond: Concatenated tensor of Ante face embedding and CLIP vision embedding
- id_vit_hidden: Hidden state of the CLIP vision model, a list of tensors.
- return_face_features_image_2: Processed face features image after normalization and parsing.
@@ -245,7 +245,7 @@ def process_face_embeddings_infer(
is_align_face: Boolean flag indicating whether face alignment should be performed (default: True).
Returns:
- Tuple:
+ tuple:
- id_cond: Concatenated tensor of Ante face embedding and CLIP vision embedding.
- id_vit_hidden: Hidden state of the CLIP vision model, a list of tensors.
- image: Processed face image after feature extraction and alignment.
diff --git a/src/diffusers/pipelines/consisid/pipeline_consisid.py b/src/diffusers/pipelines/consisid/pipeline_consisid.py
index 3e6c149d7f80..20b779bf5aaa 100644
--- a/src/diffusers/pipelines/consisid/pipeline_consisid.py
+++ b/src/diffusers/pipelines/consisid/pipeline_consisid.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -102,7 +102,7 @@ def draw_kps(image_pil, kps, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255),
Parameters:
- image_pil (PIL.Image): Input image as a PIL object.
- kps (list of tuples): A list of keypoints where each keypoint is a tuple of (x, y) coordinates.
- - color_list (list of tuples, optional): List of colors (in RGB format) for each keypoint. Default is a set of five
+ - color_list (list of tuples, optional): list of colors (in RGB format) for each keypoint. Default is a set of five
colors.
Returns:
@@ -176,10 +176,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -194,15 +194,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -235,7 +235,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -311,11 +311,11 @@ def __init__(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -354,23 +354,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -441,11 +441,11 @@ def prepare_latents(
num_frames: int = 13,
height: int = 60,
width: int = 90,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
- kps_cond: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
+ kps_cond: torch.Tensor | None = None,
):
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -572,7 +572,7 @@ def check_inputs(
and not isinstance(image, list)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+ "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
f" {type(image)}"
)
@@ -623,7 +623,7 @@ def _prepare_rotary_positional_embeddings(
width: int,
num_frames: int,
device: torch.device,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
base_size_width = self.transformer.config.sample_width // self.transformer.config.patch_size
@@ -663,8 +663,8 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 480,
width: int = 720,
num_frames: int = 49,
@@ -673,32 +673,30 @@ def __call__(
use_dynamic_cfg: bool = False,
num_videos_per_prompt: int = 1,
eta: float = 0.0,
- 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,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 226,
- id_vit_hidden: Optional[torch.Tensor] = None,
- id_cond: Optional[torch.Tensor] = None,
- kps_cond: Optional[torch.Tensor] = None,
- ) -> Union[ConsisIDPipelineOutput, Tuple]:
+ id_vit_hidden: torch.Tensor | None = None,
+ id_cond: torch.Tensor | None = None,
+ kps_cond: torch.Tensor | None = None,
+ ) -> ConsisIDPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -727,7 +725,7 @@ def __call__(
more faithful image generation, while later steps reduce it for more diverse and natural results.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -756,22 +754,22 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int`, defaults to `226`):
Maximum sequence length in encoded prompt. Must be consistent with
`self.transformer.config.max_text_seq_length` otherwise may lead to poor results.
- id_vit_hidden (`Optional[torch.Tensor]`, *optional*):
+ id_vit_hidden (`torch.Tensor | None`, *optional*):
The tensor representing the hidden features extracted from the face model, which are used to condition
the local facial extractor. This is crucial for the model to obtain high-frequency information of the
face. If not provided, the local facial extractor will not run normally.
- id_cond (`Optional[torch.Tensor]`, *optional*):
+ id_cond (`torch.Tensor | None`, *optional*):
The tensor representing the hidden features extracted from the clip model, which are used to condition
the local facial extractor. This is crucial for the model to edit facial features If not provided, the
local facial extractor will not run normally.
- kps_cond (`Optional[torch.Tensor]`, *optional*):
+ kps_cond (`torch.Tensor | None`, *optional*):
A tensor that determines whether the global facial extractor use keypoint information for conditioning.
If provided, this tensor controls whether facial keypoints such as eyes, nose, and mouth landmarks are
used during the generation process. This helps ensure the model retains more facial low-frequency
diff --git a/src/diffusers/pipelines/consisid/pipeline_output.py b/src/diffusers/pipelines/consisid/pipeline_output.py
index dd4a63aa50b9..83a5be8d230b 100644
--- a/src/diffusers/pipelines/consisid/pipeline_output.py
+++ b/src/diffusers/pipelines/consisid/pipeline_output.py
@@ -11,8 +11,8 @@ class ConsisIDPipelineOutput(BaseOutput):
Output class for ConsisID pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
diff --git a/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py b/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py
index 1fbdeb1f2741..85e59adc39a4 100644
--- a/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py
+++ b/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, List, Optional, Union
+from typing import Callable
import torch
@@ -173,27 +173,27 @@ def check_inputs(self, num_inference_steps, timesteps, latents, batch_size, img_
def __call__(
self,
batch_size: int = 1,
- class_labels: Optional[Union[torch.Tensor, List[int], int]] = None,
+ class_labels: torch.Tensor | list[int] | int | None = None,
num_inference_steps: int = 1,
- timesteps: List[int] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ timesteps: list[int] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
):
r"""
Args:
batch_size (`int`, *optional*, defaults to 1):
The number of images to generate.
- class_labels (`torch.Tensor` or `List[int]` or `int`, *optional*):
+ class_labels (`torch.Tensor` or `list[int]` or `int`, *optional*):
Optional class labels for conditioning class-conditional consistency models. Not used if the model is
not class-conditional.
num_inference_steps (`int`, *optional*, defaults to 1):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
generator (`torch.Generator`, *optional*):
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet.py
index e2fb32688392..86fa135abff4 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet.py
@@ -14,7 +14,7 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -102,10 +102,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -120,15 +120,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -189,7 +189,7 @@ class StableDiffusionControlNetPipeline(
A `CLIPTokenizer` to tokenize text.
unet ([`UNet2DConditionModel`]):
A `UNet2DConditionModel` to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -215,7 +215,7 @@ def __init__(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
@@ -269,9 +269,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -302,16 +302,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -319,7 +319,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -418,7 +418,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -908,45 +908,43 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -962,18 +960,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -981,7 +979,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -995,7 +993,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1014,16 +1012,16 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
@@ -1033,7 +1031,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py
index e0f1879405aa..4ca92b906842 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py
@@ -12,8 +12,6 @@
# 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.
-from typing import List, Optional, Union
-
import PIL.Image
import torch
from transformers import CLIPTokenizer
@@ -126,8 +124,8 @@ def __init__(
controlnet: ControlNetModel,
image_processor: BlipImageProcessor,
ctx_begin_pos: int = 2,
- mean: List[float] = None,
- std: List[float] = None,
+ mean: list[float] = None,
+ std: list[float] = None,
):
super().__init__()
@@ -241,36 +239,36 @@ def prepare_control_image(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: List[str],
+ prompt: list[str],
reference_image: PIL.Image.Image,
condtioning_image: PIL.Image.Image,
- source_subject_category: List[str],
- target_subject_category: List[str],
- latents: Optional[torch.Tensor] = None,
+ source_subject_category: list[str],
+ target_subject_category: list[str],
+ latents: torch.Tensor | None = None,
guidance_scale: float = 7.5,
height: int = 512,
width: int = 512,
num_inference_steps: int = 50,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- neg_prompt: Optional[str] = "",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ neg_prompt: str | None = "",
prompt_strength: float = 1.0,
prompt_reps: int = 20,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`List[str]`):
+ prompt (`list[str]`):
The prompt or prompts to guide the image generation.
reference_image (`PIL.Image.Image`):
The reference image to condition the generation on.
condtioning_image (`PIL.Image.Image`):
The conditioning canny edge image to condition the generation on.
- source_subject_category (`List[str]`):
+ source_subject_category (`list[str]`):
The source subject category.
- target_subject_category (`List[str]`):
+ target_subject_category (`list[str]`):
The target subject category.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
@@ -291,7 +289,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
neg_prompt (`str`, *optional*, defaults to ""):
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py
index 12cc6f630d80..f0cfabc66f25 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -101,7 +101,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -167,7 +167,7 @@ class StableDiffusionControlNetImg2ImgPipeline(
A `CLIPTokenizer` to tokenize text.
unet ([`UNet2DConditionModel`]):
A `UNet2DConditionModel` to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -193,7 +193,7 @@ def __init__(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
@@ -247,9 +247,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -280,16 +280,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -297,7 +297,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -396,7 +396,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -906,49 +906,47 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.8,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 0.8,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The initial image to be used as the starting point for the image generation process. Can also accept
image latents as `image`, and if passing latents directly they are not encoded again.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -971,7 +969,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -979,7 +977,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -993,7 +991,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1006,16 +1004,16 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
@@ -1025,7 +1023,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py
index 6de8e5747b02..d34278d0086b 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py
@@ -15,7 +15,7 @@
# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -113,7 +113,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -163,7 +163,7 @@ class StableDiffusionControlNetInpaintPipeline(
A `CLIPTokenizer` to tokenize text.
unet ([`UNet2DConditionModel`]):
A `UNet2DConditionModel` to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -196,7 +196,7 @@ def __init__(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
@@ -253,9 +253,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -286,16 +286,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -303,7 +303,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -402,7 +402,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -995,62 +995,60 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 1.0,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 0.5,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 0.5,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`,
- `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`,
+ `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, NumPy array or tensor representing an image batch to be used as the starting point. For both
NumPy array and PyTorch tensor, the expected value range is between `[0, 1]`. If it's a tensor or a
list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a NumPy array or
a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`. It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`,
- `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`,
+ `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, NumPy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a NumPy array or PyTorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for PyTorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for NumPy array, it would be for `(B, H, W, 1)`, `(B, H, W)`, `(H,
W, 1)`, or `(H, W)`.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`,
- `List[List[torch.Tensor]]`, or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`,
+ `list[list[torch.Tensor]]`, or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -1080,7 +1078,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1088,7 +1086,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1102,7 +1100,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1115,16 +1113,16 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 0.5):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 0.5):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
@@ -1134,7 +1132,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py
index fb09d04832f3..942bcb49083e 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -71,7 +71,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -244,13 +244,13 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
- feature_extractor: Optional[CLIPImageProcessor] = None,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ add_watermarker: bool | None = None,
+ feature_extractor: CLIPImageProcessor | None = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
):
super().__init__()
@@ -291,26 +291,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -319,11 +319,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -441,7 +441,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -1166,62 +1166,57 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
- control_image: Union[
- PipelineImageInput,
- List[PipelineImageInput],
- ] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ control_image: PipelineImageInput | list[PipelineImageInput] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.9999,
num_inference_steps: int = 50,
- denoising_start: Optional[float] = None,
- denoising_end: Optional[float] = None,
+ denoising_start: float | None = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
guidance_rescale: float = 0.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
image (`PIL.Image.Image`):
@@ -1275,11 +1270,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -1290,7 +1285,7 @@ def __call__(
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1324,17 +1319,17 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
@@ -1354,7 +1349,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py
index 283c3f92390c..89e8d2e54123 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py
@@ -14,7 +14,7 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -120,10 +120,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -138,15 +138,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -212,7 +212,7 @@ class StableDiffusionXLControlNetPipeline(
A `CLIPTokenizer` to tokenize text.
unet ([`UNet2DConditionModel`]):
A `UNet2DConditionModel` to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -257,10 +257,10 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
):
@@ -299,26 +299,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -327,11 +327,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -449,7 +449,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -998,59 +998,57 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -1068,11 +1066,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1086,10 +1084,10 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 5.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2`
and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders.
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1097,7 +1095,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1118,7 +1116,7 @@ def __call__(
weighting). If not provided, pooled `negative_prompt_embeds` are generated from `negative_prompt` input
argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1131,42 +1129,42 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1179,7 +1177,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py
index 2ea7307fec32..8c39856ec7cb 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py
@@ -14,7 +14,7 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -147,7 +147,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -199,7 +199,7 @@ class StableDiffusionXLControlNetImg2ImgPipeline(
Second Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the unet during the denoising process. If you set multiple ControlNets
as a list, the outputs from each ControlNet are added together to create one combined additional
conditioning.
@@ -248,11 +248,11 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
):
@@ -292,26 +292,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -320,11 +320,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -442,7 +442,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -1077,65 +1077,63 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.8,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 0.8,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The initial image will be used as the starting point for the image generation process. Can also accept
image latents as `image`, if passing latents directly, it will not be encoded again.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If
the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also
be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height
@@ -1165,11 +1163,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1177,7 +1175,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1199,7 +1197,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1214,42 +1212,42 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original unet. If multiple ControlNets are specified in init, you can set the
corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
In this mode, the ControlNet encoder will try best to recognize the content of the input image even if
you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the controlnet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the controlnet stops applying.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1270,7 +1268,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py
index e234015f8616..4b7ca284d636 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -76,7 +76,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -231,15 +231,16 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[
- ControlNetUnionModel, List[ControlNetUnionModel], Tuple[ControlNetUnionModel], MultiControlNetUnionModel
- ],
+ controlnet: ControlNetUnionModel
+ | list[ControlNetUnionModel]
+ | tuple[ControlNetUnionModel]
+ | MultiControlNetUnionModel,
scheduler: KarrasDiffusionSchedulers,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
- feature_extractor: Optional[CLIPImageProcessor] = None,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ add_watermarker: bool | None = None,
+ feature_extractor: CLIPImageProcessor | None = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
):
super().__init__()
@@ -280,26 +281,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -308,11 +309,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -430,7 +431,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -1145,60 +1146,58 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
- control_image: Union[PipelineImageInput, List[PipelineImageInput]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ control_image: PipelineImageInput | list[PipelineImageInput] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.9999,
num_inference_steps: int = 50,
- denoising_start: Optional[float] = None,
- denoising_end: Optional[float] = None,
+ denoising_start: float | None = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- control_mode: Optional[Union[int, List[int], List[List[int]]]] = None,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ control_mode: int | list[int] | list[list[int]] | None = None,
guidance_rescale: float = 0.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
image (`PIL.Image.Image`):
@@ -1209,7 +1208,7 @@ def __call__(
repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted
to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)
instead of 3, so the expected shape would be `(B, H, W, 1)`.
- control_image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*):
+ control_image (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -1259,11 +1258,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -1274,7 +1273,7 @@ def __call__(
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1308,33 +1307,33 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- control_mode (`int` or `List[int]` or `List[List[int]], *optional*):
+ control_mode (`int` or `list[int]` or `list[list[int]], *optional*):
The control condition types for the ControlNet. See the ControlNet's model card forinformation on the
available control modes. If multiple ControlNets are specified in `init`, control_mode should be a list
where each ControlNet should have its corresponding control mode list. Should reflect the order of
conditions in control_image.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
@@ -1354,7 +1353,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py
index 99f2958b320e..87057c2392df 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py
@@ -14,7 +14,7 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -115,10 +115,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -133,15 +133,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -246,12 +246,13 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[
- ControlNetUnionModel, List[ControlNetUnionModel], Tuple[ControlNetUnionModel], MultiControlNetUnionModel
- ],
+ controlnet: ControlNetUnionModel
+ | list[ControlNetUnionModel]
+ | tuple[ControlNetUnionModel]
+ | MultiControlNetUnionModel,
scheduler: KarrasDiffusionSchedulers,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
):
@@ -290,26 +291,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -318,11 +319,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -440,7 +441,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -974,58 +975,56 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- control_image: Union[PipelineImageInput, List[PipelineImageInput]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ control_image: PipelineImageInput | list[PipelineImageInput] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- control_mode: Optional[Union[int, List[int], List[List[int]]]] = None,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ control_mode: int | list[int] | list[list[int]] | None = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders.
- control_image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*):
+ control_image (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -1043,11 +1042,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1061,10 +1060,10 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 5.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2`
and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders.
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1072,7 +1071,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1093,7 +1092,7 @@ def __call__(
weighting). If not provided, pooled `negative_prompt_embeds` are generated from `negative_prompt` input
argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1106,47 +1105,47 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- control_mode (`int` or `List[int]` or `List[List[int]], *optional*):
+ control_mode (`int` or `list[int]` or `list[list[int]], *optional*):
The control condition types for the ControlNet. See the ControlNet's model card forinformation on the
available control modes. If multiple ControlNets are specified in `init`, control_mode should be a list
where each ControlNet should have its corresponding control mode list. Should reflect the order of
conditions in control_image.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1159,7 +1158,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py
index 4d0093132b9c..2fb87e57ea8a 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py
@@ -14,7 +14,7 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -173,7 +173,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -263,13 +263,14 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[
- ControlNetUnionModel, List[ControlNetUnionModel], Tuple[ControlNetUnionModel], MultiControlNetUnionModel
- ],
+ controlnet: ControlNetUnionModel
+ | list[ControlNetUnionModel]
+ | tuple[ControlNetUnionModel]
+ | MultiControlNetUnionModel,
scheduler: KarrasDiffusionSchedulers,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
):
@@ -309,26 +310,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -337,11 +338,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -459,7 +460,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -1066,65 +1067,63 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
- control_image: Union[PipelineImageInput, List[PipelineImageInput]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ control_image: PipelineImageInput | list[PipelineImageInput] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.8,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 0.8,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- control_mode: Optional[Union[int, List[int], List[List[int]]]] = None,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ control_mode: int | list[int] | list[list[int]] | None = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The initial image will be used as the starting point for the image generation process. Can also accept
image latents as `image`, if passing latents directly, it will not be encoded again.
- control_image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*):
+ control_image (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -1154,11 +1153,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1166,7 +1165,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1188,7 +1187,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1203,47 +1202,47 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
In this mode, the ControlNet encoder will try best to recognize the content of the input image even if
you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- control_mode (`int` or `List[int]` or `List[List[int]], *optional*):
+ control_mode (`int` or `list[int]` or `list[list[int]], *optional*):
The control condition types for the ControlNet. See the ControlNet's model card forinformation on the
available control modes. If multiple ControlNets are specified in `init`, control_mode should be a list
where each ControlNet should have its corresponding control mode list. Should reflect the order of
conditions in control_image
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1264,7 +1263,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py
index d4c6f336dfef..a26b9068afd1 100644
--- a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py
+++ b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py
@@ -14,7 +14,6 @@
import warnings
from functools import partial
-from typing import Dict, List, Optional, Union
import jax
import jax.numpy as jnp
@@ -148,9 +147,7 @@ def __init__(
tokenizer: CLIPTokenizer,
unet: FlaxUNet2DConditionModel,
controlnet: FlaxControlNetModel,
- scheduler: Union[
- FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler
- ],
+ scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler,
safety_checker: FlaxStableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
dtype: jnp.dtype = jnp.float32,
@@ -180,7 +177,7 @@ def __init__(
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
- def prepare_text_inputs(self, prompt: Union[str, List[str]]):
+ def prepare_text_inputs(self, prompt: str | list[str]):
if not isinstance(prompt, (str, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
@@ -194,7 +191,7 @@ def prepare_text_inputs(self, prompt: Union[str, List[str]]):
return text_input.input_ids
- def prepare_image_inputs(self, image: Union[Image.Image, List[Image.Image]]):
+ def prepare_image_inputs(self, image: Image.Image | list[Image.Image]):
if not isinstance(image, (Image.Image, list)):
raise ValueError(f"image has to be of type `PIL.Image.Image` or list but is {type(image)}")
@@ -243,12 +240,12 @@ def _generate(
self,
prompt_ids: jnp.ndarray,
image: jnp.ndarray,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
num_inference_steps: int,
guidance_scale: float,
- latents: Optional[jnp.ndarray] = None,
- neg_prompt_ids: Optional[jnp.ndarray] = None,
+ latents: jnp.ndarray | None = None,
+ neg_prompt_ids: jnp.ndarray | None = None,
controlnet_conditioning_scale: float = 1.0,
):
height, width = image.shape[-2:]
@@ -353,13 +350,13 @@ def __call__(
self,
prompt_ids: jnp.ndarray,
image: jnp.ndarray,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
num_inference_steps: int = 50,
- guidance_scale: Union[float, jnp.ndarray] = 7.5,
+ guidance_scale: float | jnp.ndarray = 7.5,
latents: jnp.ndarray = None,
neg_prompt_ids: jnp.ndarray = None,
- controlnet_conditioning_scale: Union[float, jnp.ndarray] = 1.0,
+ controlnet_conditioning_scale: float | jnp.ndarray = 1.0,
return_dict: bool = True,
jit: bool = False,
):
diff --git a/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py b/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py
index 29a7d6147638..8882b561f0a1 100644
--- a/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py
+++ b/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import numpy as np
import torch
@@ -176,10 +176,10 @@ class HunyuanDiTControlNetPipeline(DiffusionPipeline):
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. We use
`sdxl-vae-fp16-fix`.
- text_encoder (Optional[`~transformers.BertModel`, `~transformers.CLIPTextModel`]):
+ text_encoder (`~transformers.BertModel`, `~transformers.CLIPTextModel` | None):
Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
HunyuanDiT uses a fine-tuned [bilingual CLIP].
- tokenizer (Optional[`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer`]):
+ tokenizer (`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer` | None):
A `BertTokenizer` or `CLIPTokenizer` to tokenize text.
transformer ([`HunyuanDiT2DModel`]):
The HunyuanDiT model designed by Tencent Hunyuan.
@@ -189,7 +189,7 @@ class HunyuanDiTControlNetPipeline(DiffusionPipeline):
The tokenizer for the mT5 embedder.
scheduler ([`DDPMScheduler`]):
A scheduler to be used in combination with HunyuanDiT to denoise the encoded image latents.
- controlnet ([`HunyuanDiT2DControlNetModel`] or `List[HunyuanDiT2DControlNetModel]` or [`HunyuanDiT2DControlNetModel`]):
+ controlnet ([`HunyuanDiT2DControlNetModel`] or `list[HunyuanDiT2DControlNetModel]` or [`HunyuanDiT2DControlNetModel`]):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -222,14 +222,12 @@ def __init__(
scheduler: DDPMScheduler,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
- controlnet: Union[
- HunyuanDiT2DControlNetModel,
- List[HunyuanDiT2DControlNetModel],
- Tuple[HunyuanDiT2DControlNetModel],
- HunyuanDiT2DMultiControlNetModel,
- ],
- text_encoder_2: Optional[T5EncoderModel] = None,
- tokenizer_2: Optional[T5Tokenizer] = None,
+ controlnet: HunyuanDiT2DControlNetModel
+ | list[HunyuanDiT2DControlNetModel]
+ | tuple[HunyuanDiT2DControlNetModel]
+ | HunyuanDiT2DMultiControlNetModel,
+ text_encoder_2: T5EncoderModel | None = None,
+ tokenizer_2: T5Tokenizer | None = None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -282,19 +280,19 @@ def encode_prompt(
dtype: torch.dtype = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- max_sequence_length: Optional[int] = None,
+ negative_prompt: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ max_sequence_length: int | None = None,
text_encoder_index: int = 0,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -304,7 +302,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -393,7 +391,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -635,43 +633,41 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 5.0,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 5.0,
control_image: PipelineImageInput = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- prompt_attention_mask_2: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask_2: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_2: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ prompt_attention_mask_2: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask_2: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
guidance_rescale: float = 0.0,
- original_size: Optional[Tuple[int, int]] = (1024, 1024),
- target_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
+ original_size: tuple[int, int] | None = (1024, 1024),
+ target_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
use_resolution_binning: bool = True,
):
r"""
The call function to the pipeline for generation with HunyuanDiT.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`):
The height in pixels of the generated image.
@@ -683,23 +679,23 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,
images must be passed as a list such that each element of the list can be correctly batched for input
to a single ControlNet.
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -707,7 +703,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -735,19 +731,19 @@ def __call__(
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
plain tuple.
- callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+ callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
A callback function or a list of callback functions to be called at the end of each denoising step.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
A list of tensor inputs that should be passed to the callback function. If not defined, all tensor
inputs will be passed.
guidance_rescale (`float`, *optional*, defaults to 0.0):
Rescale the noise_cfg according to `guidance_rescale`. Based on findings of [Common Diffusion Noise
Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). See Section 3.4
- original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
+ original_size (`tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
The original size of the image. Used to calculate the time ids.
- target_size (`Tuple[int, int]`, *optional*):
+ target_size (`tuple[int, int]`, *optional*):
The target size of the image. Used to calculate the time ids.
- crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`):
+ crops_coords_top_left (`tuple[int, int]`, *optional*, defaults to `(0, 0)`):
The top left coordinates of the crop. Used to calculate the time ids.
use_resolution_binning (`bool`, *optional*, defaults to `True`):
Whether to use resolution binning or not. If `True`, the input resolution will be mapped to the closest
diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py
index d721acc77c2a..a787a34bdc01 100644
--- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py
+++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import (
@@ -83,10 +83,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -101,15 +101,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -174,7 +174,7 @@ class StableDiffusion3ControlNetPipeline(
tokenizer_3 (`T5TokenizerFast`):
Tokenizer of class
[T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
- controlnet ([`SD3ControlNetModel`] or `List[SD3ControlNetModel]` or [`SD3MultiControlNetModel`]):
+ controlnet ([`SD3ControlNetModel`] or `list[SD3ControlNetModel]` or [`SD3MultiControlNetModel`]):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -199,11 +199,12 @@ def __init__(
tokenizer_2: CLIPTokenizer,
text_encoder_3: T5EncoderModel,
tokenizer_3: T5TokenizerFast,
- controlnet: Union[
- SD3ControlNetModel, List[SD3ControlNetModel], Tuple[SD3ControlNetModel], SD3MultiControlNetModel
- ],
- image_encoder: Optional[SiglipVisionModel] = None,
- feature_extractor: Optional[SiglipImageProcessor] = None,
+ controlnet: SD3ControlNetModel
+ | list[SD3ControlNetModel]
+ | tuple[SD3ControlNetModel]
+ | SD3MultiControlNetModel,
+ image_encoder: SiglipVisionModel | None = None,
+ feature_extractor: SiglipImageProcessor | None = None,
):
super().__init__()
if isinstance(controlnet, (list, tuple)):
@@ -250,11 +251,11 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -307,10 +308,10 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- clip_skip: Optional[int] = None,
+ device: torch.device | None = None,
+ clip_skip: int | None = None,
clip_model_index: int = 0,
):
device = device or self._execution_device
@@ -363,32 +364,32 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]],
- prompt_3: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ prompt_3: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ clip_skip: int | None = None,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -397,14 +398,14 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
prompt_embeds (`torch.FloatTensor`, *optional*):
@@ -760,9 +761,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_ip_adapter_image_embeds
def prepare_ip_adapter_image_embeds(
self,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
) -> torch.Tensor:
@@ -819,50 +820,50 @@ def enable_sequential_cpu_offload(self, *args, **kwargs):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
control_image: PipelineImageInput = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- controlnet_pooled_projections: Optional[torch.FloatTensor] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ controlnet_pooled_projections: torch.FloatTensor | None = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
will be used instead
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -872,7 +873,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -882,37 +883,37 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,
images must be passed as a list such that each element of the list can be correctly batched for input
to a single ControlNet.
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
controlnet_pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`):
Embeddings projected from the embeddings of controlnet input conditions.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used instead
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used instead
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -954,7 +955,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py
index 2071305cdf10..96f53b16cbe8 100644
--- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py
+++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -106,10 +106,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -124,15 +124,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -197,7 +197,7 @@ class StableDiffusion3ControlNetInpaintingPipeline(
tokenizer_3 (`T5TokenizerFast`):
Tokenizer of class
[T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
- controlnet ([`SD3ControlNetModel`] or `List[SD3ControlNetModel]` or [`SD3MultiControlNetModel`]):
+ controlnet ([`SD3ControlNetModel`] or `list[SD3ControlNetModel]` or [`SD3MultiControlNetModel`]):
Provides additional conditioning to the `transformer` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -222,11 +222,12 @@ def __init__(
tokenizer_2: CLIPTokenizer,
text_encoder_3: T5EncoderModel,
tokenizer_3: T5TokenizerFast,
- controlnet: Union[
- SD3ControlNetModel, List[SD3ControlNetModel], Tuple[SD3ControlNetModel], SD3MultiControlNetModel
- ],
+ controlnet: SD3ControlNetModel
+ | list[SD3ControlNetModel]
+ | tuple[SD3ControlNetModel]
+ | SD3MultiControlNetModel,
image_encoder: SiglipModel = None,
- feature_extractor: Optional[SiglipImageProcessor] = None,
+ feature_extractor: SiglipImageProcessor | None = None,
):
super().__init__()
if isinstance(controlnet, (list, tuple)):
@@ -272,11 +273,11 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -329,10 +330,10 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- clip_skip: Optional[int] = None,
+ device: torch.device | None = None,
+ clip_skip: int | None = None,
clip_model_index: int = 0,
):
device = device or self._execution_device
@@ -385,32 +386,32 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]],
- prompt_3: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ prompt_3: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ clip_skip: int | None = None,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -419,14 +420,14 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
prompt_embeds (`torch.FloatTensor`, *optional*):
@@ -927,9 +928,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_ip_adapter_image_embeds
def prepare_ip_adapter_image_embeds(
self,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
) -> torch.Tensor:
@@ -986,51 +987,51 @@ def enable_sequential_cpu_offload(self, *args, **kwargs):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
control_image: PipelineImageInput = None,
control_mask: PipelineImageInput = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- controlnet_pooled_projections: Optional[torch.FloatTensor] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ controlnet_pooled_projections: torch.FloatTensor | None = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
will be used instead
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -1040,7 +1041,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1050,41 +1051,41 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be inpainted (which parts of the image to
be masked out with `control_mask` and repainted according to `prompt`). For both numpy array and
pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the
expected shape should be `(B, C, H, W)`. If it is a numpy array or a list of arrays, the expected shape
should be `(B, H, W, C)` or `(H, W, C)`.
- control_mask (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`):
+ control_mask (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`. And
for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`.
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
controlnet_pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`):
Embeddings projected from the embeddings of controlnet input conditions.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used instead
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used instead
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -1126,7 +1127,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py
index 3682ddc91156..9c81eb57e6c5 100644
--- a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py
+++ b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -150,7 +150,7 @@ def __init__(
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetControlNetXSModel],
+ unet: UNet2DConditionModel | UNetControlNetXSModel,
controlnet: ControlNetXSAdapter,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
@@ -203,9 +203,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -236,16 +236,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -253,7 +253,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -352,7 +352,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -644,39 +644,37 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
control_guidance_start: float = 0.0,
control_guidance_end: float = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -693,7 +691,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -701,7 +699,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -722,13 +720,13 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
@@ -738,7 +736,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py
index 7bf610f3a0ba..e7a5862b5b7a 100644
--- a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py
+++ b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -182,11 +182,11 @@ def __init__(
text_encoder_2: CLIPTextModelWithProjection,
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetControlNetXSModel],
+ unet: UNet2DConditionModel | UNetControlNetXSModel,
controlnet: ControlNetXSAdapter,
scheduler: KarrasDiffusionSchedulers,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
feature_extractor: CLIPImageProcessor = None,
):
super().__init__()
@@ -223,26 +223,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -251,11 +251,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -373,7 +373,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -718,52 +718,50 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
control_guidance_start: float = 0.0,
control_guidance_end: float = 1.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -784,10 +782,10 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 5.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2`
and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders.
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -795,7 +793,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -823,38 +821,38 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`.
control_guidance_start (`float`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
control_guidance_end (`float`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -867,7 +865,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cosmos/__init__.py b/src/diffusers/pipelines/cosmos/__init__.py
index 944f16553173..5fc66cdf84b6 100644
--- a/src/diffusers/pipelines/cosmos/__init__.py
+++ b/src/diffusers/pipelines/cosmos/__init__.py
@@ -25,6 +25,9 @@
_import_structure["pipeline_cosmos2_5_predict"] = [
"Cosmos2_5_PredictBasePipeline",
]
+ _import_structure["pipeline_cosmos2_5_transfer"] = [
+ "Cosmos2_5_TransferPipeline",
+ ]
_import_structure["pipeline_cosmos2_text2image"] = ["Cosmos2TextToImagePipeline"]
_import_structure["pipeline_cosmos2_video2world"] = ["Cosmos2VideoToWorldPipeline"]
_import_structure["pipeline_cosmos_text2world"] = ["CosmosTextToWorldPipeline"]
@@ -41,6 +44,7 @@
from .pipeline_cosmos2_5_predict import (
Cosmos2_5_PredictBasePipeline,
)
+ from .pipeline_cosmos2_5_transfer import Cosmos2_5_TransferPipeline
from .pipeline_cosmos2_text2image import Cosmos2TextToImagePipeline
from .pipeline_cosmos2_video2world import Cosmos2VideoToWorldPipeline
from .pipeline_cosmos_text2world import CosmosTextToWorldPipeline
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py
index 0f3f62551d35..cdea71a5ab93 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -64,7 +64,7 @@ def __init__(self, *args, **kwargs):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -246,10 +246,10 @@ def __init__(
def _get_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -287,6 +287,9 @@ def _get_prompt_embeds(
truncation=True,
padding="max_length",
)
+ input_ids = (
+ input_ids["input_ids"] if not isinstance(input_ids, list) and "input_ids" in input_ids else input_ids
+ )
input_ids = torch.LongTensor(input_ids)
input_ids_batch.append(input_ids)
@@ -313,23 +316,23 @@ def _get_prompt_embeds(
# Modified from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -398,7 +401,7 @@ def encode_prompt(
# diffusers.pipelines.cosmos.pipeline_cosmos2_video2world.Cosmos2TextToImagePipeline.prepare_latents
def prepare_latents(
self,
- video: Optional[torch.Tensor],
+ video: torch.Tensor | None,
batch_size: int,
num_channels_latents: int = 16,
height: int = 704,
@@ -406,10 +409,10 @@ def prepare_latents(
num_frames_in: int = 93,
num_frames_out: int = 93,
do_classifier_free_guidance: bool = True,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -537,25 +540,23 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput | None = None,
- video: List[PipelineImageInput] | None = None,
- prompt: Union[str, List[str]] | None = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ video: list[PipelineImageInput] | None = None,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 704,
width: int = 1280,
num_frames: int = 93,
num_inference_steps: int = 36,
guidance_scale: float = 7.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
conditional_frame_timestep: float = 0.1,
num_latent_conditional_frames: int = 2,
@@ -575,9 +576,9 @@ def __call__(
Args:
image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, *optional*):
Optional single image for Image2World conditioning. Must be `None` when `video` is provided.
- video (`List[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*):
+ video (`list[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*):
Optional input video for Video2World conditioning. Must be `None` when `image` is provided.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide generation. Required unless `prompt_embeds` is supplied.
height (`int`, defaults to `704`):
The height in pixels of the generated image.
@@ -595,7 +596,7 @@ def __call__(
`guidance_scale > 1`.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py
new file mode 100644
index 000000000000..bbe38c44355e
--- /dev/null
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py
@@ -0,0 +1,1012 @@
+# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable, Dict, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from transformers import AutoTokenizer, Qwen2_5_VLForConditionalGeneration
+
+from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...image_processor import PipelineImageInput
+from ...models import AutoencoderKLWan, CosmosControlNetModel, CosmosTransformer3DModel
+from ...schedulers import UniPCMultistepScheduler
+from ...utils import is_cosmos_guardrail_available, is_torch_xla_available, logging, replace_example_docstring
+from ...utils.torch_utils import randn_tensor
+from ...video_processor import VideoProcessor
+from ..pipeline_utils import DiffusionPipeline
+from .pipeline_output import CosmosPipelineOutput
+
+
+if is_cosmos_guardrail_available():
+ from cosmos_guardrail import CosmosSafetyChecker
+else:
+
+ class CosmosSafetyChecker:
+ def __init__(self, *args, **kwargs):
+ raise ImportError(
+ "`cosmos_guardrail` is not installed. Please install it to use the safety checker for Cosmos: `pip install cosmos_guardrail`."
+ )
+
+
+if is_torch_xla_available():
+ import torch_xla.core.xla_model as xm
+
+ XLA_AVAILABLE = True
+else:
+ XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _maybe_pad_or_trim_video(video: torch.Tensor, num_frames: int):
+ n_pad_frames = num_frames - video.shape[2]
+ if n_pad_frames > 0:
+ last_frame = video[:, :, -1:, :, :]
+ video = torch.cat((video, last_frame.repeat(1, 1, n_pad_frames, 1, 1)), dim=2)
+ elif num_frames < video.shape[2]:
+ video = video[:, :, :num_frames, :, :]
+ return video
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+DEFAULT_NEGATIVE_PROMPT = "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."
+
+EXAMPLE_DOC_STRING = """
+ Examples:
+ ```python
+ >>> import cv2
+ >>> import numpy as np
+ >>> import torch
+ >>> from diffusers import Cosmos2_5_TransferPipeline, AutoModel
+ >>> from diffusers.utils import export_to_video, load_video
+
+ >>> model_id = "nvidia/Cosmos-Transfer2.5-2B"
+ >>> # Load a Transfer2.5 controlnet variant (edge, depth, seg, or blur)
+ >>> controlnet = AutoModel.from_pretrained(model_id, revision="diffusers/controlnet/general/edge")
+ >>> pipe = Cosmos2_5_TransferPipeline.from_pretrained(
+ ... model_id, controlnet=controlnet, revision="diffusers/general", torch_dtype=torch.bfloat16
+ ... )
+ >>> pipe = pipe.to("cuda")
+
+ >>> # Video2World with edge control: Generate video guided by edge maps extracted from input video.
+ >>> prompt = (
+ ... "The video is a demonstration of robotic manipulation, likely in a laboratory or testing environment. It"
+ ... "features two robotic arms interacting with a piece of blue fabric. The setting is a room with a beige"
+ ... "couch in the background, providing a neutral backdrop for the robotic activity. The robotic arms are"
+ ... "positioned on either side of the fabric, which is placed on a yellow cushion. The left robotic arm is"
+ ... "white with a black gripper, while the right arm is black with a more complex, articulated gripper. At the"
+ ... "beginning, the fabric is laid out on the cushion. The left robotic arm approaches the fabric, its gripper"
+ ... "opening and closing as it positions itself. The right arm remains stationary initially, poised to assist."
+ ... "As the video progresses, the left arm grips the fabric, lifting it slightly off the cushion. The right arm"
+ ... "then moves in, its gripper adjusting to grasp the opposite side of the fabric. Both arms work in"
+ ... "coordination, lifting and holding the fabric between them. The fabric is manipulated with precision,"
+ ... "showcasing the dexterity and control of the robotic arms. The camera remains static throughout, focusing"
+ ... "on the interaction between the robotic arms and the fabric, allowing viewers to observe the detailed"
+ ... "movements and coordination involved in the task."
+ ... )
+ >>> negative_prompt = (
+ ... "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, "
+ ... "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, "
+ ... "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky "
+ ... "movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, "
+ ... "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. "
+ ... "Overall, the video is of poor quality."
+ ... )
+ >>> input_video = load_video(
+ ... "https://github.com/nvidia-cosmos/cosmos-transfer2.5/raw/refs/heads/main/assets/robot_example/robot_input.mp4"
+ ... )
+ >>> num_frames = 93
+
+ >>> # Extract edge maps from the input video using Canny edge detection
+ >>> edge_maps = [
+ ... cv2.Canny(cv2.cvtColor(np.array(frame.convert("RGB")), cv2.COLOR_RGB2BGR), 100, 200)
+ ... for frame in input_video[:num_frames]
+ ... ]
+ >>> edge_maps = np.stack(edge_maps)[None] # (T, H, W) -> (1, T, H, W)
+ >>> controls = torch.from_numpy(edge_maps).expand(3, -1, -1, -1) # (1, T, H, W) -> (3, T, H, W)
+ >>> controls = [Image.fromarray(x.numpy()) for x in controls.permute(1, 2, 3, 0)]
+ >>> export_to_video(controls, "edge_controlled_video_edge.mp4", fps=30)
+
+ >>> # Transfer inference with controls.
+ >>> video = pipe(
+ ... controls=controls,
+ ... controls_conditioning_scale=1.0,
+ ... prompt=prompt,
+ ... negative_prompt=negative_prompt,
+ ... num_frames=num_frames,
+ ... ).frames[0]
+ >>> export_to_video(video, "edge_controlled_video.mp4", fps=30)
+ ```
+"""
+
+
+class Cosmos2_5_TransferPipeline(DiffusionPipeline):
+ r"""
+ Pipeline for Cosmos Transfer2.5, supporting auto-regressive inference.
+
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+ implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+ Args:
+ text_encoder ([`Qwen2_5_VLForConditionalGeneration`]):
+ Frozen text-encoder. Cosmos Transfer2.5 uses the [Qwen2.5
+ VL](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) encoder.
+ tokenizer (`AutoTokenizer`):
+ Tokenizer associated with the Qwen2.5 VL encoder.
+ transformer ([`CosmosTransformer3DModel`]):
+ Conditional Transformer to denoise the encoded image latents.
+ scheduler ([`UniPCMultistepScheduler`]):
+ A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+ vae ([`AutoencoderKLWan`]):
+ Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
+ controlnet ([`CosmosControlNetModel`]):
+ ControlNet used to condition generation on control inputs.
+ """
+
+ model_cpu_offload_seq = "text_encoder->transformer->controlnet->vae"
+ _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+ # We mark safety_checker as optional here to get around some test failures, but it is not really optional
+ _optional_components = ["safety_checker"]
+ _exclude_from_cpu_offload = ["safety_checker"]
+
+ def __init__(
+ self,
+ text_encoder: Qwen2_5_VLForConditionalGeneration,
+ tokenizer: AutoTokenizer,
+ transformer: CosmosTransformer3DModel,
+ vae: AutoencoderKLWan,
+ scheduler: UniPCMultistepScheduler,
+ controlnet: CosmosControlNetModel,
+ safety_checker: Optional[CosmosSafetyChecker] = None,
+ ):
+ super().__init__()
+
+ if safety_checker is None:
+ safety_checker = CosmosSafetyChecker()
+
+ self.register_modules(
+ vae=vae,
+ text_encoder=text_encoder,
+ tokenizer=tokenizer,
+ transformer=transformer,
+ controlnet=controlnet,
+ scheduler=scheduler,
+ safety_checker=safety_checker,
+ )
+
+ self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4
+ self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+ self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+
+ latents_mean = (
+ torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).float()
+ if getattr(self.vae.config, "latents_mean", None) is not None
+ else None
+ )
+ latents_std = (
+ torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).float()
+ if getattr(self.vae.config, "latents_std", None) is not None
+ else None
+ )
+ self.latents_mean = latents_mean
+ self.latents_std = latents_std
+
+ if self.latents_mean is None or self.latents_std is None:
+ raise ValueError("VAE configuration must define both `latents_mean` and `latents_std`.")
+
+ def _get_prompt_embeds(
+ self,
+ prompt: Union[str, List[str]] = None,
+ max_sequence_length: int = 512,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+ ):
+ device = device or self._execution_device
+ dtype = dtype or self.text_encoder.dtype
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ input_ids_batch = []
+
+ for sample_idx in range(len(prompt)):
+ conversations = [
+ {
+ "role": "system",
+ "content": [
+ {
+ "type": "text",
+ "text": "You are a helpful assistant who will provide prompts to an image generator.",
+ }
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "text",
+ "text": prompt[sample_idx],
+ }
+ ],
+ },
+ ]
+ input_ids = self.tokenizer.apply_chat_template(
+ conversations,
+ tokenize=True,
+ add_generation_prompt=False,
+ add_vision_id=False,
+ max_length=max_sequence_length,
+ truncation=True,
+ padding="max_length",
+ )
+ input_ids = (
+ input_ids["input_ids"] if not isinstance(input_ids, list) and "input_ids" in input_ids else input_ids
+ )
+ input_ids = torch.LongTensor(input_ids)
+ input_ids_batch.append(input_ids)
+
+ input_ids_batch = torch.stack(input_ids_batch, dim=0)
+
+ outputs = self.text_encoder(
+ input_ids_batch.to(device),
+ output_hidden_states=True,
+ )
+ hidden_states = outputs.hidden_states
+
+ normalized_hidden_states = []
+ for layer_idx in range(1, len(hidden_states)):
+ normalized_state = (hidden_states[layer_idx] - hidden_states[layer_idx].mean(dim=-1, keepdim=True)) / (
+ hidden_states[layer_idx].std(dim=-1, keepdim=True) + 1e-8
+ )
+ normalized_hidden_states.append(normalized_state)
+
+ prompt_embeds = torch.cat(normalized_hidden_states, dim=-1)
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ return prompt_embeds
+
+ # Modified from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt
+ def encode_prompt(
+ self,
+ prompt: Union[str, List[str]],
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ do_classifier_free_guidance: bool = True,
+ num_videos_per_prompt: int = 1,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ max_sequence_length: int = 512,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+ Whether to use classifier free guidance or not.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ device: (`torch.device`, *optional*):
+ torch device
+ dtype: (`torch.dtype`, *optional*):
+ torch dtype
+ """
+ device = device or self._execution_device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ if prompt is not None:
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ prompt_embeds = self._get_prompt_embeds(
+ prompt=prompt, max_sequence_length=max_sequence_length, device=device, dtype=dtype
+ )
+
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ _, seq_len, _ = prompt_embeds.shape
+ prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
+ negative_prompt = negative_prompt or ""
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ negative_prompt_embeds = self._get_prompt_embeds(
+ prompt=negative_prompt, max_sequence_length=max_sequence_length, device=device, dtype=dtype
+ )
+
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ _, seq_len, _ = negative_prompt_embeds.shape
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+ return prompt_embeds, negative_prompt_embeds
+
+ # Modified from diffusers.pipelines.cosmos.pipeline_cosmos2_video2world.Cosmos2VideoToWorldPipeline.prepare_latents and
+ # diffusers.pipelines.cosmos.pipeline_cosmos2_video2world.Cosmos2TextToImagePipeline.prepare_latents
+ def prepare_latents(
+ self,
+ video: Optional[torch.Tensor],
+ batch_size: int,
+ num_channels_latents: int = 16,
+ height: int = 704,
+ width: int = 1280,
+ num_frames_in: int = 93,
+ num_frames_out: int = 93,
+ do_classifier_free_guidance: bool = True,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ num_cond_latent_frames: int = 0,
+ ) -> torch.Tensor:
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ B = batch_size
+ C = num_channels_latents
+ T = (num_frames_out - 1) // self.vae_scale_factor_temporal + 1
+ H = height // self.vae_scale_factor_spatial
+ W = width // self.vae_scale_factor_spatial
+ shape = (B, C, T, H, W)
+
+ if latents is not None:
+ if latents.shape[1:] != shape[1:]:
+ raise ValueError(f"Unexpected `latents` shape, got {latents.shape}, expected {shape}.")
+ latents = latents.to(device=device, dtype=dtype)
+ else:
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+ if num_frames_in == 0:
+ cond_mask = torch.zeros((B, 1, T, H, W), dtype=latents.dtype, device=latents.device)
+ cond_indicator = torch.zeros((B, 1, T, 1, 1), dtype=latents.dtype, device=latents.device)
+
+ cond_latents = torch.zeros_like(latents)
+
+ return (
+ latents,
+ cond_latents,
+ cond_mask,
+ cond_indicator,
+ )
+ else:
+ if video is None:
+ raise ValueError("`video` must be provided when `num_frames_in` is greater than 0.")
+ video = video.to(device=device, dtype=self.vae.dtype)
+ if isinstance(generator, list):
+ cond_latents = [
+ retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator=generator[i])
+ for i in range(batch_size)
+ ]
+ else:
+ cond_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video]
+
+ cond_latents = torch.cat(cond_latents, dim=0).to(dtype)
+
+ latents_mean = self.latents_mean.to(device=device, dtype=dtype)
+ latents_std = self.latents_std.to(device=device, dtype=dtype)
+ cond_latents = (cond_latents - latents_mean) / latents_std
+
+ padding_shape = (B, 1, T, H, W)
+ ones_padding = latents.new_ones(padding_shape)
+ zeros_padding = latents.new_zeros(padding_shape)
+
+ cond_indicator = latents.new_zeros(B, 1, latents.size(2), 1, 1)
+ cond_indicator[:, :, 0:num_cond_latent_frames, :, :] = 1.0
+ cond_mask = cond_indicator * ones_padding + (1 - cond_indicator) * zeros_padding
+
+ return (
+ latents,
+ cond_latents,
+ cond_mask,
+ cond_indicator,
+ )
+
+ # Modified from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.check_inputs
+ def check_inputs(
+ self,
+ prompt,
+ height,
+ width,
+ prompt_embeds=None,
+ callback_on_step_end_tensor_inputs=None,
+ num_ar_conditional_frames=None,
+ num_ar_latent_conditional_frames=None,
+ num_frames_per_chunk=None,
+ num_frames=None,
+ conditional_frame_timestep=0.1,
+ ):
+ if width <= 0 or height <= 0 or height % 16 != 0 or width % 16 != 0:
+ raise ValueError(
+ f"`height` and `width` have to be divisible by 16 (& positive) but are {height} and {width}."
+ )
+
+ if num_frames is not None and num_frames <= 0:
+ raise ValueError(f"`num_frames` has to be a positive integer when provided but is {num_frames}.")
+
+ if conditional_frame_timestep < 0 or conditional_frame_timestep > 1:
+ raise ValueError(
+ "`conditional_frame_timestep` has to be a float in the [0, 1] interval but is "
+ f"{conditional_frame_timestep}."
+ )
+
+ if callback_on_step_end_tensor_inputs is not None and not all(
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+ ):
+ raise ValueError(
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if num_ar_latent_conditional_frames is not None and num_ar_conditional_frames is not None:
+ raise ValueError(
+ "Provide only one of `num_ar_conditional_frames` or `num_ar_latent_conditional_frames`, not both."
+ )
+ if num_ar_latent_conditional_frames is None and num_ar_conditional_frames is None:
+ raise ValueError("Provide either `num_ar_conditional_frames` or `num_ar_latent_conditional_frames`.")
+ if num_ar_latent_conditional_frames is not None and num_ar_latent_conditional_frames < 0:
+ raise ValueError("`num_ar_latent_conditional_frames` must be >= 0.")
+ if num_ar_conditional_frames is not None and num_ar_conditional_frames < 0:
+ raise ValueError("`num_ar_conditional_frames` must be >= 0.")
+
+ if num_ar_latent_conditional_frames is not None:
+ num_ar_conditional_frames = max(
+ 0, (num_ar_latent_conditional_frames - 1) * self.vae_scale_factor_temporal + 1
+ )
+
+ min_chunk_len = self.vae_scale_factor_temporal + 1
+ if num_frames_per_chunk < min_chunk_len:
+ logger.warning(f"{num_frames_per_chunk=} must be larger than {min_chunk_len=}, setting to min_chunk_len")
+ num_frames_per_chunk = min_chunk_len
+
+ max_frames_by_rope = None
+ if getattr(self.transformer.config, "max_size", None) is not None:
+ max_frames_by_rope = max(
+ size // patch
+ for size, patch in zip(self.transformer.config.max_size, self.transformer.config.patch_size)
+ )
+ if num_frames_per_chunk > max_frames_by_rope:
+ raise ValueError(
+ f"{num_frames_per_chunk=} is too large for RoPE setting ({max_frames_by_rope=}). "
+ "Please reduce `num_frames_per_chunk`."
+ )
+
+ if num_ar_conditional_frames >= num_frames_per_chunk:
+ raise ValueError(
+ f"{num_ar_conditional_frames=} must be smaller than {num_frames_per_chunk=} for chunked generation."
+ )
+
+ return num_frames_per_chunk
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ @property
+ def do_classifier_free_guidance(self):
+ return self._guidance_scale > 1.0
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def current_timestep(self):
+ return self._current_timestep
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ controls: PipelineImageInput | List[PipelineImageInput],
+ controls_conditioning_scale: Union[float, List[float]] = 1.0,
+ prompt: Union[str, List[str]] | None = None,
+ negative_prompt: Union[str, List[str]] = DEFAULT_NEGATIVE_PROMPT,
+ height: int = 704,
+ width: Optional[int] = None,
+ num_frames: Optional[int] = None,
+ num_frames_per_chunk: int = 93,
+ num_inference_steps: int = 36,
+ guidance_scale: float = 3.0,
+ num_videos_per_prompt: int = 1,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ callback_on_step_end: Optional[
+ Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+ ] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ max_sequence_length: int = 512,
+ conditional_frame_timestep: float = 0.1,
+ num_ar_conditional_frames: Optional[int] = 1,
+ num_ar_latent_conditional_frames: Optional[int] = None,
+ ):
+ r"""
+ `controls` drive the conditioning through ControlNet. Controls are assumed to be pre-processed, e.g. edge maps
+ are pre-computed.
+
+ Setting `num_frames` will restrict the total number of frames output, if not provided or assigned to None
+ (default) then the number of output frames will match the input `controls`.
+
+ Auto-regressive inference is supported and thus a sliding window of `num_frames_per_chunk` frames are used per
+ denoising loop. In addition, when auto-regressive inference is performed, the previous
+ `num_ar_latent_conditional_frames` or `num_ar_conditional_frames` are used to condition the following denoising
+ inference loops.
+
+ Args:
+ controls (`PipelineImageInput`, `List[PipelineImageInput]`):
+ Control image or video input used by the ControlNet.
+ controls_conditioning_scale (`float` or `List[float]`, *optional*, defaults to `1.0`):
+ The scale factor(s) for the ControlNet outputs. A single float is broadcast to all control blocks.
+ prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts to guide generation. Required unless `prompt_embeds` is supplied.
+ height (`int`, defaults to `704`):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image. If not provided, this will be determined based on the
+ aspect ratio of the input and the provided height.
+ num_frames (`int`, *optional*):
+ Number of output frames. Defaults to `None` to output the same number of frames as the input
+ `controls`.
+ num_inference_steps (`int`, defaults to `36`):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+ expense of slower inference.
+ guidance_scale (`float`, defaults to `3.0`):
+ Guidance scale as defined in [Classifier-Free Diffusion
+ Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
+ of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
+ `guidance_scale > 1`.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+ generation deterministic.
+ latents (`torch.Tensor`, *optional*):
+ Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs. Can be used to
+ tweak the same generation with different prompts. If not provided, a latents tensor is generated by
+ sampling using the supplied random `generator`.
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
+ provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
+ output_type (`str`, *optional*, defaults to `"pil"`):
+ The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`CosmosPipelineOutput`] instead of a plain tuple.
+ callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+ A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
+ each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
+ DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
+ list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+ `._callback_tensor_inputs` attribute of your pipeline class.
+ max_sequence_length (`int`, defaults to `512`):
+ The maximum number of tokens in the prompt. If the prompt exceeds this length, it will be truncated. If
+ the prompt is shorter than this length, it will be padded.
+ num_ar_conditional_frames (`int`, *optional*, defaults to `1`):
+ Number of frames to condition on subsequent inference loops in auto-regressive inference, i.e. for the
+ second chunk and onwards. Only used if `num_ar_latent_conditional_frames` is `None`.
+
+ This is only used when auto-regressive inference is performed, i.e. when the number of frames in
+ controls is > num_frames_per_chunk
+ num_ar_latent_conditional_frames (`int`, *optional*):
+ Number of latent frames to condition on subsequent inference loops in auto-regressive inference, i.e.
+ for the second chunk and onwards. Only used if `num_ar_conditional_frames` is `None`.
+
+ This is only used when auto-regressive inference is performed, i.e. when the number of frames in
+ controls is > num_frames_per_chunk
+ Examples:
+
+ Returns:
+ [`~CosmosPipelineOutput`] or `tuple`:
+ If `return_dict` is `True`, [`CosmosPipelineOutput`] is returned, otherwise a `tuple` is returned where
+ the first element is a list with the generated images and the second element is a list of `bool`s
+ indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
+ """
+ if self.safety_checker is None:
+ raise ValueError(
+ f"You have disabled the safety checker for {self.__class__}. This is in violation of the "
+ "[NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). "
+ f"Please ensure that you are compliant with the license agreement."
+ )
+
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+ if width is None:
+ frame = controls[0] if isinstance(controls, list) else controls
+ if isinstance(frame, list):
+ frame = frame[0]
+ if isinstance(frame, (torch.Tensor, np.ndarray)):
+ if frame.ndim == 5:
+ frame = frame[0, 0]
+ elif frame.ndim == 4:
+ frame = frame[0]
+
+ if isinstance(frame, PIL.Image.Image):
+ width = int((height + 16) * (frame.width / frame.height))
+ else:
+ if frame.ndim != 3:
+ raise ValueError("`controls` must contain 3D frames in CHW format.")
+ width = int((height + 16) * (frame.shape[2] / frame.shape[1])) # NOTE: assuming C H W
+
+ num_frames_per_chunk = self.check_inputs(
+ prompt,
+ height,
+ width,
+ prompt_embeds,
+ callback_on_step_end_tensor_inputs,
+ num_ar_conditional_frames,
+ num_ar_latent_conditional_frames,
+ num_frames_per_chunk,
+ num_frames,
+ conditional_frame_timestep,
+ )
+
+ if num_ar_latent_conditional_frames is not None:
+ num_cond_latent_frames = num_ar_latent_conditional_frames
+ num_ar_conditional_frames = max(0, (num_cond_latent_frames - 1) * self.vae_scale_factor_temporal + 1)
+ else:
+ num_cond_latent_frames = max(0, (num_ar_conditional_frames - 1) // self.vae_scale_factor_temporal + 1)
+
+ self._guidance_scale = guidance_scale
+ self._current_timestep = None
+ self._interrupt = False
+
+ device = self._execution_device
+
+ if self.safety_checker is not None:
+ self.safety_checker.to(device)
+ if prompt is not None:
+ prompt_list = [prompt] if isinstance(prompt, str) else prompt
+ for p in prompt_list:
+ if not self.safety_checker.check_text_safety(p):
+ raise ValueError(
+ f"Cosmos Guardrail detected unsafe text in the prompt: {p}. Please ensure that the "
+ f"prompt abides by the NVIDIA Open Model License Agreement."
+ )
+
+ # 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]
+
+ # Encode input prompt
+ (
+ prompt_embeds,
+ negative_prompt_embeds,
+ ) = self.encode_prompt(
+ prompt=prompt,
+ negative_prompt=negative_prompt,
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
+ num_videos_per_prompt=num_videos_per_prompt,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ device=device,
+ max_sequence_length=max_sequence_length,
+ )
+
+ vae_dtype = self.vae.dtype
+ transformer_dtype = self.transformer.dtype
+
+ if getattr(self.transformer.config, "img_context_dim_in", None):
+ img_context = torch.zeros(
+ batch_size,
+ self.transformer.config.img_context_num_tokens,
+ self.transformer.config.img_context_dim_in,
+ device=prompt_embeds.device,
+ dtype=transformer_dtype,
+ )
+
+ if num_videos_per_prompt > 1:
+ img_context = img_context.repeat_interleave(num_videos_per_prompt, dim=0)
+
+ encoder_hidden_states = (prompt_embeds, img_context)
+ neg_encoder_hidden_states = (negative_prompt_embeds, img_context)
+ else:
+ encoder_hidden_states = prompt_embeds
+ neg_encoder_hidden_states = negative_prompt_embeds
+
+ control_video = self.video_processor.preprocess_video(controls, height, width)
+ if control_video.shape[0] != batch_size:
+ if control_video.shape[0] == 1:
+ control_video = control_video.repeat(batch_size, 1, 1, 1, 1)
+ else:
+ raise ValueError(
+ f"Expected controls batch size {batch_size} to match prompt batch size, but got {control_video.shape[0]}."
+ )
+
+ num_frames_out = control_video.shape[2]
+ if num_frames is not None:
+ num_frames_out = min(num_frames_out, num_frames)
+
+ control_video = _maybe_pad_or_trim_video(control_video, num_frames_out)
+
+ # chunk information
+ num_latent_frames_per_chunk = (num_frames_per_chunk - 1) // self.vae_scale_factor_temporal + 1
+ chunk_stride = num_frames_per_chunk - num_ar_conditional_frames
+ chunk_idxs = [
+ (start_idx, min(start_idx + num_frames_per_chunk, num_frames_out))
+ for start_idx in range(0, num_frames_out - num_ar_conditional_frames, chunk_stride)
+ ]
+
+ video_chunks = []
+ latents_mean = self.latents_mean.to(dtype=vae_dtype, device=device)
+ latents_std = self.latents_std.to(dtype=vae_dtype, device=device)
+
+ def decode_latents(latents):
+ latents = latents * latents_std + latents_mean
+ video = self.vae.decode(latents.to(dtype=self.vae.dtype, device=device), return_dict=False)[0]
+ return video
+
+ latents_arg = latents
+ initial_num_cond_latent_frames = 0
+ latent_chunks = []
+ num_chunks = len(chunk_idxs)
+ total_steps = num_inference_steps * num_chunks
+ with self.progress_bar(total=total_steps) as progress_bar:
+ for chunk_idx, (start_idx, end_idx) in enumerate(chunk_idxs):
+ if chunk_idx == 0:
+ prev_output = torch.zeros((batch_size, num_frames_per_chunk, 3, height, width), dtype=vae_dtype)
+ prev_output = self.video_processor.preprocess_video(prev_output, height, width)
+ else:
+ prev_output = video_chunks[-1].clone()
+ if num_ar_conditional_frames > 0:
+ prev_output[:, :, :num_ar_conditional_frames] = prev_output[:, :, -num_ar_conditional_frames:]
+ prev_output[:, :, num_ar_conditional_frames:] = -1 # -1 == 0 in processed video space
+ else:
+ prev_output.fill_(-1)
+
+ chunk_video = prev_output.to(device=device, dtype=vae_dtype)
+ chunk_video = _maybe_pad_or_trim_video(chunk_video, num_frames_per_chunk)
+ latents, cond_latent, cond_mask, cond_indicator = self.prepare_latents(
+ video=chunk_video,
+ batch_size=batch_size * num_videos_per_prompt,
+ num_channels_latents=self.transformer.config.in_channels - 1,
+ height=height,
+ width=width,
+ num_frames_in=chunk_video.shape[2],
+ num_frames_out=num_frames_per_chunk,
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
+ dtype=torch.float32,
+ device=device,
+ generator=generator,
+ num_cond_latent_frames=initial_num_cond_latent_frames
+ if chunk_idx == 0
+ else num_cond_latent_frames,
+ latents=latents_arg,
+ )
+ cond_mask = cond_mask.to(transformer_dtype)
+ cond_timestep = torch.ones_like(cond_indicator) * conditional_frame_timestep
+ padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
+
+ chunk_control_video = control_video[:, :, start_idx:end_idx, ...].to(
+ device=device, dtype=self.vae.dtype
+ )
+ chunk_control_video = _maybe_pad_or_trim_video(chunk_control_video, num_frames_per_chunk)
+ if isinstance(generator, list):
+ controls_latents = [
+ retrieve_latents(self.vae.encode(chunk_control_video[i].unsqueeze(0)), generator=generator[i])
+ for i in range(chunk_control_video.shape[0])
+ ]
+ else:
+ controls_latents = [
+ retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator=generator)
+ for vid in chunk_control_video
+ ]
+ controls_latents = torch.cat(controls_latents, dim=0).to(transformer_dtype)
+
+ controls_latents = (controls_latents - latents_mean) / latents_std
+
+ # Denoising loop
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
+ timesteps = self.scheduler.timesteps
+ self._num_timesteps = len(timesteps)
+
+ gt_velocity = (latents - cond_latent) * cond_mask
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ self._current_timestep = t.cpu().item()
+
+ # NOTE: assumes sigma(t) \in [0, 1]
+ sigma_t = (
+ torch.tensor(self.scheduler.sigmas[i].item())
+ .unsqueeze(0)
+ .to(device=device, dtype=transformer_dtype)
+ )
+
+ in_latents = cond_mask * cond_latent + (1 - cond_mask) * latents
+ in_latents = in_latents.to(transformer_dtype)
+ in_timestep = cond_indicator * cond_timestep + (1 - cond_indicator) * sigma_t
+ control_output = self.controlnet(
+ controls_latents=controls_latents,
+ latents=in_latents,
+ timestep=in_timestep,
+ encoder_hidden_states=encoder_hidden_states,
+ condition_mask=cond_mask,
+ conditioning_scale=controls_conditioning_scale,
+ padding_mask=padding_mask,
+ return_dict=False,
+ )
+ control_blocks = control_output[0]
+
+ noise_pred = self.transformer(
+ hidden_states=in_latents,
+ timestep=in_timestep,
+ encoder_hidden_states=encoder_hidden_states,
+ block_controlnet_hidden_states=control_blocks,
+ condition_mask=cond_mask,
+ padding_mask=padding_mask,
+ return_dict=False,
+ )[0]
+ noise_pred = gt_velocity + noise_pred * (1 - cond_mask)
+
+ if self.do_classifier_free_guidance:
+ control_output = self.controlnet(
+ controls_latents=controls_latents,
+ latents=in_latents,
+ timestep=in_timestep,
+ encoder_hidden_states=neg_encoder_hidden_states, # NOTE: negative prompt
+ condition_mask=cond_mask,
+ conditioning_scale=controls_conditioning_scale,
+ padding_mask=padding_mask,
+ return_dict=False,
+ )
+ control_blocks = control_output[0]
+
+ noise_pred_neg = self.transformer(
+ hidden_states=in_latents,
+ timestep=in_timestep,
+ encoder_hidden_states=neg_encoder_hidden_states, # NOTE: negative prompt
+ block_controlnet_hidden_states=control_blocks,
+ condition_mask=cond_mask,
+ padding_mask=padding_mask,
+ return_dict=False,
+ )[0]
+ # NOTE: replace velocity (noise_pred_neg) with gt_velocity for conditioning inputs only
+ noise_pred_neg = gt_velocity + noise_pred_neg * (1 - cond_mask)
+ noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_neg)
+
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+ # call the callback, if provided
+ 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)
+
+ if i == total_steps - 1 or ((i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+
+ if XLA_AVAILABLE:
+ xm.mark_step()
+
+ video_chunks.append(decode_latents(latents).detach().cpu())
+ latent_chunks.append(latents.detach().cpu())
+
+ self._current_timestep = None
+
+ if not output_type == "latent":
+ video_chunks = [
+ chunk[:, :, num_ar_conditional_frames:, ...] if chunk_idx != 0 else chunk
+ for chunk_idx, chunk in enumerate(video_chunks)
+ ]
+ video = torch.cat(video_chunks, dim=2)
+ video = video[:, :, :num_frames_out, ...]
+
+ assert self.safety_checker is not None
+ self.safety_checker.to(device)
+ video = self.video_processor.postprocess_video(video, output_type="np")
+ video = (video * 255).astype(np.uint8)
+ video_batch = []
+ for vid in video:
+ vid = self.safety_checker.check_video_safety(vid)
+ if vid is None:
+ video_batch.append(np.zeros_like(video[0]))
+ else:
+ video_batch.append(vid)
+ video = np.stack(video_batch).astype(np.float32) / 255.0 * 2 - 1
+ video = torch.from_numpy(video).permute(0, 4, 1, 2, 3)
+ video = self.video_processor.postprocess_video(video, output_type=output_type)
+ else:
+ latent_T = (num_frames_out - 1) // self.vae_scale_factor_temporal + 1
+ latent_chunks = [
+ chunk[:, :, num_cond_latent_frames:, ...] if chunk_idx != 0 else chunk
+ for chunk_idx, chunk in enumerate(latent_chunks)
+ ]
+ video = torch.cat(latent_chunks, dim=2)
+ video = video[:, :, :latent_T, ...]
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ return (video,)
+
+ return CosmosPipelineOutput(frames=video)
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py
index d5ee60962285..f24e19eea0d4 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -83,10 +83,10 @@ def __init__(self, *args, **kwargs):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -101,15 +101,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -210,10 +210,10 @@ def __init__(
# Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -253,23 +253,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt with num_videos_per_prompt->num_images_per_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -341,10 +341,10 @@ def prepare_latents(
height: int = 768,
width: int = 1360,
num_frames: int = 1,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype) * self.scheduler.config.sigma_max
@@ -418,30 +418,28 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 768,
width: int = 1360,
num_inference_steps: int = 35,
guidance_scale: float = 7.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, defaults to `768`):
@@ -458,7 +456,7 @@ def __call__(
`guidance_scale > 1`.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -480,7 +478,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py
index b2c44d0d9972..bdb13af06637 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -88,10 +88,10 @@ def __init__(self, *args, **kwargs):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -106,15 +106,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -147,7 +147,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -229,10 +229,10 @@ def __init__(
# Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -272,23 +272,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -362,10 +362,10 @@ def prepare_latents(
width: int = 1280,
num_frames: int = 93,
do_classifier_free_guidance: bool = True,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -487,26 +487,24 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput = None,
- video: List[PipelineImageInput] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ video: list[PipelineImageInput] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 704,
width: int = 1280,
num_frames: int = 93,
num_inference_steps: int = 35,
guidance_scale: float = 7.0,
fps: int = 16,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
sigma_conditioning: float = 0.0001,
):
@@ -516,9 +514,9 @@ def __call__(
Args:
image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, *optional*):
The image to be used as a conditioning input for the video generation.
- video (`List[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*):
+ video (`list[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*):
The video to be used as a conditioning input for the video generation.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, defaults to `704`):
@@ -539,7 +537,7 @@ def __call__(
The frames per second of the generated video.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -561,7 +559,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py
index 676b7c7a72b0..e144d62d5933 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -80,10 +80,10 @@ def __init__(self, *args, **kwargs):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -98,15 +98,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -196,10 +196,10 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -238,23 +238,23 @@ def _get_t5_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -326,10 +326,10 @@ def prepare_latents(
height: int = 704,
width: int = 1280,
num_frames: int = 121,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype) * self.scheduler.config.sigma_max
@@ -402,32 +402,30 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 704,
width: int = 1280,
num_frames: int = 121,
num_inference_steps: int = 36,
guidance_scale: float = 7.0,
fps: int = 30,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, defaults to `720`):
@@ -448,7 +446,7 @@ def __call__(
The frames per second of the generated video.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -470,7 +468,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py
index df507c3a4f90..377c3c05d284 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -109,10 +109,10 @@ def __init__(self, *args, **kwargs):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -127,15 +127,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -168,7 +168,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -241,10 +241,10 @@ def __init__(
# Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -284,23 +284,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -375,10 +375,10 @@ def prepare_latents(
num_frames: int = 121,
do_classifier_free_guidance: bool = True,
input_frames_guidance: bool = False,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -515,9 +515,9 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput = None,
- video: List[PipelineImageInput] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ video: list[PipelineImageInput] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 704,
width: int = 1280,
num_frames: int = 121,
@@ -526,24 +526,22 @@ def __call__(
input_frames_guidance: bool = False,
augment_sigma: float = 0.001,
fps: int = 30,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, defaults to `720`):
@@ -564,7 +562,7 @@ def __call__(
The frames per second of the generated video.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -586,7 +584,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/cosmos/pipeline_output.py b/src/diffusers/pipelines/cosmos/pipeline_output.py
index ec5f4826f62a..1ded292f8dfb 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_output.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -17,8 +16,8 @@ class CosmosPipelineOutput(BaseOutput):
Output class for Cosmos any-to-world/video pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
@@ -32,9 +31,9 @@ class CosmosImagePipelineOutput(BaseOutput):
Output class for Cosmos any-to-image pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py b/src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py
index 5a70c4f5ff9a..eb8f8106061d 100644
--- a/src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py
+++ b/src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py
@@ -13,8 +13,6 @@
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import torch
from ...models import UNet1DModel
@@ -61,10 +59,10 @@ def __call__(
self,
batch_size: int = 1,
num_inference_steps: int = 100,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- audio_length_in_s: Optional[float] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ audio_length_in_s: float | None = None,
return_dict: bool = True,
- ) -> Union[AudioPipelineOutput, Tuple]:
+ ) -> AudioPipelineOutput | tuple:
r"""
The call function to the pipeline for generation.
diff --git a/src/diffusers/pipelines/ddim/pipeline_ddim.py b/src/diffusers/pipelines/ddim/pipeline_ddim.py
index 39587ca5221d..dc92b34a7565 100644
--- a/src/diffusers/pipelines/ddim/pipeline_ddim.py
+++ b/src/diffusers/pipelines/ddim/pipeline_ddim.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import torch
from ...models import UNet2DModel
@@ -60,13 +58,13 @@ def __init__(self, unet: UNet2DModel, scheduler: DDIMScheduler):
def __call__(
self,
batch_size: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
eta: float = 0.0,
num_inference_steps: int = 50,
- use_clipped_model_output: Optional[bool] = None,
- output_type: Optional[str] = "pil",
+ use_clipped_model_output: bool | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
r"""
The call function to the pipeline for generation.
diff --git a/src/diffusers/pipelines/ddpm/pipeline_ddpm.py b/src/diffusers/pipelines/ddpm/pipeline_ddpm.py
index 0d7766a8cfd0..6d4796cbea1f 100644
--- a/src/diffusers/pipelines/ddpm/pipeline_ddpm.py
+++ b/src/diffusers/pipelines/ddpm/pipeline_ddpm.py
@@ -13,8 +13,6 @@
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import torch
from ...models import UNet2DModel
@@ -57,11 +55,11 @@ def __init__(self, unet: UNet2DModel, scheduler: DDPMScheduler):
def __call__(
self,
batch_size: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
num_inference_steps: int = 1000,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
r"""
The call function to the pipeline for generation.
diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py
index 8fa31f8504d3..b8c70fc6528c 100644
--- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py
+++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py
@@ -2,7 +2,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer
@@ -100,10 +100,10 @@ class IFPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin):
unet: UNet2DConditionModel
scheduler: DDPMScheduler
- feature_extractor: Optional[CLIPImageProcessor]
- safety_checker: Optional[IFSafetyChecker]
+ feature_extractor: CLIPImageProcessor | None
+ safety_checker: IFSafetyChecker | None
- watermarker: Optional[IFWatermarker]
+ watermarker: IFWatermarker | None
bad_punct_regex = re.compile(
r"["
@@ -131,9 +131,9 @@ def __init__(
text_encoder: T5EncoderModel,
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
- safety_checker: Optional[IFSafetyChecker],
- feature_extractor: Optional[CLIPImageProcessor],
- watermarker: Optional[IFWatermarker],
+ safety_checker: IFSafetyChecker | None,
+ feature_extractor: CLIPImageProcessor | None,
+ watermarker: IFWatermarker | None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -168,20 +168,20 @@ def __init__(
@torch.no_grad()
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
clean_caption: bool = False,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
whether to use classifier free guidance or not
@@ -189,7 +189,7 @@ def encode_prompt(
number of images that should be generated per prompt
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -269,7 +269,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif isinstance(negative_prompt, str):
@@ -548,36 +548,36 @@ def _clean_caption(self, caption):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_inference_steps: int = 100,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 7.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ height: int | None = None,
+ width: int | None = None,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
clean_caption: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 7.0):
@@ -586,7 +586,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -599,7 +599,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py
index 507927faf61b..3dadc63f4952 100644
--- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py
+++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py
@@ -2,7 +2,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -124,10 +124,10 @@ class IFImg2ImgPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin):
unet: UNet2DConditionModel
scheduler: DDPMScheduler
- feature_extractor: Optional[CLIPImageProcessor]
- safety_checker: Optional[IFSafetyChecker]
+ feature_extractor: CLIPImageProcessor | None
+ safety_checker: IFSafetyChecker | None
- watermarker: Optional[IFWatermarker]
+ watermarker: IFWatermarker | None
bad_punct_regex = re.compile(
r"["
@@ -155,9 +155,9 @@ def __init__(
text_encoder: T5EncoderModel,
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
- safety_checker: Optional[IFSafetyChecker],
- feature_extractor: Optional[CLIPImageProcessor],
- watermarker: Optional[IFWatermarker],
+ safety_checker: IFSafetyChecker | None,
+ feature_extractor: CLIPImageProcessor | None,
+ watermarker: IFWatermarker | None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -192,20 +192,20 @@ def __init__(
@torch.no_grad()
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
clean_caption: bool = False,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
whether to use classifier free guidance or not
@@ -213,7 +213,7 @@ def encode_prompt(
number of images that should be generated per prompt
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -293,7 +293,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif isinstance(negative_prompt, str):
@@ -430,7 +430,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -662,32 +662,35 @@ def prepare_intermediate_images(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- image: Union[
- PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]
- ] = None,
+ prompt: str | list[str] = None,
+ image: PIL.Image.Image
+ | torch.Tensor
+ | np.ndarray
+ | list[PIL.Image.Image]
+ | list[torch.Tensor]
+ | list[np.ndarray] = None,
strength: float = 0.7,
num_inference_steps: int = 80,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 10.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
clean_caption: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
image (`torch.Tensor` or `PIL.Image.Image`):
@@ -702,7 +705,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 80):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 10.0):
@@ -711,7 +714,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -720,7 +723,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py
index 9bc15c3c6f62..4839a0860462 100644
--- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py
+++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py
@@ -2,7 +2,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -129,10 +129,10 @@ class IFImg2ImgSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoa
scheduler: DDPMScheduler
image_noising_scheduler: DDPMScheduler
- feature_extractor: Optional[CLIPImageProcessor]
- safety_checker: Optional[IFSafetyChecker]
+ feature_extractor: CLIPImageProcessor | None
+ safety_checker: IFSafetyChecker | None
- watermarker: Optional[IFWatermarker]
+ watermarker: IFWatermarker | None
bad_punct_regex = re.compile(
r"["
@@ -161,9 +161,9 @@ def __init__(
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
image_noising_scheduler: DDPMScheduler,
- safety_checker: Optional[IFSafetyChecker],
- feature_extractor: Optional[CLIPImageProcessor],
- watermarker: Optional[IFWatermarker],
+ safety_checker: IFSafetyChecker | None,
+ feature_extractor: CLIPImageProcessor | None,
+ watermarker: IFWatermarker | None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -345,20 +345,20 @@ def _clean_caption(self, caption):
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
clean_caption: bool = False,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
whether to use classifier free guidance or not
@@ -366,7 +366,7 @@ def encode_prompt(
number of images that should be generated per prompt
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -446,7 +446,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif isinstance(negative_prompt, str):
@@ -586,7 +586,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -617,7 +617,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`original_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`original_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -745,26 +745,29 @@ def prepare_intermediate_images(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
- original_image: Union[
- PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]
- ] = None,
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
+ original_image: PIL.Image.Image
+ | torch.Tensor
+ | np.ndarray
+ | list[PIL.Image.Image]
+ | list[torch.Tensor]
+ | list[np.ndarray] = None,
strength: float = 0.8,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 4.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
noise_level: int = 250,
clean_caption: bool = True,
):
@@ -783,13 +786,13 @@ def __call__(
denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will
be maximum and the denoising process will run for the full number of iterations specified in
`num_inference_steps`. A value of 1, therefore, essentially ignores `image`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 4.0):
@@ -798,7 +801,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -807,7 +810,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py
index 9d6cf62020a9..03a9d6f7c5e8 100644
--- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py
+++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py
@@ -2,7 +2,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -127,10 +127,10 @@ class IFInpaintingPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin):
unet: UNet2DConditionModel
scheduler: DDPMScheduler
- feature_extractor: Optional[CLIPImageProcessor]
- safety_checker: Optional[IFSafetyChecker]
+ feature_extractor: CLIPImageProcessor | None
+ safety_checker: IFSafetyChecker | None
- watermarker: Optional[IFWatermarker]
+ watermarker: IFWatermarker | None
bad_punct_regex = re.compile(
r"["
@@ -158,9 +158,9 @@ def __init__(
text_encoder: T5EncoderModel,
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
- safety_checker: Optional[IFSafetyChecker],
- feature_extractor: Optional[CLIPImageProcessor],
- watermarker: Optional[IFWatermarker],
+ safety_checker: IFSafetyChecker | None,
+ feature_extractor: CLIPImageProcessor | None,
+ watermarker: IFWatermarker | None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -196,20 +196,20 @@ def __init__(
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
clean_caption: bool = False,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
whether to use classifier free guidance or not
@@ -217,7 +217,7 @@ def encode_prompt(
number of images that should be generated per prompt
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -297,7 +297,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif isinstance(negative_prompt, str):
@@ -437,7 +437,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -468,7 +468,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`mask_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`mask_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -754,35 +754,41 @@ def prepare_intermediate_images(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- image: Union[
- PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]
- ] = None,
- mask_image: Union[
- PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]
- ] = None,
+ prompt: str | list[str] = None,
+ image: PIL.Image.Image
+ | torch.Tensor
+ | np.ndarray
+ | list[PIL.Image.Image]
+ | list[torch.Tensor]
+ | list[np.ndarray] = None,
+ mask_image: PIL.Image.Image
+ | torch.Tensor
+ | np.ndarray
+ | list[PIL.Image.Image]
+ | list[torch.Tensor]
+ | list[np.ndarray] = None,
strength: float = 1.0,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 7.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
clean_caption: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
image (`torch.Tensor` or `PIL.Image.Image`):
@@ -802,7 +808,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 7.0):
@@ -811,7 +817,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -820,7 +826,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py
index 0122c164d8b8..841382ad9c63 100644
--- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py
+++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py
@@ -2,7 +2,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -131,10 +131,10 @@ class IFInpaintingSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLora
scheduler: DDPMScheduler
image_noising_scheduler: DDPMScheduler
- feature_extractor: Optional[CLIPImageProcessor]
- safety_checker: Optional[IFSafetyChecker]
+ feature_extractor: CLIPImageProcessor | None
+ safety_checker: IFSafetyChecker | None
- watermarker: Optional[IFWatermarker]
+ watermarker: IFWatermarker | None
bad_punct_regex = re.compile(
r"["
@@ -163,9 +163,9 @@ def __init__(
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
image_noising_scheduler: DDPMScheduler,
- safety_checker: Optional[IFSafetyChecker],
- feature_extractor: Optional[CLIPImageProcessor],
- watermarker: Optional[IFWatermarker],
+ safety_checker: IFSafetyChecker | None,
+ feature_extractor: CLIPImageProcessor | None,
+ watermarker: IFWatermarker | None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -347,20 +347,20 @@ def _clean_caption(self, caption):
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
clean_caption: bool = False,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
whether to use classifier free guidance or not
@@ -368,7 +368,7 @@ def encode_prompt(
number of images that should be generated per prompt
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -448,7 +448,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif isinstance(negative_prompt, str):
@@ -589,7 +589,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -620,7 +620,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`original_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`original_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -653,7 +653,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`mask_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`mask_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -833,29 +833,35 @@ def prepare_intermediate_images(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
- original_image: Union[
- PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]
- ] = None,
- mask_image: Union[
- PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]
- ] = None,
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
+ original_image: PIL.Image.Image
+ | torch.Tensor
+ | np.ndarray
+ | list[PIL.Image.Image]
+ | list[torch.Tensor]
+ | list[np.ndarray] = None,
+ mask_image: PIL.Image.Image
+ | torch.Tensor
+ | np.ndarray
+ | list[PIL.Image.Image]
+ | list[torch.Tensor]
+ | list[np.ndarray] = None,
strength: float = 0.8,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_inference_steps: int = 100,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 4.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
noise_level: int = 0,
clean_caption: bool = True,
):
@@ -879,13 +885,13 @@ def __call__(
denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will
be maximum and the denoising process will run for the full number of iterations specified in
`num_inference_steps`. A value of 1, therefore, essentially ignores `image`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 4.0):
@@ -894,7 +900,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -903,7 +909,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py
index ffa60575fe33..52ebebb6f9b4 100644
--- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py
+++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py
@@ -2,7 +2,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -87,10 +87,10 @@ class IFSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixi
scheduler: DDPMScheduler
image_noising_scheduler: DDPMScheduler
- feature_extractor: Optional[CLIPImageProcessor]
- safety_checker: Optional[IFSafetyChecker]
+ feature_extractor: CLIPImageProcessor | None
+ safety_checker: IFSafetyChecker | None
- watermarker: Optional[IFWatermarker]
+ watermarker: IFWatermarker | None
bad_punct_regex = re.compile(
r"["
@@ -119,9 +119,9 @@ def __init__(
unet: UNet2DConditionModel,
scheduler: DDPMScheduler,
image_noising_scheduler: DDPMScheduler,
- safety_checker: Optional[IFSafetyChecker],
- feature_extractor: Optional[CLIPImageProcessor],
- watermarker: Optional[IFWatermarker],
+ safety_checker: IFSafetyChecker | None,
+ feature_extractor: CLIPImageProcessor | None,
+ watermarker: IFWatermarker | None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -303,20 +303,20 @@ def _clean_caption(self, caption):
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
clean_caption: bool = False,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
whether to use classifier free guidance or not
@@ -324,7 +324,7 @@ def encode_prompt(
number of images that should be generated per prompt
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -404,7 +404,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif isinstance(negative_prompt, str):
@@ -547,7 +547,7 @@ def check_inputs(
and not isinstance(check_image_type, np.ndarray)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is"
+ "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is"
f" {type(check_image_type)}"
)
@@ -615,24 +615,24 @@ def preprocess_image(self, image, num_images_per_prompt, device):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
height: int = None,
width: int = None,
- image: Union[PIL.Image.Image, np.ndarray, torch.Tensor] = None,
+ image: PIL.Image.Image | np.ndarray | torch.Tensor = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 4.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
noise_level: int = 250,
clean_caption: bool = True,
):
@@ -640,7 +640,7 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to None):
@@ -652,7 +652,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*, defaults to None):
+ timesteps (`list[int]`, *optional*, defaults to None):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 4.0):
@@ -661,7 +661,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -670,7 +670,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_output.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_output.py
index b8bae89cec03..7fe1cd013835 100644
--- a/src/diffusers/pipelines/deepfloyd_if/pipeline_output.py
+++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -13,17 +12,17 @@ class IFPipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`):
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`):
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
- nsfw_detected (`List[bool]`):
- List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
+ nsfw_detected (`list[bool]`):
+ list of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
(nsfw) content or a watermark. `None` if safety checking could not be performed.
- watermark_detected (`List[bool]`):
- List of flags denoting whether the corresponding generated image likely has a watermark. `None` if safety
+ watermark_detected (`list[bool]`):
+ list of flags denoting whether the corresponding generated image likely has a watermark. `None` if safety
checking could not be performed.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_detected: Optional[List[bool]]
- watermark_detected: Optional[List[bool]]
+ images: list[PIL.Image.Image] | np.ndarray
+ nsfw_detected: list[bool] | None
+ watermark_detected: list[bool] | None
diff --git a/src/diffusers/pipelines/deepfloyd_if/watermark.py b/src/diffusers/pipelines/deepfloyd_if/watermark.py
index e03e3fab026a..d5fe99f681f7 100644
--- a/src/diffusers/pipelines/deepfloyd_if/watermark.py
+++ b/src/diffusers/pipelines/deepfloyd_if/watermark.py
@@ -1,5 +1,3 @@
-from typing import List
-
import PIL.Image
import torch
from PIL import Image
@@ -16,7 +14,7 @@ def __init__(self):
self.register_buffer("watermark_image", torch.zeros((62, 62, 4)))
self.watermark_image_as_pil = None
- def apply_watermark(self, images: List[PIL.Image.Image], sample_size=None):
+ def apply_watermark(self, images: list[PIL.Image.Image], sample_size=None):
# Copied from https://github.com/deep-floyd/IF/blob/b77482e36ca2031cb94dbca1001fc1e6400bf4ab/deepfloyd_if/modules/base.py#L287
h = images[0].height
diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py b/src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py
index f69f905b56c5..ed72e505b9c3 100644
--- a/src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py
+++ b/src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import Optional, Tuple
import torch
from torch import nn
@@ -18,22 +17,22 @@ class TransformationModelOutput(ModelOutput):
last_hidden_state (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`):
Sequence of hidden-states at the output of the last layer of the model.
hidden_states (`tuple(torch.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one
+ tuple of `torch.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one
for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
attentions (`tuple(torch.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+ tuple of `torch.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
sequence_length)`.
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
heads.
"""
- projection_state: Optional[torch.Tensor] = None
+ projection_state: torch.Tensor | None = None
last_hidden_state: torch.Tensor = None
- hidden_states: Optional[Tuple[torch.Tensor]] = None
- attentions: Optional[Tuple[torch.Tensor]] = None
+ hidden_states: tuple[torch.Tensor] | None = None
+ attentions: tuple[torch.Tensor] | None = None
class RobertaSeriesConfig(XLMRobertaConfig):
@@ -73,17 +72,17 @@ def __init__(self, config):
def forward(
self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ head_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ return_dict: bool | None = None,
+ output_hidden_states: bool | None = None,
):
r""" """
diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py
index 6f484aa3e298..9fab42916e9e 100644
--- a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py
+++ b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from packaging import version
@@ -93,10 +93,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -111,15 +111,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -299,9 +299,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -331,16 +331,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -348,7 +348,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -447,7 +447,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -704,35 +704,35 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -741,14 +741,14 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -756,7 +756,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -790,7 +790,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py
index d6bf90120755..b6cd51c6d203 100644
--- a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py
+++ b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -82,7 +82,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -121,10 +121,10 @@ def preprocess(image):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -139,15 +139,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -327,9 +327,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -359,16 +359,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -376,7 +376,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -475,7 +475,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -772,35 +772,35 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
strength: float = 0.8,
- num_inference_steps: Optional[int] = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- guidance_scale: Optional[float] = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ num_inference_steps: int | None = 50,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ guidance_scale: float | None = 7.5,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = 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"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -815,14 +815,14 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference. This parameter is modulated by `strength`.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -830,7 +830,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -856,7 +856,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py
index dd174ae3c21f..259a86756965 100644
--- a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py
+++ b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -16,13 +15,13 @@ class AltDiffusionPipelineOutput(BaseOutput):
Output class for Alt Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
- nsfw_content_detected (`List[bool]`)
- List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
+ nsfw_content_detected (`list[bool]`)
+ list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
`None` if safety checking could not be performed.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_content_detected: Optional[List[bool]]
+ images: list[PIL.Image.Image] | np.ndarray
+ nsfw_content_detected: list[bool] | None
diff --git a/src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py b/src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py
index 81fa999eb1fb..f63fc8aacbc8 100644
--- a/src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py
+++ b/src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py
@@ -14,7 +14,6 @@
from math import acos, sin
-from typing import List, Tuple, Union
import numpy as np
import torch
@@ -53,7 +52,7 @@ def __init__(
vqvae: AutoencoderKL,
unet: UNet2DConditionModel,
mel: Mel,
- scheduler: Union[DDIMScheduler, DDPMScheduler],
+ scheduler: DDIMScheduler | DDPMScheduler,
):
super().__init__()
self.register_modules(unet=unet, scheduler=scheduler, mel=mel, vqvae=vqvae)
@@ -84,10 +83,7 @@ def __call__(
noise: torch.Tensor = None,
encoding: torch.Tensor = None,
return_dict=True,
- ) -> Union[
- Union[AudioPipelineOutput, ImagePipelineOutput],
- Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
- ]:
+ ) -> AudioPipelineOutput | ImagePipelineOutput | tuple[list[Image.Image], tuple[int, list[np.ndarray]]]:
"""
The call function to the pipeline for generation.
@@ -170,8 +166,8 @@ def __call__(
```
Returns:
- `List[PIL Image]`:
- A list of Mel spectrograms (`float`, `List[np.ndarray]`) with the sample rate and raw audio.
+ `list[PIL Image]`:
+ A list of Mel spectrograms (`float`, `list[np.ndarray]`) with the sample rate and raw audio.
"""
steps = steps or self.get_default_steps()
@@ -268,13 +264,13 @@ def __call__(
return BaseOutput(**AudioPipelineOutput(np.array(audios)[:, np.newaxis, :]), **ImagePipelineOutput(images))
@torch.no_grad()
- def encode(self, images: List[Image.Image], steps: int = 50) -> np.ndarray:
+ def encode(self, images: list[Image.Image], steps: int = 50) -> np.ndarray:
"""
Reverse the denoising step process to recover a noisy image from the generated image.
Args:
- images (`List[PIL Image]`):
- List of images to encode.
+ images (`list[PIL Image]`):
+ list of images to encode.
steps (`int`):
Number of encoding steps to perform (defaults to `50`).
diff --git a/src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py b/src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py
index 0bb24ed0b1ce..70a65e2ef5be 100644
--- a/src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py
+++ b/src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Tuple, Union
import torch
@@ -47,13 +46,13 @@ def __init__(self, vqvae: VQModel, unet: UNet2DModel, scheduler: DDIMScheduler):
def __call__(
self,
batch_size: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
eta: float = 0.0,
num_inference_steps: int = 50,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
**kwargs,
- ) -> Union[Tuple, ImagePipelineOutput]:
+ ) -> tuple | ImagePipelineOutput:
r"""
The call function to the pipeline for generation.
diff --git a/src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py b/src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py
index 71e3e156e0e4..fb116511f727 100644
--- a/src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py
+++ b/src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py
@@ -13,8 +13,6 @@
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import torch
from ....models import UNet2DModel
@@ -52,11 +50,11 @@ def __call__(
self,
batch_size: int = 1,
num_inference_steps: int = 50,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
**kwargs,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
r"""
The call function to the pipeline for generation.
diff --git a/src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py b/src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py
index 56c6007ae886..3231d5e13049 100644
--- a/src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py
+++ b/src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py
@@ -13,8 +13,6 @@
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import numpy as np
import PIL.Image
import torch
@@ -30,7 +28,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess
-def _preprocess_image(image: Union[List, PIL.Image.Image, torch.Tensor]):
+def _preprocess_image(image: list | PIL.Image.Image | torch.Tensor):
deprecation_message = "The preprocess method is deprecated and will be removed in diffusers 1.0.0. Please use VaeImageProcessor.preprocess(...) instead"
deprecate("preprocess", "1.0.0", deprecation_message, standard_warn=False)
if isinstance(image, torch.Tensor):
@@ -53,7 +51,7 @@ def _preprocess_image(image: Union[List, PIL.Image.Image, torch.Tensor]):
return image
-def _preprocess_mask(mask: Union[List, PIL.Image.Image, torch.Tensor]):
+def _preprocess_mask(mask: list | PIL.Image.Image | torch.Tensor):
if isinstance(mask, torch.Tensor):
return mask
elif isinstance(mask, PIL.Image.Image):
@@ -98,16 +96,16 @@ def __init__(self, unet: UNet2DModel, scheduler: RePaintScheduler):
@torch.no_grad()
def __call__(
self,
- image: Union[torch.Tensor, PIL.Image.Image],
- mask_image: Union[torch.Tensor, PIL.Image.Image],
+ image: torch.Tensor | PIL.Image.Image,
+ mask_image: torch.Tensor | PIL.Image.Image,
num_inference_steps: int = 250,
eta: float = 0.0,
jump_length: int = 10,
jump_n_sample: int = 10,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
r"""
The call function to the pipeline for generation.
diff --git a/src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py b/src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py
index 3f04db7ad699..c6abdba42d3c 100644
--- a/src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py
+++ b/src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import torch
from ....models import UNet2DModel
@@ -48,11 +46,11 @@ def __call__(
self,
batch_size: int = 1,
num_inference_steps: int = 2000,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
**kwargs,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
r"""
The call function to the pipeline for generation.
diff --git a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py
index 8985a6f88800..76b8576468d2 100644
--- a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py
+++ b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py
@@ -16,7 +16,7 @@
import dataclasses
import math
import os
-from typing import Any, Callable, List, Mapping, MutableMapping, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Mapping, MutableMapping, Sequence
import numpy as np
import torch
@@ -91,10 +91,10 @@ class NoteRepresentationConfig:
@dataclasses.dataclass
class NoteEventData:
pitch: int
- velocity: Optional[int] = None
- program: Optional[int] = None
- is_drum: Optional[bool] = None
- instrument: Optional[int] = None
+ velocity: int | None = None
+ program: int | None = None
+ is_drum: bool | None = None
+ instrument: int | None = None
@dataclasses.dataclass
@@ -102,7 +102,7 @@ class NoteEncodingState:
"""Encoding state for note transcription, keeping track of active pitches."""
# velocity bin for active pitches and programs
- active_pitches: MutableMapping[Tuple[int, int], int] = dataclasses.field(default_factory=dict)
+ active_pitches: MutableMapping[tuple[int, int], int] = dataclasses.field(default_factory=dict)
@dataclasses.dataclass
@@ -153,7 +153,7 @@ class Codec:
and specified separately.
"""
- def __init__(self, max_shift_steps: int, steps_per_second: float, event_ranges: List[EventRange]):
+ def __init__(self, max_shift_steps: int, steps_per_second: float, event_ranges: list[EventRange]):
"""Define Codec.
Args:
@@ -197,7 +197,7 @@ def encode_event(self, event: Event) -> int:
raise ValueError(f"Unknown event type: {event.type}")
- def event_type_range(self, event_type: str) -> Tuple[int, int]:
+ def event_type_range(self, event_type: str) -> tuple[int, int]:
"""Return [min_id, max_id] for an event type."""
offset = 0
for er in self._event_ranges:
@@ -280,7 +280,7 @@ def audio_to_frames(
samples,
hop_size: int,
frame_rate: int,
-) -> Tuple[Sequence[Sequence[int]], torch.Tensor]:
+) -> tuple[Sequence[Sequence[int]], torch.Tensor]:
"""Convert audio samples to non-overlapping frames and frame times."""
frame_size = hop_size
samples = np.pad(samples, [0, frame_size - len(samples) % frame_size], mode="constant")
@@ -301,7 +301,7 @@ def audio_to_frames(
def note_sequence_to_onsets_and_offsets_and_programs(
ns: note_seq.NoteSequence,
-) -> Tuple[Sequence[float], Sequence[NoteEventData]]:
+) -> tuple[Sequence[float], Sequence[NoteEventData]]:
"""Extract onset & offset times and pitches & programs from a NoteSequence.
The onset & offset times will not necessarily be in sorted order.
@@ -348,7 +348,7 @@ def velocity_to_bin(velocity, num_velocity_bins):
def note_event_data_to_events(
- state: Optional[NoteEncodingState],
+ state: NoteEncodingState | None,
value: NoteEventData,
codec: Codec,
) -> Sequence[Event]:
@@ -632,7 +632,7 @@ def __init__(self):
self.tokenizer = Tokenizer(self.codec.num_classes)
self.note_representation_config = NoteRepresentationConfig(onsets_only=False, include_ties=True)
- def __call__(self, midi: Union[bytes, os.PathLike, str]):
+ def __call__(self, midi: bytes | os.PathLike | str):
if not isinstance(midi, bytes):
with open(midi, "rb") as f:
midi = f.read()
diff --git a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py
index be07b1b15ea8..269e7405d10d 100644
--- a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py
+++ b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py
@@ -14,7 +14,7 @@
# limitations under the License.
import math
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -129,14 +129,14 @@ def decode(self, encodings_and_masks, input_tokens, noise_time):
@torch.no_grad()
def __call__(
self,
- input_tokens: List[List[int]],
- generator: Optional[torch.Generator] = None,
+ input_tokens: list[list[int]],
+ generator: torch.Generator | None = None,
num_inference_steps: int = 100,
return_dict: bool = True,
output_type: str = "np",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- ) -> Union[AudioPipelineOutput, Tuple]:
+ ) -> AudioPipelineOutput | tuple:
if (callback_steps is None) or (
callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
):
@@ -148,8 +148,8 @@ def __call__(
The call function to the pipeline for generation.
Args:
- input_tokens (`List[List[int]]`):
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ input_tokens (`list[list[int]]`):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
num_inference_steps (`int`, *optional*, defaults to 100):
diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py
index 08f8c7e26fae..dae5e600d773 100644
--- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py
+++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -63,7 +63,7 @@ def preprocess(image):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -259,9 +259,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -292,16 +292,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -309,7 +309,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -408,7 +408,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -633,31 +633,31 @@ def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dt
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
- source_prompt: Union[str, List[str]],
+ prompt: str | list[str],
+ source_prompt: str | list[str],
image: PipelineImageInput = None,
strength: float = 0.8,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 7.5,
- source_guidance_scale: Optional[float] = 1,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 7.5,
+ source_guidance_scale: float | None = 1,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image` or tensor representing an image batch to be used as the starting point. Can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
strength (`float`, *optional*, defaults to 0.8):
@@ -680,7 +680,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py
index fcd8bf317adf..f526dc419cea 100644
--- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py
+++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py
@@ -1,5 +1,5 @@
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -76,7 +76,7 @@ class OnnxStableDiffusionInpaintPipelineLegacy(DiffusionPipeline):
text_encoder: OnnxRuntimeModel
tokenizer: CLIPTokenizer
unet: OnnxRuntimeModel
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler
safety_checker: OnnxRuntimeModel
feature_extractor: CLIPImageProcessor
@@ -87,7 +87,7 @@ def __init__(
text_encoder: OnnxRuntimeModel,
tokenizer: CLIPTokenizer,
unet: OnnxRuntimeModel,
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler,
safety_checker: OnnxRuntimeModel,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = True,
@@ -152,24 +152,24 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion.OnnxStableDiffusionPipeline._encode_prompt
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- num_images_per_prompt: Optional[int],
+ prompt: str | list[str],
+ num_images_per_prompt: int | None,
do_classifier_free_guidance: bool,
- negative_prompt: Optional[str],
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
+ negative_prompt: str | None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
num_images_per_prompt (`int`):
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
prompt_embeds (`np.ndarray`, *optional*):
@@ -214,7 +214,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -297,28 +297,28 @@ def check_inputs(
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[np.ndarray, PIL.Image.Image] = None,
- mask_image: Union[np.ndarray, PIL.Image.Image] = None,
+ prompt: str | list[str],
+ image: np.ndarray | PIL.Image.Image = None,
+ mask_image: np.ndarray | PIL.Image.Image = None,
strength: float = 0.8,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[np.random.RandomState] = None,
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
- output_type: Optional[str] = "pil",
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 7.5,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: np.random.RandomState | None = None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
+ callback: Callable[[int, int, np.ndarray], None] | None = None,
callback_steps: int = 1,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
image (`nd.ndarray` or `PIL.Image.Image`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
@@ -343,7 +343,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py
index ba0dd66c2938..650695b604c1 100644
--- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py
+++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -230,9 +230,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -263,16 +263,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -280,7 +280,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -379,7 +379,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -562,31 +562,31 @@ def prepare_latents(self, image, timestep, num_images_per_prompt, dtype, device,
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- image: Union[torch.Tensor, PIL.Image.Image] = None,
- mask_image: Union[torch.Tensor, PIL.Image.Image] = None,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | PIL.Image.Image = None,
+ mask_image: torch.Tensor | PIL.Image.Image = None,
strength: float = 0.8,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- add_predicted_noise: Optional[bool] = False,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 7.5,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ add_predicted_noise: bool | None = False,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
image (`torch.Tensor` or `PIL.Image.Image`):
@@ -611,7 +611,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale`
is less than `1`).
diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py
index b7a0be57c12b..851820c00aed 100644
--- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py
+++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py
@@ -13,7 +13,7 @@
import copy
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
@@ -163,9 +163,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -196,16 +196,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -213,7 +213,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -312,7 +312,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -614,30 +614,30 @@ def edit_model(
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -649,7 +649,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -657,7 +657,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py
index c236e73bf448..ea81be87a0f4 100644
--- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py
+++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
@@ -160,9 +160,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -193,16 +193,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -210,7 +210,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -309,7 +309,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -488,25 +488,25 @@ def _cumsum(self, input, dim, debug=False):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
parallel: int = 10,
tolerance: float = 0.1,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
debug: bool = False,
clip_skip: int = None,
):
@@ -514,7 +514,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -533,7 +533,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -541,7 +541,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py
index 2a461ae20cc9..0955b6fe48a1 100644
--- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py
+++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -62,13 +62,13 @@ class Pix2PixInversionPipelineOutput(BaseOutput, TextualInversionLoaderMixin):
Args:
latents (`torch.Tensor`)
inverted latents tensor
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
latents: torch.Tensor
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
EXAMPLE_DOC_STRING = """
@@ -328,7 +328,7 @@ def __init__(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
- scheduler: Union[DDPMScheduler, DDIMScheduler, EulerAncestralDiscreteScheduler, LMSDiscreteScheduler],
+ scheduler: DDPMScheduler | DDIMScheduler | EulerAncestralDiscreteScheduler | LMSDiscreteScheduler,
feature_extractor: CLIPImageProcessor,
safety_checker: StableDiffusionSafetyChecker,
inverse_scheduler: DDIMInverseScheduler,
@@ -378,9 +378,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -411,16 +411,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -428,7 +428,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -527,7 +527,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -709,7 +709,7 @@ def construct_direction(self, embs_source: torch.Tensor, embs_target: torch.Tens
return (embs_target.mean(0) - embs_source.mean(0)).unsqueeze(0)
@torch.no_grad()
- def get_embeds(self, prompt: List[str], batch_size: int = 16) -> torch.Tensor:
+ def get_embeds(self, prompt: list[str], batch_size: int = 16) -> torch.Tensor:
num_prompts = len(prompt)
embeds = []
for i in range(0, num_prompts, batch_size):
@@ -818,33 +818,33 @@ def kl_divergence(self, hidden_states):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] | None = None,
source_embeds: torch.Tensor = None,
target_embeds: torch.Tensor = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
cross_attention_guidance_amount: float = 0.1,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
source_embeds (`torch.Tensor`):
@@ -866,7 +866,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -875,7 +875,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1109,19 +1109,19 @@ def __call__(
@replace_example_docstring(EXAMPLE_INVERT_DOC_STRING)
def invert(
self,
- prompt: Optional[str] = None,
+ prompt: str | None = None,
image: PipelineImageInput = None,
num_inference_steps: int = 50,
guidance_scale: float = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
cross_attention_guidance_amount: float = 0.1,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
+ cross_attention_kwargs: dict[str, Any] | None = None,
lambda_auto_corr: float = 20.0,
lambda_kl: float = 20.0,
num_reg_steps: int = 5,
@@ -1131,10 +1131,10 @@ def invert(
Function used to generate inverted latents given a prompt and image.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, or tensor representing an image batch which will be used for conditioning. Can also accept
image latents as `image`, if passing latents directly, it will not be encoded again.
num_inference_steps (`int`, *optional*, defaults to 50):
@@ -1146,7 +1146,7 @@ def invert(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py b/src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py
index 50b8b0bcbc1d..ce2c785c8d98 100644
--- a/src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py
+++ b/src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Tuple, Union
-
import torch
from ....models import UNet2DModel
@@ -46,11 +44,11 @@ def __call__(
self,
batch_size: int = 1,
num_inference_steps: int = 50,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
**kwargs,
- ) -> Union[Tuple, ImagePipelineOutput]:
+ ) -> tuple | ImagePipelineOutput:
r"""
The call function to the pipeline for generation.
diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py
index 7c25713cd1d7..7be159d77af5 100644
--- a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py
+++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -277,7 +277,7 @@ class UNetFlatConditionModel(ModelMixin, ConfigMixin):
for all models (such as downloading or saving).
Parameters:
- sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+ sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`):
Height and width of input/output sample.
in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
@@ -285,17 +285,17 @@ class UNetFlatConditionModel(ModelMixin, ConfigMixin):
flip_sin_to_cos (`bool`, *optional*, defaults to `False`):
Whether to flip the sin to cos in the time embedding.
freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
- down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "DownBlockFlat")`):
+ down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "DownBlockFlat")`):
The tuple of downsample blocks to use.
mid_block_type (`str`, *optional*, defaults to `"UNetMidBlockFlatCrossAttn"`):
Block type for middle of UNet, it can be one of `UNetMidBlockFlatCrossAttn`, `UNetMidBlockFlat`, or
`UNetMidBlockFlatSimpleCrossAttn`. If `None`, the mid block layer is skipped.
- up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat")`):
+ up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat")`):
The tuple of upsample blocks to use.
- only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
+ only_cross_attention(`bool` or `tuple[bool]`, *optional*, default to `False`):
Whether to include self-attention in the basic transformer blocks, see
[`~models.attention.BasicTransformerBlock`].
- block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+ block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
The tuple of output channels for each block.
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
@@ -305,15 +305,15 @@ class UNetFlatConditionModel(ModelMixin, ConfigMixin):
norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
If `None`, normalization and activation layers is skipped in post-processing.
norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
- cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+ cross_attention_dim (`int` or `tuple[int]`, *optional*, defaults to 1280):
The dimension of the cross attention features.
- transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+ transformer_layers_per_block (`int`, `tuple[int]`, or `tuple[tuple]` , *optional*, defaults to 1):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
[`~models.unet_2d_blocks.CrossAttnDownBlockFlat`], [`~models.unet_2d_blocks.CrossAttnUpBlockFlat`],
[`~models.unet_2d_blocks.UNetMidBlockFlatCrossAttn`].
- reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None):
+ reverse_transformer_layers_per_block : (`tuple[tuple]`, *optional*, defaults to None):
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling
- blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for
+ blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `tuple[tuple]` and for
[`~models.unet_2d_blocks.CrossAttnDownBlockFlat`], [`~models.unet_2d_blocks.CrossAttnUpBlockFlat`],
[`~models.unet_2d_blocks.UNetMidBlockFlatCrossAttn`].
encoder_hid_dim (`int`, *optional*, defaults to None):
@@ -368,63 +368,63 @@ class conditioning with `class_embed_type` equal to `None`.
@register_to_config
def __init__(
self,
- sample_size: Optional[int] = None,
+ sample_size: int | None = None,
in_channels: int = 4,
out_channels: int = 4,
center_input_sample: bool = False,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
- down_block_types: Tuple[str, ...] = (
+ down_block_types: tuple[str] = (
"CrossAttnDownBlockFlat",
"CrossAttnDownBlockFlat",
"CrossAttnDownBlockFlat",
"DownBlockFlat",
),
- mid_block_type: Optional[str] = "UNetMidBlockFlatCrossAttn",
- up_block_types: Tuple[str, ...] = (
+ mid_block_type: str = "UNetMidBlockFlatCrossAttn",
+ up_block_types: tuple[str] = (
"UpBlockFlat",
"CrossAttnUpBlockFlat",
"CrossAttnUpBlockFlat",
"CrossAttnUpBlockFlat",
),
- only_cross_attention: Union[bool, Tuple[bool]] = False,
- block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
- layers_per_block: Union[int, Tuple[int]] = 2,
+ only_cross_attention: bool | tuple[bool] = False,
+ block_out_channels: tuple[int] = (320, 640, 1280, 1280),
+ layers_per_block: int | tuple[int] = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
dropout: float = 0.0,
act_fn: str = "silu",
- norm_num_groups: Optional[int] = 32,
+ norm_num_groups: int | None = 32,
norm_eps: float = 1e-5,
- cross_attention_dim: Union[int, Tuple[int]] = 1280,
- transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
- reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
- encoder_hid_dim: Optional[int] = None,
- encoder_hid_dim_type: Optional[str] = None,
- attention_head_dim: Union[int, Tuple[int]] = 8,
- num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+ cross_attention_dim: int | tuple[int] = 1280,
+ transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1,
+ reverse_transformer_layers_per_block: tuple[tuple[int]] | None = None,
+ encoder_hid_dim: int | None = None,
+ encoder_hid_dim_type: str | None = None,
+ attention_head_dim: int | tuple[int] = 8,
+ num_attention_heads: int | tuple[int] | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
- class_embed_type: Optional[str] = None,
- addition_embed_type: Optional[str] = None,
- addition_time_embed_dim: Optional[int] = None,
- num_class_embeds: Optional[int] = None,
+ class_embed_type: str | None = None,
+ addition_embed_type: str | None = None,
+ addition_time_embed_dim: int | None = None,
+ num_class_embeds: int | None = None,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
resnet_skip_time_act: bool = False,
resnet_out_scale_factor: int = 1.0,
time_embedding_type: str = "positional",
- time_embedding_dim: Optional[int] = None,
- time_embedding_act_fn: Optional[str] = None,
- timestep_post_act: Optional[str] = None,
- time_cond_proj_dim: Optional[int] = None,
+ time_embedding_dim: int | None = None,
+ time_embedding_act_fn: str | None = None,
+ timestep_post_act: str | None = None,
+ time_cond_proj_dim: int | None = None,
conv_in_kernel: int = 3,
conv_out_kernel: int = 3,
- projection_class_embeddings_input_dim: Optional[int] = None,
+ projection_class_embeddings_input_dim: int | None = None,
attention_type: str = "default",
class_embeddings_concat: bool = False,
- mid_block_only_cross_attention: Optional[bool] = None,
- cross_attention_norm: Optional[str] = None,
+ mid_block_only_cross_attention: bool | None = None,
+ cross_attention_norm: str | None = None,
addition_embed_type_num_heads=64,
):
super().__init__()
@@ -826,7 +826,7 @@ def __init__(
)
@property
- def attn_processors(self) -> Dict[str, AttentionProcessor]:
+ def attn_processors(self) -> dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
@@ -835,7 +835,7 @@ def attn_processors(self) -> Dict[str, AttentionProcessor]:
# set recursively
processors = {}
- def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict[str, AttentionProcessor]):
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
@@ -849,7 +849,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors:
return processors
- def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+ def set_attn_processor(self, processor: AttentionProcessor | dict[str, AttentionProcessor]):
r"""
Sets the attention processor to use to compute attention.
@@ -952,7 +952,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
# Recursively walk through all the children.
# Any children which exposes the set_attention_slice method
# gets the message
- def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]):
if hasattr(module, "set_attention_slice"):
module.set_attention_slice(slice_size.pop())
@@ -1037,19 +1037,19 @@ def unload_lora(self):
def forward(
self,
sample: torch.Tensor,
- timestep: Union[torch.Tensor, float, int],
+ timestep: torch.Tensor | float | int,
encoder_hidden_states: torch.Tensor,
- class_labels: Optional[torch.Tensor] = None,
- timestep_cond: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
- down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
- mid_block_additional_residual: Optional[torch.Tensor] = None,
- down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ class_labels: torch.Tensor | None = None,
+ timestep_cond: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ added_cond_kwargs: dict[str, torch.Tensor] | None = None,
+ down_block_additional_residuals: tuple[torch.Tensor] | None = None,
+ mid_block_additional_residual: torch.Tensor | None = None,
+ down_intrablock_additional_residuals: tuple[torch.Tensor] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[UNet2DConditionOutput, Tuple]:
+ ) -> UNet2DConditionOutput | tuple:
r"""
The [`UNetFlatConditionModel`] forward method.
@@ -1579,8 +1579,8 @@ def __init__(
self.gradient_checkpointing = False
def forward(
- self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
output_states = ()
for resnet in self.resnets:
@@ -1608,7 +1608,7 @@ def __init__(
temb_channels: int,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
@@ -1695,13 +1695,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- additional_residuals: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ additional_residuals: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
output_states = ()
blocks = list(zip(self.resnets, self.attentions))
@@ -1751,7 +1751,7 @@ def __init__(
prev_output_channel: int,
out_channels: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
@@ -1797,9 +1797,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- upsample_size: Optional[int] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ upsample_size: int | None = None,
*args,
**kwargs,
) -> torch.Tensor:
@@ -1853,10 +1853,10 @@ def __init__(
out_channels: int,
prev_output_channel: int,
temb_channels: int,
- resolution_idx: Optional[int] = None,
+ resolution_idx: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
@@ -1940,13 +1940,13 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- res_hidden_states_tuple: Tuple[torch.Tensor, ...],
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ res_hidden_states_tuple: tuple[torch.Tensor, ...],
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ upsample_size: int | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
@@ -2023,7 +2023,7 @@ class UNetMidBlockFlat(nn.Module):
resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.
resnet_groups (`int`, *optional*, defaults to 32):
The number of groups to use in the group normalization layers of the resnet blocks.
- attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.
+ attn_groups (`int | None`, *optional*, defaults to None): The number of groups for the attention blocks.
resnet_pre_norm (`bool`, *optional*, defaults to `True`):
Whether to use pre-normalization for the resnet blocks.
add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.
@@ -2048,7 +2048,7 @@ def __init__(
resnet_time_scale_shift: str = "default", # default, spatial
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
- attn_groups: Optional[int] = None,
+ attn_groups: int | None = None,
resnet_pre_norm: bool = True,
add_attention: bool = True,
attention_head_dim: int = 1,
@@ -2154,7 +2154,7 @@ def __init__(
self.gradient_checkpointing = False
- def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
hidden_states = self.resnets[0](hidden_states, temb)
for attn, resnet in zip(self.attentions, self.resnets[1:]):
if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -2175,15 +2175,15 @@ def __init__(
self,
in_channels: int,
temb_channels: int,
- out_channels: Optional[int] = None,
+ out_channels: int | None = None,
dropout: float = 0.0,
num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+ transformer_layers_per_block: int | tuple[int] = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
- resnet_groups_out: Optional[int] = None,
+ resnet_groups_out: int | None = None,
resnet_pre_norm: bool = True,
num_attention_heads: int = 1,
output_scale_factor: float = 1.0,
@@ -2276,11 +2276,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
if cross_attention_kwargs is not None:
if cross_attention_kwargs.get("scale", None) is not None:
@@ -2330,7 +2330,7 @@ def __init__(
cross_attention_dim: int = 1280,
skip_time_act: bool = False,
only_cross_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
+ cross_attention_norm: str | None = None,
):
super().__init__()
@@ -2401,11 +2401,11 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- temb: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
+ temb: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
if cross_attention_kwargs.get("scale", None) is not None:
diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py
index 9ff8e9857791..101a1b72e7f9 100644
--- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py
+++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py
@@ -1,5 +1,5 @@
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import PIL.Image
import torch
@@ -81,26 +81,26 @@ def __init__(
@torch.no_grad()
def image_variation(
self,
- image: Union[torch.Tensor, PIL.Image.Image],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ image: torch.Tensor | PIL.Image.Image,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
):
r"""
The call function to the pipeline for generation.
Args:
- image (`PIL.Image.Image`, `List[PIL.Image.Image]` or `torch.Tensor`):
+ image (`PIL.Image.Image`, `list[PIL.Image.Image]` or `torch.Tensor`):
The image prompt or prompts to guide the image generation.
height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -112,7 +112,7 @@ def image_variation(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -193,26 +193,26 @@ def image_variation(
@torch.no_grad()
def text_to_image(
self,
- prompt: Union[str, List[str]],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str],
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation.
height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -224,7 +224,7 @@ def text_to_image(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -301,27 +301,27 @@ def text_to_image(
@torch.no_grad()
def dual_guided(
self,
- prompt: Union[PIL.Image.Image, List[PIL.Image.Image]],
- image: Union[str, List[str]],
+ prompt: PIL.Image.Image | list[PIL.Image.Image],
+ image: str | list[str],
text_to_image_strength: float = 0.5,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- num_images_per_prompt: Optional[int] = 1,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation.
height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -333,7 +333,7 @@ def dual_guided(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -341,7 +341,7 @@ def dual_guided(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py
index 0252f4f6af7f..8f8fb712e023 100644
--- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py
+++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -158,7 +158,7 @@ def _encode_text_prompt(self, prompt, device, num_images_per_prompt, do_classifi
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -249,7 +249,7 @@ def _encode_image_prompt(self, prompt, device, num_images_per_prompt, do_classif
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -369,7 +369,7 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
latents = latents * self.scheduler.init_noise_sigma
return latents
- def set_transformer_params(self, mix_ratio: float = 0.5, condition_types: Tuple = ("text", "image")):
+ def set_transformer_params(self, mix_ratio: float = 0.5, condition_types: tuple = ("text", "image")):
for name, module in self.image_unet.named_modules():
if isinstance(module, DualTransformer2DModel):
module.mix_ratio = mix_ratio
@@ -385,20 +385,20 @@ def set_transformer_params(self, mix_ratio: float = 0.5, condition_types: Tuple
@torch.no_grad()
def __call__(
self,
- prompt: Union[PIL.Image.Image, List[PIL.Image.Image]],
- image: Union[str, List[str]],
+ prompt: PIL.Image.Image | list[PIL.Image.Image],
+ image: str | list[str],
text_to_image_strength: float = 0.5,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- num_images_per_prompt: Optional[int] = 1,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
**kwargs,
):
@@ -406,7 +406,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation.
height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -418,7 +418,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -426,7 +426,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py
index 034a0226419b..348417ad11df 100644
--- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py
+++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -84,7 +84,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -92,7 +92,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
"""
@@ -122,7 +122,7 @@ def normalize_embeddings(encoder_output):
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance:
- uncond_images: List[str]
+ uncond_images: list[str]
if negative_prompt is None:
uncond_images = [np.zeros((512, 512, 3)) + 0.5] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -196,7 +196,7 @@ def check_inputs(self, image, height, width, callback_steps):
and not isinstance(image, list)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+ "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
f" {type(image)}"
)
@@ -237,19 +237,19 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@torch.no_grad()
def __call__(
self,
- image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
**kwargs,
):
@@ -257,7 +257,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- image (`PIL.Image.Image`, `List[PIL.Image.Image]` or `torch.Tensor`):
+ image (`PIL.Image.Image`, `list[PIL.Image.Image]` or `torch.Tensor`):
The image prompt or prompts to guide the image generation.
height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -269,7 +269,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py
index 2f54f4fc98a4..8e2af1063421 100644
--- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py
+++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer
@@ -108,7 +108,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -116,7 +116,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
"""
@@ -164,7 +164,7 @@ def normalize_embeddings(encoder_output):
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -323,19 +323,19 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str],
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
**kwargs,
):
@@ -343,7 +343,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation.
height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -355,7 +355,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
diff --git a/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py b/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py
index e8617a54b691..8f3b9512e888 100644
--- a/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py
+++ b/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import CLIPTextModel, CLIPTokenizer
@@ -33,7 +33,7 @@ class LearnedClassifierFreeSamplingEmbeddings(ModelMixin, ConfigMixin):
"""
@register_to_config
- def __init__(self, learnable: bool, hidden_size: Optional[int] = None, length: Optional[int] = None):
+ def __init__(self, learnable: bool, hidden_size: int | None = None, length: int | None = None):
super().__init__()
self.learnable = learnable
@@ -163,23 +163,23 @@ def _encode_prompt(self, prompt, num_images_per_prompt, do_classifier_free_guida
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_inference_steps: int = 100,
guidance_scale: float = 5.0,
truncation_rate: float = 1.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation.
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
diff --git a/src/diffusers/pipelines/dit/pipeline_dit.py b/src/diffusers/pipelines/dit/pipeline_dit.py
index 68ff6c9b559a..505360fc78a7 100644
--- a/src/diffusers/pipelines/dit/pipeline_dit.py
+++ b/src/diffusers/pipelines/dit/pipeline_dit.py
@@ -18,8 +18,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Dict, List, Optional, Tuple, Union
-
import torch
from ...models import AutoencoderKL, DiTTransformer2DModel
@@ -62,7 +60,7 @@ def __init__(
transformer: DiTTransformer2DModel,
vae: AutoencoderKL,
scheduler: KarrasDiffusionSchedulers,
- id2label: Optional[Dict[int, str]] = None,
+ id2label: dict[int, str] | None = None,
):
super().__init__()
self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
@@ -75,7 +73,7 @@ def __init__(
self.labels[label.lstrip().rstrip()] = int(key)
self.labels = dict(sorted(self.labels.items()))
- def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
+ def get_label_ids(self, label: str | list[str]) -> list[int]:
r"""
Map label strings from ImageNet to corresponding class ids.
@@ -103,19 +101,19 @@ def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
@torch.no_grad()
def __call__(
self,
- class_labels: List[int],
+ class_labels: list[int],
guidance_scale: float = 4.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
num_inference_steps: int = 50,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
r"""
The call function to the pipeline for generation.
Args:
- class_labels (List[int]):
- List of ImageNet class labels for the images to be generated.
+ class_labels (list[int]):
+ list of ImageNet class labels for the images to be generated.
guidance_scale (`float`, *optional*, defaults to 4.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
diff --git a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py
index 86c4d6812130..6ec8f44e6d1a 100755
--- a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py
+++ b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py
@@ -14,7 +14,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import torch
from transformers import (
@@ -126,10 +126,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -144,15 +144,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -195,9 +195,9 @@ class EasyAnimatePipeline(DiffusionPipeline):
Args:
vae ([`AutoencoderKLMagvit`]):
Variational Auto-Encoder (VAE) Model to encode and decode video to and from latent representations.
- text_encoder (Optional[`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel`]):
+ text_encoder (`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel` | None):
EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1.
- tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]):
+ tokenizer (`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer` | None):
A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text.
transformer ([`EasyAnimateTransformer3DModel`]):
The EasyAnimate model designed by EasyAnimate Team.
@@ -211,8 +211,8 @@ class EasyAnimatePipeline(DiffusionPipeline):
def __init__(
self,
vae: AutoencoderKLMagvit,
- text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel],
- tokenizer: Union[Qwen2Tokenizer, BertTokenizer],
+ text_encoder: Qwen2VLForConditionalGeneration | BertModel,
+ tokenizer: Qwen2Tokenizer | BertTokenizer,
transformer: EasyAnimateTransformer3DModel,
scheduler: FlowMatchEulerDiscreteScheduler,
):
@@ -240,23 +240,23 @@ def __init__(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -266,7 +266,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -525,35 +525,33 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- num_frames: Optional[int] = 49,
- height: Optional[int] = 512,
- width: Optional[int] = 512,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- timesteps: Optional[List[int]] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ prompt: str | list[str] = None,
+ num_frames: int | None = 49,
+ height: int | None = 512,
+ width: int | None = 512,
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 5.0,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ timesteps: list[int] | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
guidance_rescale: float = 0.0,
):
r"""
Generates images or video using the EasyAnimate pipeline based on the provided prompts.
Examples:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
Text prompts to guide the image or video generation. If not provided, use `prompt_embeds` instead.
num_frames (`int`, *optional*):
Length of the generated video (in frames).
@@ -566,13 +564,13 @@ def __call__(
down inference.
guidance_scale (`float`, *optional*, defaults to 5.0):
Encourages the model to align outputs with prompts. A higher value may decrease image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
Prompts indicating what to exclude in generation. If not specified, use `negative_prompt_embeds`.
num_images_per_prompt (`int`, *optional*, defaults to 1):
Number of images to generate for each prompt.
eta (`float`, *optional*, defaults to 0.0):
Applies to DDIM scheduling. Controlled by the eta parameter from the related literature.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A generator to ensure reproducibility in image generation.
latents (`torch.Tensor`, *optional*):
Predefined latent tensors to condition generation.
@@ -590,15 +588,15 @@ def __call__(
If `True`, returns a structured output. Otherwise returns a simple tuple.
callback_on_step_end (`Callable`, *optional*):
Functions called at the end of each denoising step.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
Tensor names to be included in callback function calls.
guidance_rescale (`float`, *optional*, defaults to 0.0):
Adjusts noise levels based on guidance scale.
- original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
+ original_size (`tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
Original dimensions of the output.
- target_size (`Tuple[int, int]`, *optional*):
+ target_size (`tuple[int, int]`, *optional*):
Desired output dimensions for calculations.
- crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`):
+ crops_coords_top_left (`tuple[int, int]`, *optional*, defaults to `(0, 0)`):
Coordinates for cropping.
Returns:
diff --git a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py
index b28a2c9fb273..5e07996a661c 100755
--- a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py
+++ b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py
@@ -14,7 +14,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -228,10 +228,10 @@ def resize_mask(mask, latent, process_first_frame_only=True):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -246,15 +246,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -297,9 +297,9 @@ class EasyAnimateControlPipeline(DiffusionPipeline):
Args:
vae ([`AutoencoderKLMagvit`]):
Variational Auto-Encoder (VAE) Model to encode and decode video to and from latent representations.
- text_encoder (Optional[`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel`]):
+ text_encoder (`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel` | None):
EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1.
- tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]):
+ tokenizer (`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer` | None):
A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text.
transformer ([`EasyAnimateTransformer3DModel`]):
The EasyAnimate model designed by EasyAnimate Team.
@@ -313,8 +313,8 @@ class EasyAnimateControlPipeline(DiffusionPipeline):
def __init__(
self,
vae: AutoencoderKLMagvit,
- text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel],
- tokenizer: Union[Qwen2Tokenizer, BertTokenizer],
+ text_encoder: Qwen2VLForConditionalGeneration | BertModel,
+ tokenizer: Qwen2Tokenizer | BertTokenizer,
transformer: EasyAnimateTransformer3DModel,
scheduler: FlowMatchEulerDiscreteScheduler,
):
@@ -351,23 +351,23 @@ def __init__(
# Copied from diffusers.pipelines.easyanimate.pipeline_easyanimate.EasyAnimatePipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -377,7 +377,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -671,38 +671,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- num_frames: Optional[int] = 49,
- height: Optional[int] = 512,
- width: Optional[int] = 512,
- control_video: Union[torch.FloatTensor] = None,
- control_camera_video: Union[torch.FloatTensor] = None,
- ref_image: Union[torch.FloatTensor] = None,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ prompt: str | list[str] = None,
+ num_frames: int | None = 49,
+ height: int | None = 512,
+ width: int | None = 512,
+ control_video: torch.FloatTensor = None,
+ control_camera_video: torch.FloatTensor = None,
+ ref_image: torch.FloatTensor = None,
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 5.0,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
guidance_rescale: float = 0.0,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
):
r"""
Generates images or video using the EasyAnimate pipeline based on the provided prompts.
Examples:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
Text prompts to guide the image or video generation. If not provided, use `prompt_embeds` instead.
num_frames (`int`, *optional*):
Length of the generated video (in frames).
@@ -715,13 +713,13 @@ def __call__(
down inference.
guidance_scale (`float`, *optional*, defaults to 5.0):
Encourages the model to align outputs with prompts. A higher value may decrease image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
Prompts indicating what to exclude in generation. If not specified, use `negative_prompt_embeds`.
num_images_per_prompt (`int`, *optional*, defaults to 1):
Number of images to generate for each prompt.
eta (`float`, *optional*, defaults to 0.0):
Applies to DDIM scheduling. Controlled by the eta parameter from the related literature.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A generator to ensure reproducibility in image generation.
latents (`torch.Tensor`, *optional*):
Predefined latent tensors to condition generation.
@@ -739,7 +737,7 @@ def __call__(
If `True`, returns a structured output. Otherwise returns a simple tuple.
callback_on_step_end (`Callable`, *optional*):
Functions called at the end of each denoising step.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
Tensor names to be included in callback function calls.
guidance_rescale (`float`, *optional*, defaults to 0.0):
Adjusts noise levels based on guidance scale.
diff --git a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py
index ec394315ee93..872313898008 100755
--- a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py
+++ b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py
@@ -14,7 +14,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -270,10 +270,10 @@ def add_noise_to_reference_video(image, ratio=None, generator=None):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -288,15 +288,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -339,9 +339,9 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
Args:
vae ([`AutoencoderKLMagvit`]):
Variational Auto-Encoder (VAE) Model to encode and decode video to and from latent representations.
- text_encoder (Optional[`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel`]):
+ text_encoder (`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel` | None):
EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1.
- tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]):
+ tokenizer (`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer` | None):
A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text.
transformer ([`EasyAnimateTransformer3DModel`]):
The EasyAnimate model designed by EasyAnimate Team.
@@ -355,8 +355,8 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
def __init__(
self,
vae: AutoencoderKLMagvit,
- text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel],
- tokenizer: Union[Qwen2Tokenizer, BertTokenizer],
+ text_encoder: Qwen2VLForConditionalGeneration | BertModel,
+ tokenizer: Qwen2Tokenizer | BertTokenizer,
transformer: EasyAnimateTransformer3DModel,
scheduler: FlowMatchEulerDiscreteScheduler,
):
@@ -393,23 +393,23 @@ def __init__(
# Copied from diffusers.pipelines.easyanimate.pipeline_easyanimate.EasyAnimatePipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -419,7 +419,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -789,40 +789,38 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- num_frames: Optional[int] = 49,
- video: Union[torch.FloatTensor] = None,
- mask_video: Union[torch.FloatTensor] = None,
- masked_video_latents: Union[torch.FloatTensor] = None,
- height: Optional[int] = 512,
- width: Optional[int] = 512,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ prompt: str | list[str] = None,
+ num_frames: int | None = 49,
+ video: torch.FloatTensor = None,
+ mask_video: torch.FloatTensor = None,
+ masked_video_latents: torch.FloatTensor = None,
+ height: int | None = 512,
+ width: int | None = 512,
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 5.0,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
guidance_rescale: float = 0.0,
strength: float = 1.0,
noise_aug_strength: float = 0.0563,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
):
r"""
The call function to the pipeline for generation with HunyuanDiT.
Examples:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
num_frames (`int`, *optional*):
Length of the video to be generated in seconds. This parameter influences the number of frames and
@@ -843,7 +841,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 5.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is effective when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to exclude in image generation. If not defined, you need to provide
`negative_prompt_embeds`. This parameter is ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -852,7 +850,7 @@ def __call__(
A parameter defined in the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the
[`~schedulers.DDIMScheduler`] and is ignored in other schedulers. It adjusts noise level during the
inference process.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) for setting
random seeds which helps in making generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -874,11 +872,11 @@ def __call__(
return_dict (`bool`, *optional*, defaults to `True`):
If set to `True`, a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] will be returned;
otherwise, a tuple containing the generated images and safety flags will be returned.
- callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`,
+ callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`,
*optional*):
A callback function (or a list of them) that will be executed at the end of each denoising step,
allowing for custom processing during generation.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
Specifies which tensor inputs should be included in the callback function. If not defined, all tensor
inputs will be passed, facilitating enhanced logging or monitoring of the generation process.
guidance_rescale (`float`, *optional*, defaults to 0.0):
diff --git a/src/diffusers/pipelines/easyanimate/pipeline_output.py b/src/diffusers/pipelines/easyanimate/pipeline_output.py
index c761a3b1079f..4fd904ae7dfe 100644
--- a/src/diffusers/pipelines/easyanimate/pipeline_output.py
+++ b/src/diffusers/pipelines/easyanimate/pipeline_output.py
@@ -11,8 +11,8 @@ class EasyAnimatePipelineOutput(BaseOutput):
Output class for EasyAnimate pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
diff --git a/src/diffusers/pipelines/flux/modeling_flux.py b/src/diffusers/pipelines/flux/modeling_flux.py
index d7f2f45812b3..916e3fbe5953 100644
--- a/src/diffusers/pipelines/flux/modeling_flux.py
+++ b/src/diffusers/pipelines/flux/modeling_flux.py
@@ -14,7 +14,6 @@
from dataclasses import dataclass
-from typing import Optional
import torch
import torch.nn as nn
@@ -26,7 +25,7 @@
@dataclass
class ReduxImageEncoderOutput(BaseOutput):
- image_embeds: Optional[torch.Tensor] = None
+ image_embeds: torch.Tensor | None = None
class ReduxImageEncoder(ModelMixin, ConfigMixin):
diff --git a/src/diffusers/pipelines/flux/pipeline_flux.py b/src/diffusers/pipelines/flux/pipeline_flux.py
index 9562722dbee3..be2bbe2acc6a 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -87,10 +87,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -105,15 +105,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -217,11 +217,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -266,9 +266,9 @@ def _get_t5_prompt_embeds(
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -310,21 +310,21 @@ def _get_clip_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -653,49 +653,49 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] | None = None,
true_cfg_scale: float = 1.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 3.5,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_ip_adapter_image: Optional[PipelineImageInput] = None,
- negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_ip_adapter_image: PipelineImageInput | None = None,
+ negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
true_cfg_scale (`float`, *optional*, defaults to 1.0):
@@ -708,7 +708,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -720,7 +720,7 @@ def __call__(
the [paper](https://huggingface.co/papers/2210.03142) to learn more.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -734,13 +734,13 @@ def __call__(
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
If not provided, pooled text embeddings will be generated from `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
negative_ip_adapter_image:
(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
@@ -766,7 +766,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control.py b/src/diffusers/pipelines/flux/pipeline_flux_control.py
index 77f971d57a80..84e21b59ce58 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_control.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_control.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -100,10 +100,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -118,15 +118,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -229,11 +229,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -279,9 +279,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -324,21 +324,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -636,38 +636,38 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 3.5,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -681,7 +681,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -693,7 +693,7 @@ def __call__(
the [paper](https://huggingface.co/papers/2210.03142) to learn more.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -720,7 +720,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py
index e1bbc6735051..b455c611e0ae 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -103,7 +103,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -118,10 +118,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -136,15 +136,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -239,11 +239,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -289,9 +289,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -334,21 +334,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -635,46 +635,46 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -694,7 +694,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -706,7 +706,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -733,7 +733,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py
index b02e74d3b2d6..2d1e05493a11 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -130,7 +130,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -145,10 +145,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -163,15 +163,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -279,11 +279,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -329,9 +329,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -374,21 +374,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -806,62 +806,62 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,
images must be passed as a list such that each element of the list can be correctly batched for input
to a single ControlNet.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
1)`, or `(H, W)`.
- mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`):
+ mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`):
`Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask
latents tensor will be generated by `mask_image`.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -877,7 +877,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -889,7 +889,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -916,7 +916,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py
index 78de4f617f84..d8dcdfcd4640 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -102,7 +102,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -117,10 +117,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -135,15 +135,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -214,9 +214,10 @@ def __init__(
text_encoder_2: T5EncoderModel,
tokenizer_2: T5TokenizerFast,
transformer: FluxTransformer2DModel,
- controlnet: Union[
- FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel
- ],
+ controlnet: FluxControlNetModel
+ | list[FluxControlNetModel]
+ | tuple[FluxControlNetModel]
+ | FluxMultiControlNetModel,
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
):
@@ -247,11 +248,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -296,9 +297,9 @@ def _get_t5_prompt_embeds(
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -340,21 +341,21 @@ def _get_clip_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -678,47 +679,47 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] | None = None,
true_cfg_scale: float = 1.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
control_image: PipelineImageInput = None,
- control_mode: Optional[Union[int, List[int]]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_ip_adapter_image: Optional[PipelineImageInput] = None,
- negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ control_mode: int | list[int] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_ip_adapter_image: PipelineImageInput | None = None,
+ negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -728,7 +729,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -738,27 +739,27 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,
images must be passed as a list such that each element of the list can be correctly batched for input
to a single ControlNet.
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
- control_mode (`int` or `List[int]`,, *optional*, defaults to None):
+ control_mode (`int` or `list[int]`,, *optional*, defaults to None):
The control mode when applying ControlNet-Union.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -772,13 +773,13 @@ def __call__(
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
If not provided, pooled text embeddings will be generated from `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
negative_ip_adapter_image:
(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
@@ -796,7 +797,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py
index 5bf593258f49..fdaff9b0af8a 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py
@@ -1,5 +1,5 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -97,7 +97,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -112,10 +112,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -130,15 +130,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -209,9 +209,10 @@ def __init__(
text_encoder_2: T5EncoderModel,
tokenizer_2: T5TokenizerFast,
transformer: FluxTransformer2DModel,
- controlnet: Union[
- FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel
- ],
+ controlnet: FluxControlNetModel
+ | list[FluxControlNetModel]
+ | tuple[FluxControlNetModel]
+ | FluxMultiControlNetModel,
):
super().__init__()
if isinstance(controlnet, (list, tuple)):
@@ -239,11 +240,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -289,9 +290,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -334,21 +335,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -635,43 +636,43 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- control_mode: Optional[Union[int, List[int]]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ control_mode: int | list[int] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`.
- image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
+ image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`):
The image(s) to modify with the pipeline.
- control_image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
+ control_image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`):
The ControlNet input condition. Image to control the generation.
height (`int`, *optional*, defaults to self.default_sample_size * self.vae_scale_factor):
The height in pixels of the generated image.
@@ -682,21 +683,21 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 28):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.0):
Guidance scale as defined in [Classifier-Free Diffusion
Guidance](https://huggingface.co/papers/2207.12598).
- control_mode (`int` or `List[int]`, *optional*):
+ control_mode (`int` or `list[int]`, *optional*):
The mode for the ControlNet. If multiple ControlNets are used, this should be a list.
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original transformer.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or more [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to
make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -714,7 +715,7 @@ def __call__(
Additional keyword arguments to be passed to the joint attention mechanism.
callback_on_step_end (`Callable`, *optional*):
A function that calls at the end of each denoising step during the inference.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function.
max_sequence_length (`int`, *optional*, defaults to 512):
The maximum length of the sequence to be generated.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py
index a1e1f5f5e9e5..eed671152bc9 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py
@@ -1,5 +1,5 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -99,7 +99,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -114,10 +114,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -132,15 +132,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -211,9 +211,10 @@ def __init__(
text_encoder_2: T5EncoderModel,
tokenizer_2: T5TokenizerFast,
transformer: FluxTransformer2DModel,
- controlnet: Union[
- FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel
- ],
+ controlnet: FluxControlNetModel
+ | list[FluxControlNetModel]
+ | tuple[FluxControlNetModel]
+ | FluxMultiControlNetModel,
):
super().__init__()
if isinstance(controlnet, (list, tuple)):
@@ -250,11 +251,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -300,9 +301,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -345,21 +346,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -739,51 +740,51 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.6,
- padding_mask_crop: Optional[int] = None,
- sigmas: Optional[List[float]] = None,
+ padding_mask_crop: int | None = None,
+ sigmas: list[float] | None = None,
num_inference_steps: int = 28,
guidance_scale: float = 7.0,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- control_mode: Optional[Union[int, List[int]]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ control_mode: int | list[int] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`.
- image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
+ image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`):
The image(s) to inpaint.
- mask_image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
+ mask_image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`):
The mask image(s) to use for inpainting. White pixels in the mask will be repainted, while black pixels
will be preserved.
masked_image_latents (`torch.FloatTensor`, *optional*):
Pre-generated masked image latents.
- control_image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
+ control_image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`):
The ControlNet input condition. Image to control the generation.
height (`int`, *optional*, defaults to self.default_sample_size * self.vae_scale_factor):
The height in pixels of the generated image.
@@ -796,25 +797,25 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 28):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.0):
Guidance scale as defined in [Classifier-Free Diffusion
Guidance](https://huggingface.co/papers/2207.12598).
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- control_mode (`int` or `List[int]`, *optional*):
+ control_mode (`int` or `list[int]`, *optional*):
The mode for the ControlNet. If multiple ControlNets are used, this should be a list.
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original transformer.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or more [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to
make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -832,7 +833,7 @@ def __call__(
Additional keyword arguments to be passed to the joint attention mechanism.
callback_on_step_end (`Callable`, *optional*):
A function that calls at the end of each denoising step during the inference.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function.
max_sequence_length (`int`, *optional*, defaults to 512):
The maximum length of the sequence to be generated.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_fill.py b/src/diffusers/pipelines/flux/pipeline_flux_fill.py
index 8ec9871d2579..cf929f53fc6d 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_fill.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_fill.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -94,10 +94,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -112,15 +112,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -153,7 +153,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -244,11 +244,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -294,9 +294,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -419,21 +419,21 @@ def prepare_mask_latents(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -753,52 +753,52 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- image: Optional[torch.FloatTensor] = None,
- mask_image: Optional[torch.FloatTensor] = None,
- masked_image_latents: Optional[torch.FloatTensor] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ image: torch.FloatTensor | None = None,
+ mask_image: torch.FloatTensor | None = None,
+ masked_image_latents: torch.FloatTensor | None = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 1.0,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 30.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
1)`, or `(H, W)`.
- mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`):
+ mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`):
`Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask
latents tensor will be generated by `mask_image`.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -814,7 +814,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -826,7 +826,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -853,7 +853,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py
index 5166a6497e01..cadff7736ff4 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -95,7 +95,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -110,10 +110,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -128,15 +128,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -238,11 +238,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -288,9 +288,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -333,21 +333,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -735,47 +735,47 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] | None = None,
true_cfg_scale: float = 1.0,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_ip_adapter_image: Optional[PipelineImageInput] = None,
- negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_ip_adapter_image: PipelineImageInput | None = None,
+ negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -794,7 +794,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -806,7 +806,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -820,13 +820,13 @@ def __call__(
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
If not provided, pooled text embeddings will be generated from `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
negative_ip_adapter_image:
(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
@@ -844,7 +844,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py
index 64a81fb0699f..b8ce25a4f5a9 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -91,7 +91,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -106,10 +106,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -124,15 +124,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -241,11 +241,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -291,9 +291,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -336,21 +336,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -776,63 +776,63 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] | None = None,
true_cfg_scale: float = 1.0,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_ip_adapter_image: Optional[PipelineImageInput] = None,
- negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_ip_adapter_image: PipelineImageInput | None = None,
+ negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
1)`, or `(H, W)`.
- mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`):
+ mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`):
`Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask
latents tensor will be generated by `mask_image`.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -855,7 +855,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -867,7 +867,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -881,13 +881,13 @@ def __call__(
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
If not provided, pooled text embeddings will be generated from `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
negative_ip_adapter_image:
(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
@@ -905,7 +905,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py
index 94ae460afcd0..f4bbe42ef850 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -117,10 +117,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -135,15 +135,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -176,7 +176,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -263,11 +263,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -313,9 +313,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -358,21 +358,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -668,15 +668,15 @@ def disable_vae_tiling(self):
def prepare_latents(
self,
- image: Optional[torch.Tensor],
+ image: torch.Tensor | None,
batch_size: int,
num_channels_latents: int,
height: int,
width: int,
dtype: torch.dtype,
device: torch.device,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
):
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -752,33 +752,33 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[PipelineImageInput] = None,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
+ image: PipelineImageInput | None = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] | None = None,
true_cfg_scale: float = 1.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 3.5,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_ip_adapter_image: Optional[PipelineImageInput] = None,
- negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_ip_adapter_image: PipelineImageInput | None = None,
+ negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
max_area: int = 1024**2,
_auto_resize: bool = True,
@@ -787,23 +787,23 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
true_cfg_scale (`float`, *optional*, defaults to 1.0):
@@ -815,7 +815,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -827,7 +827,7 @@ def __call__(
the [paper](https://huggingface.co/papers/2210.03142) to learn more.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -842,13 +842,13 @@ def __call__(
If not provided, pooled text embeddings will be generated from `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
negative_ip_adapter_image:
(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
@@ -874,7 +874,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py
index b6f957981e14..313682dc7e33 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py
@@ -2,7 +2,7 @@
# author: @vuongminh1907
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -141,10 +141,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -159,15 +159,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -200,7 +200,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -296,11 +296,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -346,9 +346,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -391,21 +391,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -742,7 +742,7 @@ def disable_vae_tiling(self):
def prepare_latents(
self,
- image: Optional[torch.Tensor],
+ image: torch.Tensor | None,
timestep: int,
batch_size: int,
num_channels_latents: int,
@@ -750,9 +750,9 @@ def prepare_latents(
width: int,
dtype: torch.dtype,
device: torch.device,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- image_reference: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ image_reference: torch.Tensor | None = None,
):
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -941,37 +941,37 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[PipelineImageInput] = None,
- image_reference: Optional[PipelineImageInput] = None,
+ image: PipelineImageInput | None = None,
+ image_reference: PipelineImageInput | None = None,
mask_image: PipelineImageInput = None,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] | None = None,
true_cfg_scale: float = 1.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 1.0,
- padding_mask_crop: Optional[int] = None,
+ padding_mask_crop: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 3.5,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_ip_adapter_image: Optional[PipelineImageInput] = None,
- negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_ip_adapter_image: PipelineImageInput | None = None,
+ negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
max_area: int = 1024**2,
_auto_resize: bool = True,
@@ -980,37 +980,37 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be be inpainted (which parts of the image
to be masked out with `mask_image` and repainted according to `prompt` and `image_reference`). For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- image_reference (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image_reference (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point for the
masked area. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If
it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)` If it is
a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can
also accept image latents as `image`, but if passing latents directly it is not encoded again.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
1)`, or `(H, W)`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
true_cfg_scale (`float`, *optional*, defaults to 1.0):
@@ -1036,7 +1036,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1048,7 +1048,7 @@ def __call__(
the [paper](https://huggingface.co/papers/2210.03142) to learn more.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -1063,13 +1063,13 @@ def __call__(
If not provided, pooled text embeddings will be generated from `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
negative_ip_adapter_image:
(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
@@ -1095,7 +1095,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py
index e79db337b2e3..330e2623b287 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py
@@ -13,8 +13,6 @@
# limitations under the License.
-from typing import List, Optional, Union
-
import torch
from PIL import Image
from transformers import (
@@ -196,11 +194,11 @@ def encode_image(self, image, device, num_images_per_prompt):
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -246,9 +244,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -291,21 +289,21 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -373,28 +371,28 @@ def encode_prompt(
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- prompt_embeds_scale: Optional[Union[float, List[float]]] = 1.0,
- pooled_prompt_embeds_scale: Optional[Union[float, List[float]]] = 1.0,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ prompt_embeds_scale: float | list[float] | None = 1.0,
+ pooled_prompt_embeds_scale: float | list[float] | None = 1.0,
return_dict: bool = True,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. **experimental feature**: to use this feature,
make sure to explicitly load text encoders to the pipeline. Prompts will be ignored if text encoders
are not loaded.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`.
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
diff --git a/src/diffusers/pipelines/flux/pipeline_output.py b/src/diffusers/pipelines/flux/pipeline_output.py
index 69e742d3e026..7f24bdbf5d60 100644
--- a/src/diffusers/pipelines/flux/pipeline_output.py
+++ b/src/diffusers/pipelines/flux/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -14,14 +13,14 @@ class FluxPipelineOutput(BaseOutput):
Output class for Flux image generation pipelines.
Args:
- images (`List[PIL.Image.Image]` or `torch.Tensor` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size,
+ images (`list[PIL.Image.Image]` or `torch.Tensor` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size,
height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion
pipeline. Torch tensors can represent either the denoised images or the intermediate latents ready to be
passed to the decoder.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
@dataclass
@@ -30,8 +29,8 @@ class FluxPriorReduxPipelineOutput(BaseOutput):
Output class for Flux Prior Redux pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
diff --git a/src/diffusers/pipelines/flux2/image_processor.py b/src/diffusers/pipelines/flux2/image_processor.py
index f1a8742491f7..e0a1b80ce533 100644
--- a/src/diffusers/pipelines/flux2/image_processor.py
+++ b/src/diffusers/pipelines/flux2/image_processor.py
@@ -13,7 +13,6 @@
# limitations under the License.
import math
-from typing import List
import PIL.Image
@@ -148,7 +147,7 @@ def _resize_and_crop(
# Taken from
# https://github.com/black-forest-labs/flux2/blob/5a5d316b1b42f6b59a8c9194b77c8256be848432/src/flux2/sampling.py#L310C1-L339C19
@staticmethod
- def concatenate_images(images: List[PIL.Image.Image]) -> PIL.Image.Image:
+ def concatenate_images(images: list[PIL.Image.Image]) -> PIL.Image.Image:
"""
Concatenate a list of PIL images horizontally with center alignment and white background.
"""
diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2.py b/src/diffusers/pipelines/flux2/pipeline_flux2.py
index c01b7137e086..6cd0563fcc19 100644
--- a/src/diffusers/pipelines/flux2/pipeline_flux2.py
+++ b/src/diffusers/pipelines/flux2/pipeline_flux2.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -63,9 +63,9 @@
# Adapted from
# https://github.com/black-forest-labs/flux2/blob/5a5d316b1b42f6b59a8c9194b77c8256be848432/src/flux2/text_encoder.py#L68
def format_input(
- prompts: List[str],
+ prompts: list[str],
system_message: str = SYSTEM_MESSAGE,
- images: Optional[Union[List[PIL.Image.Image], List[List[PIL.Image.Image]]]] = None,
+ images: list[PIL.Image.Image, list[list[PIL.Image.Image]]] | None = None,
):
"""
Format a batch of text prompts into the conversation format expected by apply_chat_template. Optionally, add images
@@ -130,10 +130,10 @@ def format_input(
# Adapted from
# https://github.com/black-forest-labs/flux2/blob/5a5d316b1b42f6b59a8c9194b77c8256be848432/src/flux2/text_encoder.py#L49C5-L66C19
def _validate_and_process_images(
- images: List[List[PIL.Image.Image]] | List[PIL.Image.Image],
+ images: list[list[PIL.Image.Image]] | list[PIL.Image.Image],
image_processor: Flux2ImageProcessor,
upsampling_max_image_size: int,
-) -> List[List[PIL.Image.Image]]:
+) -> list[list[PIL.Image.Image]]:
# Simple validation: ensure it's a list of PIL images or list of lists of PIL images
if not images:
return []
@@ -177,10 +177,10 @@ def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -195,15 +195,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -236,7 +236,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -304,12 +304,12 @@ def __init__(
def _get_mistral_3_small_prompt_embeds(
text_encoder: Mistral3ForConditionalGeneration,
tokenizer: AutoProcessor,
- prompt: Union[str, List[str]],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
max_sequence_length: int = 512,
system_message: str = SYSTEM_MESSAGE,
- hidden_states_layers: List[int] = (10, 20, 30),
+ hidden_states_layers: list[int] = (10, 20, 30),
):
dtype = text_encoder.dtype if dtype is None else dtype
device = text_encoder.device if device is None else device
@@ -355,7 +355,7 @@ def _get_mistral_3_small_prompt_embeds(
@staticmethod
def _prepare_text_ids(
x: torch.Tensor, # (B, L, D) or (L, D)
- t_coord: Optional[torch.Tensor] = None,
+ t_coord: torch.Tensor | None = None,
):
B, L, _ = x.shape
out_ids = []
@@ -405,7 +405,7 @@ def _prepare_latent_ids(
@staticmethod
def _prepare_image_ids(
- image_latents: List[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...]
+ image_latents: list[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...]
scale: int = 10,
):
r"""
@@ -415,7 +415,7 @@ def _prepare_image_ids(
dimensions.
Args:
- image_latents (List[torch.Tensor]):
+ image_latents (list[torch.Tensor]):
A list of image latent feature tensors, typically of shape (C, H, W).
scale (int, optional):
A factor used to define the time separation (T-coordinate) between latents. T-coordinate for the i-th
@@ -508,11 +508,11 @@ def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> list[torch
def upsample_prompt(
self,
- prompt: Union[str, List[str]],
- images: Union[List[PIL.Image.Image], List[List[PIL.Image.Image]]] = None,
+ prompt: str | list[str],
+ images: list[PIL.Image.Image, list[list[PIL.Image.Image]]] = None,
temperature: float = 0.15,
device: torch.device = None,
- ) -> List[str]:
+ ) -> list[str]:
prompt = [prompt] if isinstance(prompt, str) else prompt
device = self.text_encoder.device if device is None else device
@@ -570,12 +570,12 @@ def upsample_prompt(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 512,
- text_encoder_out_layers: Tuple[int] = (10, 20, 30),
+ text_encoder_out_layers: tuple[int] = (10, 20, 30),
):
device = device or self._execution_device
@@ -625,7 +625,7 @@ def prepare_latents(
dtype,
device,
generator: torch.Generator,
- latents: Optional[torch.Tensor] = None,
+ latents: torch.Tensor | None = None,
):
# VAE applies 8x compression on images but we must also account for packing which requires
# latent height and width to be divisible by 2.
@@ -651,7 +651,7 @@ def prepare_latents(
def prepare_image_latents(
self,
- images: List[torch.Tensor],
+ images: list[torch.Tensor],
batch_size,
generator: torch.Generator,
device,
@@ -744,37 +744,37 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[Union[List[PIL.Image.Image], PIL.Image.Image]] = None,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ image: list[PIL.Image.Image, PIL.Image.Image] | None = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = 4.0,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = 4.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
- text_encoder_out_layers: Tuple[int] = (10, 20, 30),
+ text_encoder_out_layers: tuple[int] = (10, 20, 30),
caption_upsample_temperature: float = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
guidance_scale (`float`, *optional*, defaults to 1.0):
@@ -790,13 +790,13 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -825,7 +825,7 @@ def __call__(
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
- text_encoder_out_layers (`Tuple[int]`):
+ text_encoder_out_layers (`tuple[int]`):
Layer indices to use in the `text_encoder` to derive the final prompt embeddings.
caption_upsample_temperature (`float`):
When specified, we will try to perform caption upsampling for potentially improved outputs. We
diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py
index efb0aebf8593..936d2c3804ab 100644
--- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py
+++ b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -81,10 +81,10 @@ def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -99,15 +99,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -140,7 +140,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -208,11 +208,11 @@ def __init__(
def _get_qwen3_prompt_embeds(
text_encoder: Qwen3ForCausalLM,
tokenizer: Qwen2TokenizerFast,
- prompt: Union[str, List[str]],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
max_sequence_length: int = 512,
- hidden_states_layers: List[int] = (9, 18, 27),
+ hidden_states_layers: list[int] = (9, 18, 27),
):
dtype = text_encoder.dtype if dtype is None else dtype
device = text_encoder.device if device is None else device
@@ -265,7 +265,7 @@ def _get_qwen3_prompt_embeds(
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_text_ids
def _prepare_text_ids(
x: torch.Tensor, # (B, L, D) or (L, D)
- t_coord: Optional[torch.Tensor] = None,
+ t_coord: torch.Tensor | None = None,
):
B, L, _ = x.shape
out_ids = []
@@ -317,7 +317,7 @@ def _prepare_latent_ids(
@staticmethod
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_image_ids
def _prepare_image_ids(
- image_latents: List[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...]
+ image_latents: list[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...]
scale: int = 10,
):
r"""
@@ -327,7 +327,7 @@ def _prepare_image_ids(
dimensions.
Args:
- image_latents (List[torch.Tensor]):
+ image_latents (list[torch.Tensor]):
A list of image latent feature tensors, typically of shape (C, H, W).
scale (int, optional):
A factor used to define the time separation (T-coordinate) between latents. T-coordinate for the i-th
@@ -424,12 +424,12 @@ def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> list[torch
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 512,
- text_encoder_out_layers: Tuple[int] = (9, 18, 27),
+ text_encoder_out_layers: tuple[int] = (9, 18, 27),
):
device = device or self._execution_device
@@ -480,7 +480,7 @@ def prepare_latents(
dtype,
device,
generator: torch.Generator,
- latents: Optional[torch.Tensor] = None,
+ latents: torch.Tensor | None = None,
):
# VAE applies 8x compression on images but we must also account for packing which requires
# latent height and width to be divisible by 2.
@@ -507,7 +507,7 @@ def prepare_latents(
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline.prepare_image_latents
def prepare_image_latents(
self,
- images: List[torch.Tensor],
+ images: list[torch.Tensor],
batch_size,
generator: torch.Generator,
device,
@@ -608,25 +608,25 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[Union[List[PIL.Image.Image], PIL.Image.Image]] = None,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ image: list[PIL.Image.Image] | PIL.Image.Image | None = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = 4.0,
+ sigmas: list[float] | None = None,
+ guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[Union[str, List[str]]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: str | list[str] | None = None,
+ output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int, dict], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
- text_encoder_out_layers: Tuple[int] = (9, 18, 27),
+ text_encoder_out_layers: tuple[int] = (9, 18, 27),
):
r"""
Function invoked when calling the pipeline for generation.
@@ -693,7 +693,7 @@ def __call__(
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
- text_encoder_out_layers (`Tuple[int]`):
+ text_encoder_out_layers (`tuple[int]`):
Layer indices to use in the `text_encoder` to derive the final prompt embeddings.
Examples:
diff --git a/src/diffusers/pipelines/flux2/pipeline_output.py b/src/diffusers/pipelines/flux2/pipeline_output.py
index 58e8ad49c210..34ae9b574453 100644
--- a/src/diffusers/pipelines/flux2/pipeline_output.py
+++ b/src/diffusers/pipelines/flux2/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,11 +12,11 @@ class Flux2PipelineOutput(BaseOutput):
Output class for Flux2 image generation pipelines.
Args:
- images (`List[PIL.Image.Image]` or `torch.Tensor` or `np.ndarray`)
+ images (`list[PIL.Image.Image]` or `torch.Tensor` or `np.ndarray`)
List of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size,
height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion
pipeline. Torch tensors can represent either the denoised images or the intermediate latents ready to be
passed to the decoder.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image, np.ndarray]
diff --git a/src/diffusers/pipelines/free_init_utils.py b/src/diffusers/pipelines/free_init_utils.py
index 4495c5ea2683..04a385edd819 100644
--- a/src/diffusers/pipelines/free_init_utils.py
+++ b/src/diffusers/pipelines/free_init_utils.py
@@ -13,7 +13,6 @@
# limitations under the License.
import math
-from typing import Tuple, Union
import torch
import torch.fft as fft
@@ -73,8 +72,8 @@ def free_init_enabled(self):
def _get_free_init_freq_filter(
self,
- shape: Tuple[int, ...],
- device: Union[str, torch.dtype],
+ shape: tuple[int, ...],
+ device: str | torch.dtype,
filter_type: str,
order: float,
spatial_stop_frequency: float,
diff --git a/src/diffusers/pipelines/free_noise_utils.py b/src/diffusers/pipelines/free_noise_utils.py
index 2910afaf237b..5990e680ba07 100644
--- a/src/diffusers/pipelines/free_noise_utils.py
+++ b/src/diffusers/pipelines/free_noise_utils.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import torch
import torch.nn as nn
@@ -48,7 +48,7 @@ class SplitInferenceModule(nn.Module):
The size of each chunk after splitting the input tensor.
split_dim (`int`, defaults to `0`):
The dimension along which the input tensors are split.
- input_kwargs_to_split (`List[str]`, defaults to `["hidden_states"]`):
+ input_kwargs_to_split (`list[str]`, defaults to `["hidden_states"]`):
A list of keyword arguments (strings) that represent the input tensors to be split.
Workflow:
@@ -80,7 +80,7 @@ def __init__(
module: nn.Module,
split_size: int = 1,
split_dim: int = 0,
- input_kwargs_to_split: List[str] = ["hidden_states"],
+ input_kwargs_to_split: list[str] = ["hidden_states"],
) -> None:
super().__init__()
@@ -89,7 +89,7 @@ def __init__(
self.split_dim = split_dim
self.input_kwargs_to_split = set(input_kwargs_to_split)
- def forward(self, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor]]:
+ def forward(self, *args, **kwargs) -> torch.Tensor | tuple[torch.Tensor]:
r"""Forward method for the `SplitInferenceModule`.
This method processes the input by splitting specified keyword arguments along a given dimension, running the
@@ -99,13 +99,13 @@ def forward(self, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor]]:
Args:
*args (`Any`):
Positional arguments that are passed directly to the `module` without modification.
- **kwargs (`Dict[str, torch.Tensor]`):
+ **kwargs (`dict[str, torch.Tensor]`):
Keyword arguments passed to the underlying `module`. Only keyword arguments whose names match the
entries in `input_kwargs_to_split` and are of type `torch.Tensor` will be split. The remaining keyword
arguments are passed unchanged.
Returns:
- `Union[torch.Tensor, Tuple[torch.Tensor]]`:
+ `torch.Tensor | tuple[torch.Tensor]`:
The outputs obtained from `SplitInferenceModule` are the same as if the underlying module was inferred
without it.
- If the underlying module returns a single tensor, the result will be a single concatenated tensor
@@ -145,7 +145,7 @@ def forward(self, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor]]:
class AnimateDiffFreeNoiseMixin:
r"""Mixin class for [FreeNoise](https://huggingface.co/papers/2310.15169)."""
- def _enable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, DownBlockMotion, UpBlockMotion]):
+ def _enable_free_noise_in_block(self, block: CrossAttnDownBlockMotion | DownBlockMotion | UpBlockMotion):
r"""Helper function to enable FreeNoise in transformer blocks."""
for motion_module in block.motion_modules:
@@ -186,7 +186,7 @@ def _enable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, Dow
basic_transfomer_block._chunk_size, basic_transfomer_block._chunk_dim
)
- def _disable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, DownBlockMotion, UpBlockMotion]):
+ def _disable_free_noise_in_block(self, block: CrossAttnDownBlockMotion | DownBlockMotion | UpBlockMotion):
r"""Helper function to disable FreeNoise in transformer blocks."""
for motion_module in block.motion_modules:
@@ -255,16 +255,16 @@ def _check_inputs_free_noise(
def _encode_prompt_free_noise(
self,
- prompt: Union[str, Dict[int, str]],
+ prompt: str | dict[int, str],
num_frames: int,
device: torch.device,
num_videos_per_prompt: int,
do_classifier_free_guidance: bool,
- negative_prompt: Optional[Union[str, Dict[int, str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | dict[int, str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
) -> torch.Tensor:
if negative_prompt is None:
negative_prompt = ""
@@ -362,8 +362,8 @@ def _prepare_latents_free_noise(
width: int,
dtype: torch.dtype,
device: torch.device,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
):
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -443,13 +443,14 @@ def _lerp(
def enable_free_noise(
self,
- context_length: Optional[int] = 16,
+ context_length: int | None = 16,
context_stride: int = 4,
weighting_scheme: str = "pyramid",
noise_type: str = "shuffle_context",
- prompt_interpolation_callback: Optional[
- Callable[[DiffusionPipeline, int, int, torch.Tensor, torch.Tensor], torch.Tensor]
- ] = None,
+ prompt_interpolation_callback: Callable[
+ [DiffusionPipeline, int, int, torch.Tensor, torch.Tensor], torch.Tensor
+ ]
+ | None = None,
) -> None:
r"""
Enable long video generation using FreeNoise.
@@ -529,7 +530,7 @@ def disable_free_noise(self) -> None:
self._disable_free_noise_in_block(block)
def _enable_split_inference_motion_modules_(
- self, motion_modules: List[AnimateDiffTransformer3D], spatial_split_size: int
+ self, motion_modules: list[AnimateDiffTransformer3D], spatial_split_size: int
) -> None:
for motion_module in motion_modules:
motion_module.proj_in = SplitInferenceModule(motion_module.proj_in, spatial_split_size, 0, ["input"])
@@ -545,19 +546,19 @@ def _enable_split_inference_motion_modules_(
motion_module.proj_out = SplitInferenceModule(motion_module.proj_out, spatial_split_size, 0, ["input"])
def _enable_split_inference_attentions_(
- self, attentions: List[Transformer2DModel], temporal_split_size: int
+ self, attentions: list[Transformer2DModel], temporal_split_size: int
) -> None:
for i in range(len(attentions)):
attentions[i] = SplitInferenceModule(
attentions[i], temporal_split_size, 0, ["hidden_states", "encoder_hidden_states"]
)
- def _enable_split_inference_resnets_(self, resnets: List[ResnetBlock2D], temporal_split_size: int) -> None:
+ def _enable_split_inference_resnets_(self, resnets: list[ResnetBlock2D], temporal_split_size: int) -> None:
for i in range(len(resnets)):
resnets[i] = SplitInferenceModule(resnets[i], temporal_split_size, 0, ["input_tensor", "temb"])
def _enable_split_inference_samplers_(
- self, samplers: Union[List[Downsample2D], List[Upsample2D]], temporal_split_size: int
+ self, samplers: list[Downsample2D] | list[Upsample2D], temporal_split_size: int
) -> None:
for i in range(len(samplers)):
samplers[i] = SplitInferenceModule(samplers[i], temporal_split_size, 0, ["hidden_states"])
diff --git a/src/diffusers/pipelines/glm_image/pipeline_glm_image.py b/src/diffusers/pipelines/glm_image/pipeline_glm_image.py
index 589b3be47b2c..859b371b2514 100644
--- a/src/diffusers/pipelines/glm_image/pipeline_glm_image.py
+++ b/src/diffusers/pipelines/glm_image/pipeline_glm_image.py
@@ -15,7 +15,7 @@
import inspect
import re
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -79,10 +79,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.cogview4.pipeline_cogview4.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -97,15 +97,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
@@ -146,7 +146,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -262,9 +262,9 @@ def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) ->
@staticmethod
def _validate_and_normalize_images(
- image: Optional[Union[List[PIL.Image.Image], List[List[PIL.Image.Image]]]],
+ image: list[PIL.Image.Image] | list[list[PIL.Image.Image]],
batch_size: int,
- ) -> Optional[List[List[PIL.Image.Image]]]:
+ ) -> list[list[PIL.Image.Image]]:
"""
Validate and normalize image inputs to List[List[PIL.Image]].
@@ -320,12 +320,12 @@ def _validate_and_normalize_images(
def generate_prior_tokens(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
height: int,
width: int,
- image: Optional[List[List[PIL.Image.Image]]] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
+ image: list[list[PIL.Image.Image]] | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
):
"""
Generate prior tokens for the DiT model using the AR model.
@@ -490,10 +490,10 @@ def get_glyph_texts(self, prompt):
def _get_glyph_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 2048,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
"""Get glyph embeddings for each prompt in the batch."""
device = device or self._execution_device
@@ -544,20 +544,20 @@ def _get_glyph_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 2048,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
Whether to use classifier free guidance or not.
@@ -658,12 +658,7 @@ def check_inputs(
)
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
- if prompt is not None and prior_token_ids is not None:
- raise ValueError(
- f"Cannot forward both `prompt`: {prompt} and `prior_token_ids`: {prior_token_ids}. Please make sure to"
- " only forward one of the two."
- )
- elif prompt is None and prior_token_ids is None:
+ if prompt is None and prior_token_ids is None:
raise ValueError(
"Provide either `prompt` or `prior_token_ids`. Cannot leave both `prompt` and `prior_token_ids` undefined."
)
@@ -694,8 +689,8 @@ def check_inputs(
"for i2i mode, as the images are needed for VAE encoding to build the KV cache."
)
- if prior_token_ids is not None and prompt_embeds is None:
- raise ValueError("`prompt_embeds` must also be provided with `prior_token_ids`.")
+ if prior_token_ids is not None and prompt_embeds is None and prompt is None:
+ raise ValueError("`prompt_embeds` or `prompt` must also be provided with `prior_token_ids`.")
@property
def guidance_scale(self):
@@ -725,41 +720,44 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- image: Optional[
- Union[
- torch.Tensor, PIL.Image.Image, np.ndarray, List[torch.Tensor], List[PIL.Image.Image], List[np.ndarray]
- ]
- ] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ image: torch.Tensor
+ | PIL.Image.Image
+ | np.ndarray
+ | list[torch.Tensor]
+ | list[PIL.Image.Image]
+ | list[np.ndarray]
+ | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 1.5,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prior_token_ids: Optional[torch.FloatTensor] = None,
- prior_token_image_ids: Optional[List[torch.Tensor]] = None,
- source_image_grid_thw: Optional[List[torch.Tensor]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prior_token_ids: torch.Tensor | None = None,
+ prior_token_image_ids: list[torch.Tensor] | None = None,
+ source_image_grid_thw: list[torch.Tensor] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int, dict], None]
+ | PipelineCallback
+ | MultiPipelineCallbacks
+ | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 2048,
- ) -> Union[GlmImagePipelineOutput, Tuple]:
+ ) -> GlmImagePipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. Must contain shape info in the format 'H
W' where H and W are token dimensions (d32). Example: "A beautiful sunset36 24"
generates a 1152x768 image.
diff --git a/src/diffusers/pipelines/glm_image/pipeline_output.py b/src/diffusers/pipelines/glm_image/pipeline_output.py
index aec5a5454ea8..d4fd061335d4 100644
--- a/src/diffusers/pipelines/glm_image/pipeline_output.py
+++ b/src/diffusers/pipelines/glm_image/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -18,4 +17,4 @@ class GlmImagePipelineOutput(BaseOutput):
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py
index b41d9772a7cc..8e5e078cc2af 100644
--- a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py
+++ b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import (
@@ -99,10 +99,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -117,15 +117,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -201,10 +201,10 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder_3.dtype
@@ -240,10 +240,10 @@ def _get_clip_prompt_embeds(
self,
tokenizer,
text_encoder,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or text_encoder.dtype
@@ -275,10 +275,10 @@ def _get_clip_prompt_embeds(
def _get_llama3_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder_4.dtype
@@ -319,26 +319,26 @@ def _get_llama3_prompt_embeds(
def encode_prompt(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_4: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] | None = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ prompt_4: str | list[str] | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- negative_prompt_4: Optional[Union[str, List[str]]] = None,
- prompt_embeds_t5: Optional[List[torch.FloatTensor]] = None,
- prompt_embeds_llama3: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds_t5: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds_llama3: Optional[List[torch.FloatTensor]] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ negative_prompt_4: str | list[str] | None = None,
+ prompt_embeds_t5: list[torch.FloatTensor] | None = None,
+ prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 128,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
if prompt is not None:
@@ -728,33 +728,33 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_4: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ prompt_4: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- negative_prompt_4: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds_t5: Optional[torch.FloatTensor] = None,
- prompt_embeds_llama3: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds_t5: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds_llama3: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ negative_prompt_4: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds_t5: torch.FloatTensor | None = None,
+ prompt_embeds_llama3: torch.FloatTensor | None = None,
+ negative_prompt_embeds_t5: torch.FloatTensor | None = None,
+ negative_prompt_embeds_llama3: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 128,
**kwargs,
):
@@ -762,16 +762,16 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead.
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
will be used instead.
- prompt_4 (`str` or `List[str]`, *optional*):
+ prompt_4 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_4` and `text_encoder_4`. If not defined, `prompt` is
will be used instead.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -781,7 +781,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -791,22 +791,22 @@ def __call__(
Guidance-distilled models approximates true classifer-free guidance for `guidance_scale` > 1. Refer to
the [paper](https://huggingface.co/papers/2210.03142) to learn more.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_4 (`str` or `List[str]`, *optional*):
+ negative_prompt_4 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_4` and
`text_encoder_4`. If not defined, `negative_prompt` is used in all the text-encoders.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -841,7 +841,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/hidream_image/pipeline_output.py b/src/diffusers/pipelines/hidream_image/pipeline_output.py
index 66f0f1260d18..1802c7220691 100644
--- a/src/diffusers/pipelines/hidream_image/pipeline_output.py
+++ b/src/diffusers/pipelines/hidream_image/pipeline_output.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -27,9 +26,9 @@ class HiDreamImagePipelineOutput(BaseOutput):
Output class for HiDreamImage pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py
index 658935ccd886..50239e9afa22 100644
--- a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py
+++ b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py
@@ -12,9 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import inspect
import re
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -98,10 +100,10 @@ def extract_glyph_text(prompt: str):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -116,15 +118,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -193,8 +195,8 @@ def __init__(
text_encoder_2: T5EncoderModel,
tokenizer_2: ByT5Tokenizer,
transformer: HunyuanImageTransformer2DModel,
- guider: Optional[AdaptiveProjectedMixGuidance] = None,
- ocr_guider: Optional[AdaptiveProjectedMixGuidance] = None,
+ guider: AdaptiveProjectedMixGuidance | None = None,
+ ocr_guider: AdaptiveProjectedMixGuidance | None = None,
):
super().__init__()
@@ -222,9 +224,9 @@ def _get_qwen_prompt_embeds(
self,
tokenizer: Qwen2Tokenizer,
text_encoder: Qwen2_5_VLForConditionalGeneration,
- prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
tokenizer_max_length: int = 1000,
template: str = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>",
drop_idx: int = 34,
@@ -260,8 +262,8 @@ def _get_byt5_prompt_embeds(
tokenizer: ByT5Tokenizer,
text_encoder: T5EncoderModel,
prompt: str,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
tokenizer_max_length: int = 128,
):
device = device or self._execution_device
@@ -293,19 +295,19 @@ def _get_byt5_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
batch_size: int = 1,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_embeds_mask_2: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ prompt_embeds_mask_2: torch.Tensor | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -503,38 +505,38 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- distilled_guidance_scale: Optional[float] = 3.25,
- sigmas: Optional[List[float]] = None,
+ distilled_guidance_scale: float | None = 3.25,
+ sigmas: list[float] | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_embeds_mask_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask_2: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ prompt_embeds_mask_2: torch.Tensor | None = None,
+ negative_prompt_embeds_2: torch.Tensor | None = None,
+ negative_prompt_embeds_mask_2: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined and negative_prompt_embeds is
not provided, will use an empty negative prompt. Ignored when not using guidance. ).
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -544,7 +546,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -558,7 +560,7 @@ def __call__(
ignored.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py
index f38f53d9a562..93e4deb2974a 100644
--- a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py
+++ b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -62,10 +64,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -80,15 +82,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -121,7 +123,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -161,7 +163,7 @@ def __init__(
text_encoder: Qwen2_5_VLForConditionalGeneration,
tokenizer: Qwen2Tokenizer,
transformer: HunyuanImageTransformer2DModel,
- guider: Optional[AdaptiveProjectedMixGuidance] = None,
+ guider: AdaptiveProjectedMixGuidance | None = None,
):
super().__init__()
@@ -187,9 +189,9 @@ def _get_qwen_prompt_embeds(
self,
tokenizer: Qwen2Tokenizer,
text_encoder: Qwen2_5_VLForConditionalGeneration,
- prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
tokenizer_max_length: int = 1000,
template: str = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>",
drop_idx: int = 34,
@@ -222,17 +224,17 @@ def _get_qwen_prompt_embeds(
def encode_prompt(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str] | None = None,
+ device: torch.device | None = None,
batch_size: int = 1,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -435,35 +437,35 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- distilled_guidance_scale: Optional[float] = 3.25,
- image: Optional[PipelineImageInput] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ distilled_guidance_scale: float | None = 3.25,
+ image: PipelineImageInput | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 4,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, will use an empty negative
prompt. Ignored when not using guidance.
distilled_guidance_scale (`float`, *optional*, defaults to None):
@@ -482,13 +484,13 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_output.py b/src/diffusers/pipelines/hunyuan_image/pipeline_output.py
index 1e76892a0e81..169436b7d86f 100644
--- a/src/diffusers/pipelines/hunyuan_image/pipeline_output.py
+++ b/src/diffusers/pipelines/hunyuan_image/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class HunyuanImagePipelineOutput(BaseOutput):
Output class for HunyuanImage pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image, np.ndarray]
diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py
index 6bb7a4344da5..1a7cae256d63 100644
--- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py
+++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -95,10 +95,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -113,15 +113,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -154,7 +154,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -225,14 +225,14 @@ def __init__(
# Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline._get_llama_prompt_embeds
def _get_llama_prompt_embeds(
self,
- prompt: Union[str, List[str]],
- prompt_template: Dict[str, Any],
+ prompt: str | list[str],
+ prompt_template: dict[str, Any],
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
num_hidden_layers_to_skip: int = 2,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -292,10 +292,10 @@ def _get_llama_prompt_embeds(
# Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 77,
) -> torch.Tensor:
device = device or self._execution_device
@@ -332,15 +332,15 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]] = None,
- prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+ prompt: str | list[str],
+ prompt_2: str | list[str] = None,
+ prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
):
if prompt_embeds is None:
@@ -422,10 +422,10 @@ def prepare_latents(
height: int = 544,
width: int = 960,
num_frames: int = 97,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -536,51 +536,49 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
- prompt_2: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] = None,
height: int = 544,
width: int = 960,
num_frames: int = 97,
num_inference_steps: int = 50,
- sigmas: List[float] = None,
+ sigmas: list[float] = None,
true_cfg_scale: float = 6.0,
guidance_scale: float = 1.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
max_sequence_length: int = 256,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
height (`int`, defaults to `720`):
@@ -592,7 +590,7 @@ def __call__(
num_inference_steps (`int`, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -608,7 +606,7 @@ def __call__(
conditional latent is not applied.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -645,7 +643,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py
index 42ab090f1cba..3c6ec39398ef 100644
--- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py
+++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -84,10 +84,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -102,15 +102,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -198,14 +198,14 @@ def __init__(
def _get_llama_prompt_embeds(
self,
- prompt: Union[str, List[str]],
- prompt_template: Dict[str, Any],
+ prompt: str | list[str],
+ prompt_template: dict[str, Any],
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
num_hidden_layers_to_skip: int = 2,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -264,10 +264,10 @@ def _get_llama_prompt_embeds(
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 77,
) -> torch.Tensor:
device = device or self._execution_device
@@ -303,15 +303,15 @@ def _get_clip_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]] = None,
- prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+ prompt: str | list[str],
+ prompt_2: str | list[str] = None,
+ prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
):
if prompt_embeds is None:
@@ -391,10 +391,10 @@ def prepare_latents(
height: int = 720,
width: int = 1280,
num_frames: int = 129,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -492,51 +492,49 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] = None,
height: int = 720,
width: int = 1280,
num_frames: int = 129,
num_inference_steps: int = 50,
- sigmas: List[float] = None,
+ sigmas: list[float] = None,
true_cfg_scale: float = 1.0,
guidance_scale: float = 6.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
max_sequence_length: int = 256,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
height (`int`, defaults to `720`):
@@ -548,7 +546,7 @@ def __call__(
num_inference_steps (`int`, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -563,7 +561,7 @@ def __call__(
the [paper](https://huggingface.co/papers/2210.03142) to learn more.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -600,7 +598,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py
index 8006514f47ea..f82f26eea5b9 100644
--- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py
+++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py
@@ -15,7 +15,7 @@
import inspect
import math
from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -178,10 +178,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -196,15 +196,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -302,14 +302,14 @@ def __init__(
# Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline._get_llama_prompt_embeds
def _get_llama_prompt_embeds(
self,
- prompt: Union[str, List[str]],
- prompt_template: Dict[str, Any],
+ prompt: str | list[str],
+ prompt_template: dict[str, Any],
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
num_hidden_layers_to_skip: int = 2,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -369,10 +369,10 @@ def _get_llama_prompt_embeds(
# Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 77,
) -> torch.Tensor:
device = device or self._execution_device
@@ -409,15 +409,15 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]] = None,
- prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+ prompt: str | list[str],
+ prompt_2: str | list[str] = None,
+ prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
):
if prompt_embeds is None:
@@ -443,9 +443,7 @@ def encode_prompt(
return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask
- def encode_image(
- self, image: torch.Tensor, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None
- ):
+ def encode_image(self, image: torch.Tensor, device: torch.device | None = None, dtype: torch.dtype | None = None):
device = device or self._execution_device
image = (image + 1) / 2.0 # [-1, 1] -> [0, 1]
image = self.feature_extractor(images=image, return_tensors="pt", do_rescale=False).to(
@@ -528,10 +526,10 @@ def prepare_latents(
height: int = 720,
width: int = 1280,
num_frames: int = 129,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -553,10 +551,10 @@ def prepare_latents(
def prepare_image_latents(
self,
image: torch.Tensor,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
device = device or self._execution_device
if latents is None:
@@ -643,37 +641,35 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput,
- last_image: Optional[PipelineImageInput] = None,
- prompt: Union[str, List[str]] = None,
- prompt_2: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Union[str, List[str]] = None,
+ last_image: PipelineImageInput | None = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] = None,
height: int = 720,
width: int = 1280,
num_frames: int = 129,
latent_window_size: int = 9,
num_inference_steps: int = 50,
- sigmas: List[float] = None,
+ sigmas: list[float] = None,
true_cfg_scale: float = 1.0,
guidance_scale: float = 6.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- image_latents: Optional[torch.Tensor] = None,
- last_image_latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ image_latents: torch.Tensor | None = None,
+ last_image_latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
max_sequence_length: int = 256,
sampling_type: FramepackSamplingType = FramepackSamplingType.INVERTED_ANTI_DRIFTING,
):
@@ -686,17 +682,17 @@ def __call__(
last_image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`, *optional*):
The optional last image to be used as the ending point for the video generation. This is useful for
generating transitions between two images.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
height (`int`, defaults to `720`):
@@ -708,7 +704,7 @@ def __call__(
num_inference_steps (`int`, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -724,7 +720,7 @@ def __call__(
conditional latent is not applied.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
image_latents (`torch.Tensor`, *optional*):
@@ -761,7 +757,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py
index aa04e6509730..c599488c2379 100644
--- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py
+++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -147,10 +147,10 @@ def _expand_input_ids_with_image_tokens(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -165,15 +165,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -206,7 +206,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -279,15 +279,15 @@ def __init__(
def _get_llama_prompt_embeds(
self,
image: torch.Tensor,
- prompt: Union[str, List[str]],
- prompt_template: Dict[str, Any],
+ prompt: str | list[str],
+ prompt_template: dict[str, Any],
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
num_hidden_layers_to_skip: int = 2,
image_embed_interleave: int = 2,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -417,10 +417,10 @@ def _get_llama_prompt_embeds(
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 77,
) -> torch.Tensor:
device = device or self._execution_device
@@ -451,18 +451,18 @@ def _get_clip_prompt_embeds(
def encode_prompt(
self,
image: torch.Tensor,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]] = None,
- prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+ prompt: str | list[str],
+ prompt_2: str | list[str] = None,
+ prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
max_sequence_length: int = 256,
image_embed_interleave: int = 2,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
if prompt_embeds is None:
prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds(
image,
@@ -552,10 +552,10 @@ def prepare_latents(
height: int = 720,
width: int = 1280,
num_frames: int = 129,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
image_condition_type: str = "latent_concat",
) -> torch.Tensor:
if isinstance(generator, list) and len(generator) != batch_size:
@@ -671,52 +671,50 @@ def interrupt(self):
def __call__(
self,
image: PIL.Image.Image,
- prompt: Union[str, List[str]] = None,
- prompt_2: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- negative_prompt_2: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ negative_prompt_2: str | list[str] = None,
height: int = 720,
width: int = 1280,
num_frames: int = 129,
num_inference_steps: int = 50,
- sigmas: List[float] = None,
+ sigmas: list[float] = None,
true_cfg_scale: float = 1.0,
guidance_scale: float = 1.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
max_sequence_length: int = 256,
- image_embed_interleave: Optional[int] = None,
+ image_embed_interleave: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
height (`int`, defaults to `720`):
@@ -728,7 +726,7 @@ def __call__(
num_inference_steps (`int`, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -744,7 +742,7 @@ def __call__(
conditional latent is not applied.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -781,7 +779,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_output.py b/src/diffusers/pipelines/hunyuan_video/pipeline_output.py
index fae0370a53b7..4cf966262770 100644
--- a/src/diffusers/pipelines/hunyuan_video/pipeline_output.py
+++ b/src/diffusers/pipelines/hunyuan_video/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -14,8 +13,8 @@ class HunyuanVideoPipelineOutput(BaseOutput):
Output class for HunyuanVideo pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
@@ -29,11 +28,11 @@ class HunyuanVideoFramepackPipelineOutput(BaseOutput):
Output class for HunyuanVideo pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`. Or, a list of torch tensors where each tensor
corresponds to a latent that decodes to multiple frames.
"""
- frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]], List[torch.Tensor]]
+ frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] | list[torch.Tensor]
diff --git a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5.py b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5.py
index 00a703939004..a0adff493ac0 100644
--- a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5.py
+++ b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5.py
@@ -14,7 +14,7 @@
import inspect
import re
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -61,16 +61,16 @@
"""
-def format_text_input(prompt: List[str], system_message: str) -> List[Dict[str, Any]]:
+def format_text_input(prompt: list[str], system_message: str) -> list[dict[str, Any]]:
"""
Apply text to template.
Args:
- prompt (List[str]): Input text.
+ prompt (list[str]): Input text.
system_message (str): System message.
Returns:
- List[Dict[str, Any]]: List of chat conversation.
+ list[dict[str, Any]]: List of chat conversation.
"""
template = [
@@ -80,7 +80,7 @@ def format_text_input(prompt: List[str], system_message: str) -> List[Dict[str,
return template
-def extract_glyph_texts(prompt: str) -> List[str]:
+def extract_glyph_texts(prompt: str) -> list[str]:
"""
Extract glyph texts from prompt using regex pattern.
@@ -106,10 +106,10 @@ def extract_glyph_texts(prompt: str) -> List[str]:
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -124,15 +124,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -241,7 +241,7 @@ def __init__(
def _get_mllm_prompt_embeds(
text_encoder: Qwen2_5_VLTextModel,
tokenizer: Qwen2Tokenizer,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
tokenizer_max_length: int = 1000,
num_hidden_layers_to_skip: int = 2,
@@ -254,7 +254,7 @@ def _get_mllm_prompt_embeds(
5. camera angles, movements, and transitions used in the video.",
# fmt: on
crop_start: int = 108,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
prompt = [prompt] if isinstance(prompt, str) else prompt
prompt = format_text_input(prompt, system_message)
@@ -289,7 +289,7 @@ def _get_mllm_prompt_embeds(
def _get_byt5_prompt_embeds(
tokenizer: ByT5Tokenizer,
text_encoder: T5EncoderModel,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
tokenizer_max_length: int = 256,
):
@@ -333,20 +333,20 @@ def _get_byt5_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
batch_size: int = 1,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_embeds_mask_2: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ prompt_embeds_mask_2: torch.Tensor | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -481,10 +481,10 @@ def prepare_latents(
height: int = 720,
width: int = 1280,
num_frames: int = 129,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -505,7 +505,7 @@ def prepare_latents(
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
return latents
- def prepare_cond_latents_and_mask(self, latents, dtype: Optional[torch.dtype], device: Optional[torch.device]):
+ def prepare_cond_latents_and_mask(self, latents, dtype: torch.dtype | None, device: torch.device | None):
"""
Prepare conditional latents and mask for t2v generation.
@@ -543,36 +543,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_frames: int = 121,
num_inference_steps: int = 50,
- sigmas: List[float] = None,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_embeds_mask_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask_2: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ sigmas: list[float] = None,
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ prompt_embeds_mask_2: torch.Tensor | None = None,
+ negative_prompt_embeds_2: torch.Tensor | None = None,
+ negative_prompt_embeds_mask_2: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead.
height (`int`, *optional*):
@@ -584,13 +584,13 @@ def __call__(
num_inference_steps (`int`, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality video at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5_image2video.py b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5_image2video.py
index 8c555eabba11..791dec073524 100644
--- a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5_image2video.py
+++ b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5_image2video.py
@@ -14,7 +14,7 @@
import inspect
import re
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import PIL
@@ -73,16 +73,16 @@
# Copied from diffusers.pipelines.hunyuan_video1_5.pipeline_hunyuan_video1_5.format_text_input
-def format_text_input(prompt: List[str], system_message: str) -> List[Dict[str, Any]]:
+def format_text_input(prompt: list[str], system_message: str) -> list[dict[str, Any]]:
"""
Apply text to template.
Args:
- prompt (List[str]): Input text.
+ prompt (list[str]): Input text.
system_message (str): System message.
Returns:
- List[Dict[str, Any]]: List of chat conversation.
+ list[dict[str, Any]]: List of chat conversation.
"""
template = [
@@ -93,7 +93,7 @@ def format_text_input(prompt: List[str], system_message: str) -> List[Dict[str,
# Copied from diffusers.pipelines.hunyuan_video1_5.pipeline_hunyuan_video1_5.extract_glyph_texts
-def extract_glyph_texts(prompt: str) -> List[str]:
+def extract_glyph_texts(prompt: str) -> list[str]:
"""
Extract glyph texts from prompt using regex pattern.
@@ -118,7 +118,7 @@ def extract_glyph_texts(prompt: str) -> List[str]:
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -133,10 +133,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -151,15 +151,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -280,7 +280,7 @@ def __init__(
def _get_mllm_prompt_embeds(
text_encoder: Qwen2_5_VLTextModel,
tokenizer: Qwen2Tokenizer,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
tokenizer_max_length: int = 1000,
num_hidden_layers_to_skip: int = 2,
@@ -293,7 +293,7 @@ def _get_mllm_prompt_embeds(
5. camera angles, movements, and transitions used in the video.",
# fmt: on
crop_start: int = 108,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
prompt = [prompt] if isinstance(prompt, str) else prompt
prompt = format_text_input(prompt, system_message)
@@ -329,7 +329,7 @@ def _get_mllm_prompt_embeds(
def _get_byt5_prompt_embeds(
tokenizer: ByT5Tokenizer,
text_encoder: T5EncoderModel,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
tokenizer_max_length: int = 256,
):
@@ -421,20 +421,20 @@ def encode_image(
# Copied from diffusers.pipelines.hunyuan_video1_5.pipeline_hunyuan_video1_5.HunyuanVideo15Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
batch_size: int = 1,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_embeds_mask_2: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ prompt_embeds_mask_2: torch.Tensor | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -567,10 +567,10 @@ def prepare_latents(
height: int = 720,
width: int = 1280,
num_frames: int = 129,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -652,25 +652,25 @@ def interrupt(self):
def __call__(
self,
image: PIL.Image.Image,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
num_frames: int = 121,
num_inference_steps: int = 50,
- sigmas: List[float] = None,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_embeds_mask_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask_2: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ sigmas: list[float] = None,
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ prompt_embeds_mask_2: torch.Tensor | None = None,
+ negative_prompt_embeds_2: torch.Tensor | None = None,
+ negative_prompt_embeds_mask_2: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
+ attention_kwargs: dict[str, Any] | None = None,
):
r"""
The call function to the pipeline for generation.
@@ -678,10 +678,10 @@ def __call__(
Args:
image (`PIL.Image.Image`):
The input image to condition video generation on.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the video generation. If not defined, one has to pass
`negative_prompt_embeds` instead.
num_frames (`int`, defaults to `121`):
@@ -689,13 +689,13 @@ def __call__(
num_inference_steps (`int`, defaults to `50`):
The number of denoising steps. More denoising steps usually lead to a higher quality video at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_output.py b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_output.py
index 441164db5a09..ab8767299e55 100644
--- a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_output.py
+++ b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_output.py
@@ -11,7 +11,7 @@ class HunyuanVideo15PipelineOutput(BaseOutput):
Output class for HunyuanVideo1.5 pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
diff --git a/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py b/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
index 052c7b473915..b908dd5dfe83 100644
--- a/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
+++ b/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import numpy as np
import torch
@@ -160,10 +160,10 @@ class HunyuanDiTPipeline(DiffusionPipeline):
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. We use
`sdxl-vae-fp16-fix`.
- text_encoder (Optional[`~transformers.BertModel`, `~transformers.CLIPTextModel`]):
+ text_encoder (`~transformers.BertModel`, `~transformers.CLIPTextModel` | None):
Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
HunyuanDiT uses a fine-tuned [bilingual CLIP].
- tokenizer (Optional[`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer`]):
+ tokenizer (`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer` | None):
A `BertTokenizer` or `CLIPTokenizer` to tokenize text.
transformer ([`HunyuanDiT2DModel`]):
The HunyuanDiT model designed by Tencent Hunyuan.
@@ -203,8 +203,8 @@ def __init__(
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = True,
- text_encoder_2: Optional[T5EncoderModel] = None,
- tokenizer_2: Optional[T5Tokenizer] = None,
+ text_encoder_2: T5EncoderModel | None = None,
+ tokenizer_2: T5Tokenizer | None = None,
):
super().__init__()
@@ -252,19 +252,19 @@ def encode_prompt(
dtype: torch.dtype = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- max_sequence_length: Optional[int] = None,
+ negative_prompt: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ max_sequence_length: int | None = None,
text_encoder_index: int = 0,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -274,7 +274,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -363,7 +363,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -569,41 +569,39 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- prompt_attention_mask_2: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask_2: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 5.0,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_2: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ prompt_attention_mask_2: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask_2: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
guidance_rescale: float = 0.0,
- original_size: Optional[Tuple[int, int]] = (1024, 1024),
- target_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
+ original_size: tuple[int, int] | None = (1024, 1024),
+ target_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
use_resolution_binning: bool = True,
):
r"""
The call function to the pipeline for generation with HunyuanDiT.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`):
The height in pixels of the generated image.
@@ -615,7 +613,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -623,7 +621,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -651,19 +649,19 @@ def __call__(
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
plain tuple.
- callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+ callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
A callback function or a list of callback functions to be called at the end of each denoising step.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
A list of tensor inputs that should be passed to the callback function. If not defined, all tensor
inputs will be passed.
guidance_rescale (`float`, *optional*, defaults to 0.0):
Rescale the noise_cfg according to `guidance_rescale`. Based on findings of [Common Diffusion Noise
Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). See Section 3.4
- original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
+ original_size (`tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
The original size of the image. Used to calculate the time ids.
- target_size (`Tuple[int, int]`, *optional*):
+ target_size (`tuple[int, int]`, *optional*):
The target size of the image. Used to calculate the time ids.
- crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`):
+ crops_coords_top_left (`tuple[int, int]`, *optional*, defaults to `(0, 0)`):
The top left coordinates of the crop. Used to calculate the time ids.
use_resolution_binning (`bool`, *optional*, defaults to `True`):
Whether to use resolution binning or not. If `True`, the input resolution will be mapped to the closest
diff --git a/src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py b/src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py
index c6cc724a71f0..731ac27a0ff5 100644
--- a/src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py
+++ b/src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import PIL
@@ -86,14 +86,14 @@ class I2VGenXLPipelineOutput(BaseOutput):
Output class for image-to-video pipeline.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised
PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`
"""
- frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]]
+ frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]
class I2VGenXLPipeline(
@@ -165,15 +165,15 @@ def encode_prompt(
device,
num_videos_per_prompt,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -181,7 +181,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -263,7 +263,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if self.do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -445,7 +445,7 @@ def check_inputs(
and not isinstance(image, list)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+ "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
f" {type(image)}"
)
@@ -511,34 +511,34 @@ def prepare_latents(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
- height: Optional[int] = 704,
- width: Optional[int] = 1280,
- target_fps: Optional[int] = 16,
+ height: int | None = 704,
+ width: int | None = 1280,
+ target_fps: int | None = 16,
num_frames: int = 16,
num_inference_steps: int = 50,
guidance_scale: float = 9.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
eta: float = 0.0,
- num_videos_per_prompt: Optional[int] = 1,
- decode_chunk_size: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ decode_chunk_size: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = 1,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = 1,
):
r"""
The call function to the pipeline for image-to-video generation with [`I2VGenXLPipeline`].
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`):
+ image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.Tensor`):
Image or images to guide image generation. If you provide a tensor, it needs to be compatible with
[`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
@@ -555,7 +555,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*):
@@ -567,7 +567,7 @@ def __call__(
The number of frames to decode at a time. The higher the chunk size, the higher the temporal
consistency between frames, but also the higher the memory consumption. By default, the decoder will
decode all frames at once for maximal quality. Reduce `decode_chunk_size` to reduce memory usage.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -745,7 +745,7 @@ def __call__(
# https://github.com/ali-vilab/i2vgen-xl/blob/main/utils/transforms.py.
-def _convert_pt_to_pil(image: Union[torch.Tensor, List[torch.Tensor]]):
+def _convert_pt_to_pil(image: torch.Tensor | list[torch.Tensor]):
if isinstance(image, list) and isinstance(image[0], torch.Tensor):
image = torch.cat(image, 0)
@@ -761,7 +761,7 @@ def _convert_pt_to_pil(image: Union[torch.Tensor, List[torch.Tensor]]):
def _resize_bilinear(
- image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]], resolution: Tuple[int, int]
+ image: torch.Tensor | list[torch.Tensor] | PIL.Image.Image | list[PIL.Image.Image], resolution: tuple[int, int]
):
# First convert the images to PIL in case they are float tensors (only relevant for tests now).
image = _convert_pt_to_pil(image)
@@ -774,7 +774,7 @@ def _resize_bilinear(
def _center_crop_wide(
- image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]], resolution: Tuple[int, int]
+ image: torch.Tensor | list[torch.Tensor] | PIL.Image.Image | list[PIL.Image.Image], resolution: tuple[int, int]
):
# First convert the images to PIL in case they are float tensors (only relevant for tests now).
image = _convert_pt_to_pil(image)
diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py
index 33529f5d0954..d635057f2b05 100644
--- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py
+++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, List, Optional, Union
+from typing import Callable
import torch
from transformers import (
@@ -90,7 +90,7 @@ class KandinskyPipeline(DiffusionPipeline):
Frozen text-encoder.
tokenizer ([`XLMRobertaTokenizer`]):
Tokenizer of class
- scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
+ scheduler (`DDIMScheduler` | `DDPMScheduler`):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
@@ -105,7 +105,7 @@ def __init__(
text_encoder: MultilingualCLIP,
tokenizer: XLMRobertaTokenizer,
unet: UNet2DConditionModel,
- scheduler: Union[DDIMScheduler, DDPMScheduler],
+ scheduler: DDIMScheduler | DDPMScheduler,
movq: VQModel,
):
super().__init__()
@@ -173,7 +173,7 @@ def _encode_prompt(
text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -237,19 +237,19 @@ def _encode_prompt(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- image_embeds: Union[torch.Tensor, List[torch.Tensor]],
- negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ image_embeds: torch.Tensor | list[torch.Tensor],
+ negative_image_embeds: torch.Tensor | list[torch.Tensor],
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
):
@@ -257,13 +257,13 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for text prompt, that will be used to condition the image generation.
- negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for negative text prompt, will be used to condition the image generation.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
height (`int`, *optional*, defaults to 512):
@@ -281,7 +281,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py
index 7286bcbee17b..21d5c98efa90 100644
--- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py
+++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py
@@ -11,7 +11,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.
-from typing import Callable, List, Optional, Union
+from typing import Callable
import PIL.Image
import torch
@@ -122,7 +122,7 @@ class KandinskyCombinedPipeline(DiffusionPipeline):
Frozen text-encoder.
tokenizer ([`XLMRobertaTokenizer`]):
Tokenizer of class
- scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
+ scheduler (`DDIMScheduler` | `DDPMScheduler`):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
@@ -150,7 +150,7 @@ def __init__(
text_encoder: MultilingualCLIP,
tokenizer: XLMRobertaTokenizer,
unet: UNet2DConditionModel,
- scheduler: Union[DDIMScheduler, DDPMScheduler],
+ scheduler: DDIMScheduler | DDPMScheduler,
movq: VQModel,
prior_prior: PriorTransformer,
prior_image_encoder: CLIPVisionModelWithProjection,
@@ -190,10 +190,10 @@ def __init__(
movq=movq,
)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models (`unet`, `text_encoder`, `vae`, and `safety checker` state dicts) to CPU using 🤗
Accelerate, significantly reducing memory usage. Models are moved to a `torch.device('meta')` and loaded on a
@@ -216,8 +216,8 @@ def set_progress_bar_config(self, **kwargs):
@replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
@@ -225,10 +225,10 @@ def __call__(
width: int = 512,
prior_guidance_scale: float = 4.0,
prior_num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
):
@@ -236,9 +236,9 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -265,7 +265,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -340,7 +340,7 @@ class KandinskyImg2ImgCombinedPipeline(DiffusionPipeline):
Frozen text-encoder.
tokenizer ([`XLMRobertaTokenizer`]):
Tokenizer of class
- scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
+ scheduler (`DDIMScheduler` | `DDPMScheduler`):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
@@ -368,7 +368,7 @@ def __init__(
text_encoder: MultilingualCLIP,
tokenizer: XLMRobertaTokenizer,
unet: UNet2DConditionModel,
- scheduler: Union[DDIMScheduler, DDPMScheduler],
+ scheduler: DDIMScheduler | DDPMScheduler,
movq: VQModel,
prior_prior: PriorTransformer,
prior_image_encoder: CLIPVisionModelWithProjection,
@@ -408,10 +408,10 @@ def __init__(
movq=movq,
)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
@@ -435,9 +435,9 @@ def set_progress_bar_config(self, **kwargs):
@replace_example_docstring(IMAGE2IMAGE_EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
+ negative_prompt: str | list[str] | None = None,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
@@ -446,10 +446,10 @@ def __call__(
width: int = 512,
prior_guidance_scale: float = 4.0,
prior_num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
):
@@ -457,13 +457,13 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded
again.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -496,7 +496,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -581,7 +581,7 @@ class KandinskyInpaintCombinedPipeline(DiffusionPipeline):
Frozen text-encoder.
tokenizer ([`XLMRobertaTokenizer`]):
Tokenizer of class
- scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
+ scheduler (`DDIMScheduler` | `DDPMScheduler`):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
@@ -609,7 +609,7 @@ def __init__(
text_encoder: MultilingualCLIP,
tokenizer: XLMRobertaTokenizer,
unet: UNet2DConditionModel,
- scheduler: Union[DDIMScheduler, DDPMScheduler],
+ scheduler: DDIMScheduler | DDPMScheduler,
movq: VQModel,
prior_prior: PriorTransformer,
prior_image_encoder: CLIPVisionModelWithProjection,
@@ -649,10 +649,10 @@ def __init__(
movq=movq,
)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
@@ -676,10 +676,10 @@ def set_progress_bar_config(self, **kwargs):
@replace_example_docstring(INPAINT_EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
- mask_image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
+ mask_image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
+ negative_prompt: str | list[str] | None = None,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
@@ -687,10 +687,10 @@ def __call__(
width: int = 512,
prior_guidance_scale: float = 4.0,
prior_num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
):
@@ -698,9 +698,9 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded
again.
@@ -709,7 +709,7 @@ def __call__(
black pixels will be preserved. If `mask_image` is a PIL image, it will be converted to a single
channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3,
so the expected shape would be `(B, H, W, 1)`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -736,7 +736,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py
index f5e41d499dc3..f33317447b49 100644
--- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py
+++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, List, Optional, Union
+from typing import Callable
import PIL.Image
import torch
@@ -207,7 +207,7 @@ def _encode_prompt(
text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -298,20 +298,20 @@ def add_noise(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
+ prompt: str | list[str],
+ image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
image_embeds: torch.Tensor,
negative_image_embeds: torch.Tensor,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
strength: float = 0.3,
guidance_scale: float = 7.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
):
@@ -319,16 +319,16 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
image (`torch.Tensor`, `PIL.Image.Image`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
process.
- image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for text prompt, that will be used to condition the image generation.
- negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for negative text prompt, will be used to condition the image generation.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
height (`int`, *optional*, defaults to 512):
@@ -352,7 +352,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
output_type (`str`, *optional*, defaults to `"pil"`):
diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py
index 731fce499859..4dba85446db5 100644
--- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py
+++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
from copy import deepcopy
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -134,7 +134,7 @@ def prepare_mask_and_masked_image(image, mask, height, width):
binarized (``mask > 0.5``) and cast to ``torch.float32`` too.
Args:
- image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint.
+ image (np.array | PIL.Image | torch.Tensor): The image to inpaint.
It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width``
``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``.
mask (_type_): The mask to apply to the image, i.e. regions to inpaint.
@@ -338,7 +338,7 @@ def _encode_prompt(
text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -402,21 +402,21 @@ def _encode_prompt(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[torch.Tensor, PIL.Image.Image],
- mask_image: Union[torch.Tensor, PIL.Image.Image, np.ndarray],
+ prompt: str | list[str],
+ image: torch.Tensor | PIL.Image.Image,
+ mask_image: torch.Tensor | PIL.Image.Image | np.ndarray,
image_embeds: torch.Tensor,
negative_image_embeds: torch.Tensor,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
):
@@ -424,7 +424,7 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
image (`torch.Tensor`, `PIL.Image.Image` or `np.ndarray`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
@@ -437,11 +437,11 @@ def __call__(
image or numpy array, mask should also be a either PIL image or numpy array. If it is a PIL image, it
will be converted to a single channel (luminance) before use. If it is a nummpy array, the expected
shape is `(H, W)`.
- image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for text prompt, that will be used to condition the image generation.
- negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for negative text prompt, will be used to condition the image generation.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
height (`int`, *optional*, defaults to 512):
@@ -459,7 +459,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py
index 10ea8005c90d..523fd010eb7f 100644
--- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py
+++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -126,12 +125,12 @@ class KandinskyPriorPipelineOutput(BaseOutput):
Args:
image_embeds (`torch.Tensor`)
clip image embeddings for text prompt
- negative_image_embeds (`List[PIL.Image.Image]` or `np.ndarray`)
+ negative_image_embeds (`list[PIL.Image.Image]` or `np.ndarray`)
clip image embeddings for unconditional tokens
"""
- image_embeds: Union[torch.Tensor, np.ndarray]
- negative_image_embeds: Union[torch.Tensor, np.ndarray]
+ image_embeds: torch.Tensor | np.ndarray
+ negative_image_embeds: torch.Tensor | np.ndarray
class KandinskyPriorPipeline(DiffusionPipeline):
@@ -182,13 +181,13 @@ def __init__(
@replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING)
def interpolate(
self,
- images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]],
- weights: List[float],
+ images_and_prompts: list[str | PIL.Image.Image | torch.Tensor],
+ weights: list[float],
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- negative_prior_prompt: Optional[str] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ negative_prior_prompt: str | None = None,
negative_prompt: str = "",
guidance_scale: float = 4.0,
device=None,
@@ -197,16 +196,16 @@ def interpolate(
Function invoked when using the prior pipeline for interpolation.
Args:
- images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`):
+ images_and_prompts (`list[str | PIL.Image.Image | torch.Tensor]`):
list of prompts and images to guide the image generation.
- weights: (`List[float]`):
+ weights: (`list[float]`):
list of weights for each condition in `images_and_prompts`
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
num_inference_steps (`int`, *optional*, defaults to 25):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -216,7 +215,7 @@ def interpolate(
negative_prior_prompt (`str`, *optional*):
The prompt not to guide the prior diffusion process. Ignored when not using guidance (i.e., ignored if
`guidance_scale` is less than `1`).
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if
`guidance_scale` is less than `1`).
guidance_scale (`float`, *optional*, defaults to 4.0):
@@ -346,7 +345,7 @@ def _encode_prompt(
text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -407,23 +406,23 @@ def _encode_prompt(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
guidance_scale: float = 4.0,
- output_type: Optional[str] = "pt",
+ output_type: str | None = "pt",
return_dict: bool = True,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -431,7 +430,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 25):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/kandinsky/text_encoder.py b/src/diffusers/pipelines/kandinsky/text_encoder.py
index caa0029f00ca..58cc9ac4d3ed 100644
--- a/src/diffusers/pipelines/kandinsky/text_encoder.py
+++ b/src/diffusers/pipelines/kandinsky/text_encoder.py
@@ -20,6 +20,8 @@ def __init__(self, config, *args, **kwargs):
self.LinearTransformation = torch.nn.Linear(
in_features=config.transformerDimensions, out_features=config.numDims
)
+ if hasattr(self, "post_init"):
+ self.post_init()
def forward(self, input_ids, attention_mask):
embs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)[0]
diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py
index 429253e99898..5129b3f548e8 100644
--- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py
+++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import torch
@@ -77,7 +77,7 @@ class KandinskyV22Pipeline(DiffusionPipeline):
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
- scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
+ scheduler (`DDIMScheduler` | `DDPMScheduler`):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
@@ -131,28 +131,28 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image_embeds: Union[torch.Tensor, List[torch.Tensor]],
- negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
+ image_embeds: torch.Tensor | list[torch.Tensor],
+ negative_image_embeds: torch.Tensor | list[torch.Tensor],
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for text prompt, that will be used to condition the image generation.
- negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for negative text prompt, will be used to condition the image generation.
height (`int`, *optional*, defaults to 512):
The height in pixels of the generated image.
@@ -169,7 +169,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -186,7 +186,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py
index fc2083247bb0..26e163a70142 100644
--- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py
+++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import PIL.Image
import torch
@@ -112,7 +112,7 @@ class KandinskyV22CombinedPipeline(DiffusionPipeline):
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
- scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
+ scheduler (`DDIMScheduler` | `DDPMScheduler`):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
@@ -176,10 +176,10 @@ def __init__(
movq=movq,
)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
@@ -203,8 +203,8 @@ def set_progress_bar_config(self, **kwargs):
@replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
@@ -212,24 +212,24 @@ def __call__(
width: int = 512,
prior_guidance_scale: float = 4.0,
prior_num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
- prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ prior_callback_on_step_end: Callable[[int, int], None] | None = None,
+ prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -256,7 +256,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -272,7 +272,7 @@ def __call__(
A function that calls at the end of each denoising steps during the inference of the prior pipeline.
The function is called with the following arguments: `prior_callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`.
- prior_callback_on_step_end_tensor_inputs (`List`, *optional*):
+ prior_callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `prior_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 prior pipeline class.
@@ -281,7 +281,7 @@ def __call__(
The function is called with the following arguments: `callback_on_step_end(self: DiffusionPipeline,
step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors
as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
@@ -340,7 +340,7 @@ class KandinskyV22Img2ImgCombinedPipeline(DiffusionPipeline):
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
- scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
+ scheduler (`DDIMScheduler` | `DDPMScheduler`):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
@@ -404,10 +404,10 @@ def __init__(
movq=movq,
)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
- def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
@@ -417,7 +417,7 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t
self.prior_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device)
self.decoder_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device)
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
@@ -441,9 +441,9 @@ def set_progress_bar_config(self, **kwargs):
@replace_example_docstring(IMAGE2IMAGE_EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
+ negative_prompt: str | list[str] | None = None,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
strength: float = 0.3,
@@ -452,28 +452,28 @@ def __call__(
width: int = 512,
prior_guidance_scale: float = 4.0,
prior_num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
- prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ prior_callback_on_step_end: Callable[[int, int], None] | None = None,
+ prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded
again.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -506,7 +506,7 @@ def __call__(
prior_num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -589,7 +589,7 @@ class KandinskyV22InpaintCombinedPipeline(DiffusionPipeline):
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
- scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]):
+ scheduler (`DDIMScheduler` | `DDPMScheduler`):
A scheduler to be used in combination with `unet` to generate image latents.
unet ([`UNet2DConditionModel`]):
Conditional U-Net architecture to denoise the image embedding.
@@ -653,10 +653,10 @@ def __init__(
movq=movq,
)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
@@ -680,10 +680,10 @@ def set_progress_bar_config(self, **kwargs):
@replace_example_docstring(INPAINT_EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
- mask_image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
+ mask_image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
+ negative_prompt: str | list[str] | None = None,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
@@ -691,23 +691,23 @@ def __call__(
width: int = 512,
prior_guidance_scale: float = 4.0,
prior_num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ prior_callback_on_step_end: Callable[[int, int], None] | None = None,
+ prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded
again.
@@ -716,7 +716,7 @@ def __call__(
black pixels will be preserved. If `mask_image` is a PIL image, it will be converted to a single
channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3,
so the expected shape would be `(B, H, W, 1)`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -743,7 +743,7 @@ def __call__(
prior_num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -759,7 +759,7 @@ def __call__(
A function that calls at the end of each denoising steps during the inference. The function is called
with the following arguments: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep:
int, callback_kwargs: Dict)`.
- prior_callback_on_step_end_tensor_inputs (`List`, *optional*):
+ prior_callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `prior_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.
@@ -768,7 +768,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py
index c5faae82796b..01001e0c9eba 100644
--- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py
+++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, List, Optional, Union
+from typing import Callable
import torch
@@ -160,18 +160,18 @@ def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):
@torch.no_grad()
def __call__(
self,
- image_embeds: Union[torch.Tensor, List[torch.Tensor]],
- negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
+ image_embeds: torch.Tensor | list[torch.Tensor],
+ negative_image_embeds: torch.Tensor | list[torch.Tensor],
hint: torch.Tensor,
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
):
@@ -179,15 +179,15 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
hint (`torch.Tensor`):
The controlnet condition.
- image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for text prompt, that will be used to condition the image generation.
- negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for negative text prompt, will be used to condition the image generation.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
height (`int`, *optional*, defaults to 512):
@@ -205,7 +205,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py
index 54154c6ec1f2..891235558330 100644
--- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py
+++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, List, Optional, Union
+from typing import Callable
import PIL.Image
import torch
@@ -200,9 +200,9 @@ def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dt
@torch.no_grad()
def __call__(
self,
- image_embeds: Union[torch.Tensor, List[torch.Tensor]],
- image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
- negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
+ image_embeds: torch.Tensor | list[torch.Tensor],
+ image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
+ negative_image_embeds: torch.Tensor | list[torch.Tensor],
hint: torch.Tensor,
height: int = 512,
width: int = 512,
@@ -210,9 +210,9 @@ def __call__(
guidance_scale: float = 4.0,
strength: float = 0.3,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
return_dict: bool = True,
):
@@ -220,9 +220,9 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for text prompt, that will be used to condition the image generation.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded
again.
@@ -234,7 +234,7 @@ def __call__(
`num_inference_steps`. A value of 1, therefore, essentially ignores `image`.
hint (`torch.Tensor`):
The controlnet condition.
- negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for negative text prompt, will be used to condition the image generation.
height (`int`, *optional*, defaults to 512):
The height in pixels of the generated image.
@@ -251,7 +251,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
output_type (`str`, *optional*, defaults to `"pil"`):
diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py
index 3b2509098fd1..a43c085a8921 100644
--- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py
+++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import PIL.Image
import torch
@@ -183,29 +183,29 @@ def num_timesteps(self):
@torch.no_grad()
def __call__(
self,
- image_embeds: Union[torch.Tensor, List[torch.Tensor]],
- image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]],
- negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
+ image_embeds: torch.Tensor | list[torch.Tensor],
+ image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image],
+ negative_image_embeds: torch.Tensor | list[torch.Tensor],
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
strength: float = 0.3,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for text prompt, that will be used to condition the image generation.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded
again.
@@ -215,7 +215,7 @@ def __call__(
denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will
be maximum and the denoising process will run for the full number of iterations specified in
`num_inference_steps`. A value of 1, therefore, essentially ignores `image`.
- negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for negative text prompt, will be used to condition the image generation.
height (`int`, *optional*, defaults to 512):
The height in pixels of the generated image.
@@ -232,7 +232,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
output_type (`str`, *optional*, defaults to `"pil"`):
@@ -245,7 +245,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py
index a61673293e1f..31bd88103a06 100644
--- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py
+++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py
@@ -13,7 +13,7 @@
# limitations under the License.
from copy import deepcopy
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -132,7 +132,7 @@ def prepare_mask_and_masked_image(image, mask, height, width):
binarized (``mask > 0.5``) and cast to ``torch.float32`` too.
Args:
- image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint.
+ image (np.array | PIL.Image | torch.Tensor): The image to inpaint.
It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width``
``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``.
mask (_type_): The mask to apply to the image, i.e. regions to inpaint.
@@ -302,28 +302,28 @@ def num_timesteps(self):
@torch.no_grad()
def __call__(
self,
- image_embeds: Union[torch.Tensor, List[torch.Tensor]],
- image: Union[torch.Tensor, PIL.Image.Image],
- mask_image: Union[torch.Tensor, PIL.Image.Image, np.ndarray],
- negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]],
+ image_embeds: torch.Tensor | list[torch.Tensor],
+ image: torch.Tensor | PIL.Image.Image,
+ mask_image: torch.Tensor | PIL.Image.Image | np.ndarray,
+ negative_image_embeds: torch.Tensor | list[torch.Tensor],
height: int = 512,
width: int = 512,
num_inference_steps: int = 100,
guidance_scale: float = 4.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for text prompt, that will be used to condition the image generation.
image (`PIL.Image.Image`):
`Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will
@@ -333,7 +333,7 @@ def __call__(
black pixels will be preserved. If `mask_image` is a PIL image, it will be converted to a single
channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3,
so the expected shape would be `(B, H, W, 1)`.
- negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`):
+ negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`):
The clip image embeddings for negative text prompt, will be used to condition the image generation.
height (`int`, *optional*, defaults to 512):
The height in pixels of the generated image.
@@ -350,7 +350,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -367,7 +367,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py
index bc67847831a5..41f4474c3906 100644
--- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py
+++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py
@@ -1,4 +1,4 @@
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import PIL.Image
import torch
@@ -137,13 +137,13 @@ def __init__(
@replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING)
def interpolate(
self,
- images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]],
- weights: List[float],
+ images_and_prompts: list[str | PIL.Image.Image | torch.Tensor],
+ weights: list[float],
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- negative_prior_prompt: Optional[str] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ negative_prior_prompt: str | None = None,
negative_prompt: str = "",
guidance_scale: float = 4.0,
device=None,
@@ -152,16 +152,16 @@ def interpolate(
Function invoked when using the prior pipeline for interpolation.
Args:
- images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`):
+ images_and_prompts (`list[str | PIL.Image.Image | torch.Tensor]`):
list of prompts and images to guide the image generation.
- weights: (`List[float]`):
+ weights: (`list[float]`):
list of weights for each condition in `images_and_prompts`
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -171,7 +171,7 @@ def interpolate(
negative_prior_prompt (`str`, *optional*):
The prompt not to guide the prior diffusion process. Ignored when not using guidance (i.e., ignored if
`guidance_scale` is less than `1`).
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if
`guidance_scale` is less than `1`).
guidance_scale (`float`, *optional*, defaults to 4.0):
@@ -303,7 +303,7 @@ def _encode_prompt(
text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -376,25 +376,25 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
guidance_scale: float = 4.0,
- output_type: Optional[str] = "pt", # pt only
+ output_type: str | None = "pt", # pt only
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -402,7 +402,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -425,7 +425,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py
index b586d166118b..adbc3a5badc5 100644
--- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py
+++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py
@@ -1,5 +1,3 @@
-from typing import List, Optional, Union
-
import PIL.Image
import torch
from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection
@@ -161,13 +159,13 @@ def get_timesteps(self, num_inference_steps, strength, device):
@replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING)
def interpolate(
self,
- images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]],
- weights: List[float],
+ images_and_prompts: list[str | PIL.Image.Image | torch.Tensor],
+ weights: list[float],
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- negative_prior_prompt: Optional[str] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ negative_prior_prompt: str | None = None,
negative_prompt: str = "",
guidance_scale: float = 4.0,
device=None,
@@ -176,16 +174,16 @@ def interpolate(
Function invoked when using the prior pipeline for interpolation.
Args:
- images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`):
+ images_and_prompts (`list[str | PIL.Image.Image | torch.Tensor]`):
list of prompts and images to guide the image generation.
- weights: (`List[float]`):
+ weights: (`list[float]`):
list of weights for each condition in `images_and_prompts`
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -195,7 +193,7 @@ def interpolate(
negative_prior_prompt (`str`, *optional*):
The prompt not to guide the prior diffusion process. Ignored when not using guidance (i.e., ignored if
`guidance_scale` is less than `1`).
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if
`guidance_scale` is less than `1`).
guidance_scale (`float`, *optional*, defaults to 4.0):
@@ -249,7 +247,7 @@ def interpolate(
def _encode_image(
self,
- image: Union[torch.Tensor, List[PIL.Image.Image]],
+ image: torch.Tensor | list[PIL.Image.Image],
device,
num_images_per_prompt,
):
@@ -341,7 +339,7 @@ def _encode_prompt(
text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -402,22 +400,22 @@ def _encode_prompt(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]],
+ prompt: str | list[str],
+ image: torch.Tensor | list[torch.Tensor] | PIL.Image.Image | list[PIL.Image.Image],
strength: float = 0.3,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
guidance_scale: float = 4.0,
- output_type: Optional[str] = "pt", # pt only
+ output_type: str | None = "pt", # pt only
return_dict: bool = True,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
strength (`float`, *optional*, defaults to 0.8):
Conceptually, indicates how much to transform the reference `emb`. Must be between 0 and 1. `image`
@@ -425,7 +423,7 @@ def __call__(
denoising steps depends on the amount of noise initially added.
emb (`torch.Tensor`):
The image embedding.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -433,7 +431,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
guidance_scale (`float`, *optional*, defaults to 4.0):
@@ -480,7 +478,7 @@ def __call__(
prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
)
- if not isinstance(image, List):
+ if not isinstance(image, list):
image = [image]
if isinstance(image[0], torch.Tensor):
diff --git a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py
index 57cc0270442d..97353c95c9c7 100644
--- a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py
+++ b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py
@@ -1,4 +1,4 @@
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import torch
from transformers import T5EncoderModel, T5Tokenizer
@@ -96,17 +96,17 @@ def encode_prompt(
num_images_per_prompt=1,
device=None,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
_cut_context=False,
- attention_mask: Optional[torch.Tensor] = None,
- negative_attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ negative_attention_mask: torch.Tensor | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
@@ -114,7 +114,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -181,7 +181,7 @@ def encode_prompt(
attention_mask = attention_mask.repeat(num_images_per_prompt, 1)
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
@@ -335,36 +335,36 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_inference_steps: int = 25,
guidance_scale: float = 3.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- height: Optional[int] = 1024,
- width: Optional[int] = 1024,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- negative_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ height: int | None = 1024,
+ width: int | None = 1024,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ negative_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
latents=None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
num_inference_steps (`int`, *optional*, defaults to 25):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 3.0):
@@ -373,7 +373,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -386,7 +386,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py
index 73c268897502..beb4caafb6d3 100644
--- a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py
+++ b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py
@@ -1,5 +1,5 @@
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import PIL
import PIL.Image
@@ -111,17 +111,17 @@ def encode_prompt(
num_images_per_prompt=1,
device=None,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
_cut_context=True,
- attention_mask: Optional[torch.Tensor] = None,
- negative_attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ negative_attention_mask: torch.Tensor | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
@@ -129,7 +129,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -196,7 +196,7 @@ def encode_prompt(
attention_mask = attention_mask.repeat(num_images_per_prompt, 1)
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
@@ -401,32 +401,32 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]] = None,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image] = None,
strength: float = 0.3,
num_inference_steps: int = 25,
guidance_scale: float = 3.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- negative_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ negative_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
process.
strength (`float`, *optional*, defaults to 0.8):
@@ -444,13 +444,13 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -474,7 +474,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py
index 2b666f0ec697..1c94a8219e2a 100644
--- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py
+++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py
@@ -13,7 +13,7 @@
# limitations under the License.
import html
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import regex as re
import torch
@@ -312,10 +312,10 @@ def get_sparse_params(self, sample, device):
def _encode_prompt_qwen(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
max_sequence_length: int = 256,
- dtype: Optional[torch.dtype] = None,
+ dtype: torch.dtype | None = None,
):
"""
Encode prompt using Qwen2.5-VL text encoder.
@@ -324,14 +324,14 @@ def _encode_prompt_qwen(
video generation.
Args:
- prompt (Union[str, List[str]]): Input prompt or list of prompts
+ prompt (str | list[str]): Input prompt or list of prompts
device (torch.device): Device to run encoding on
num_videos_per_prompt (int): Number of videos to generate per prompt
max_sequence_length (int): Maximum sequence length for tokenization
dtype (torch.dtype): Data type for embeddings
Returns:
- Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths
+ tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths
"""
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -382,9 +382,9 @@ def _encode_prompt_qwen(
def _encode_prompt_clip(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
"""
Encode prompt using CLIP text encoder.
@@ -393,7 +393,7 @@ def _encode_prompt_clip(
semantic information.
Args:
- prompt (Union[str, List[str]]): Input prompt or list of prompts
+ prompt (str | list[str]): Input prompt or list of prompts
device (torch.device): Device to run encoding on
num_videos_per_prompt (int): Number of videos to generate per prompt
dtype (torch.dtype): Data type for embeddings
@@ -419,11 +419,11 @@ def _encode_prompt_clip(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_videos_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes a single prompt (positive or negative) into text encoder hidden states.
@@ -432,7 +432,7 @@ def encode_prompt(
representations for video generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
Prompt to be encoded.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos to generate per prompt.
@@ -444,7 +444,7 @@ def encode_prompt(
Torch dtype.
Returns:
- Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- Qwen text embeddings of shape (batch_size * num_videos_per_prompt, sequence_length, embedding_dim)
- CLIP pooled embeddings of shape (batch_size * num_videos_per_prompt, clip_embedding_dim)
- Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size *
@@ -603,10 +603,10 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Prepare initial latent variables for video generation.
@@ -683,37 +683,35 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 768,
num_frames: int = 121,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds_qwen: Optional[torch.Tensor] = None,
- prompt_embeds_clip: Optional[torch.Tensor] = None,
- negative_prompt_embeds_qwen: Optional[torch.Tensor] = None,
- negative_prompt_embeds_clip: Optional[torch.Tensor] = None,
- prompt_cu_seqlens: Optional[torch.Tensor] = None,
- negative_prompt_cu_seqlens: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds_qwen: torch.Tensor | None = None,
+ prompt_embeds_clip: torch.Tensor | None = None,
+ negative_prompt_embeds_qwen: torch.Tensor | None = None,
+ negative_prompt_embeds_clip: torch.Tensor | None = None,
+ prompt_cu_seqlens: torch.Tensor | None = None,
+ negative_prompt_cu_seqlens: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, pass `prompt_embeds` instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to avoid during video generation. If not defined, pass `negative_prompt_embeds`
instead. Ignored when not using guidance (`guidance_scale` < `1`).
height (`int`, defaults to `512`):
@@ -728,7 +726,7 @@ def __call__(
Guidance scale as defined in classifier-free guidance.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A torch generator to make generation deterministic.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents.
@@ -742,7 +740,7 @@ def __call__(
Whether or not to return a [`KandinskyPipelineOutput`].
callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
A function that is called at the end of each denoising step.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function.
max_sequence_length (`int`, defaults to `512`):
The maximum sequence length for text encoding.
diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py
index f965cdad8f3e..244db7300767 100644
--- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py
+++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py
@@ -13,7 +13,7 @@
# limitations under the License.
import html
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import regex as re
@@ -183,11 +183,11 @@ def __init__(
def _encode_prompt_qwen(
self,
- prompt: List[str],
- image: Optional[PipelineImageInput] = None,
- device: Optional[torch.device] = None,
+ prompt: list[str],
+ image: PipelineImageInput | None = None,
+ device: torch.device | None = None,
max_sequence_length: int = 1024,
- dtype: Optional[torch.dtype] = None,
+ dtype: torch.dtype | None = None,
):
"""
Encode prompt using Qwen2.5-VL text encoder.
@@ -196,14 +196,14 @@ def _encode_prompt_qwen(
image generation.
Args:
- prompt List[str]: Input list of prompts
+ prompt list[str]: Input list of prompts
image (PipelineImageInput): Input list of images to condition the generation on
device (torch.device): Device to run encoding on
max_sequence_length (int): Maximum sequence length for tokenization
dtype (torch.dtype): Data type for embeddings
Returns:
- Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths
+ tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths
"""
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -258,9 +258,9 @@ def _encode_prompt_qwen(
def _encode_prompt_clip(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
"""
Encode prompt using CLIP text encoder.
@@ -269,7 +269,7 @@ def _encode_prompt_clip(
semantic information.
Args:
- prompt (Union[str, List[str]]): Input prompt or list of prompts
+ prompt (str | list[str]): Input prompt or list of prompts
device (torch.device): Device to run encoding on
dtype (torch.dtype): Data type for embeddings
@@ -294,12 +294,12 @@ def _encode_prompt_clip(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
image: torch.Tensor,
num_images_per_prompt: int = 1,
max_sequence_length: int = 1024,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes a single prompt (positive or negative) into text encoder hidden states.
@@ -308,7 +308,7 @@ def encode_prompt(
representations for image generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
Prompt to be encoded.
num_images_per_prompt (`int`, *optional*, defaults to 1):
Number of images to generate per prompt.
@@ -320,7 +320,7 @@ def encode_prompt(
Torch dtype.
Returns:
- Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- Qwen text embeddings of shape (batch_size * num_images_per_prompt, sequence_length, embedding_dim)
- CLIP pooled embeddings of shape (batch_size * num_images_per_prompt, clip_embedding_dim)
- Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size *
@@ -486,10 +486,10 @@ def prepare_latents(
num_channels_latents: int = 16,
height: int = 1024,
width: int = 1024,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Prepare initial latent variables for image-to-image generation.
@@ -568,27 +568,25 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 3.5,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds_qwen: Optional[torch.Tensor] = None,
- prompt_embeds_clip: Optional[torch.Tensor] = None,
- negative_prompt_embeds_qwen: Optional[torch.Tensor] = None,
- negative_prompt_embeds_clip: Optional[torch.Tensor] = None,
- prompt_cu_seqlens: Optional[torch.Tensor] = None,
- negative_prompt_cu_seqlens: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds_qwen: torch.Tensor | None = None,
+ prompt_embeds_clip: torch.Tensor | None = None,
+ negative_prompt_embeds_qwen: torch.Tensor | None = None,
+ negative_prompt_embeds_clip: torch.Tensor | None = None,
+ prompt_cu_seqlens: torch.Tensor | None = None,
+ negative_prompt_cu_seqlens: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 1024,
):
r"""
@@ -597,9 +595,9 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds`
instead. Ignored when not using guidance (`guidance_scale` < `1`).
height (`int`):
@@ -612,7 +610,7 @@ def __call__(
Guidance scale as defined in classifier-free guidance.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A torch generator to make generation deterministic.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents.
diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py
index d457c9b69657..ad4bb182d248 100644
--- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py
+++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py
@@ -13,7 +13,7 @@
# limitations under the License.
import html
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import regex as re
import torch
@@ -309,10 +309,10 @@ def get_sparse_params(self, sample, device):
def _encode_prompt_qwen(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
max_sequence_length: int = 256,
- dtype: Optional[torch.dtype] = None,
+ dtype: torch.dtype | None = None,
):
"""
Encode prompt using Qwen2.5-VL text encoder.
@@ -321,13 +321,13 @@ def _encode_prompt_qwen(
video generation.
Args:
- prompt (Union[str, List[str]]): Input prompt or list of prompts
+ prompt (str | list[str]): Input prompt or list of prompts
device (torch.device): Device to run encoding on
max_sequence_length (int): Maximum sequence length for tokenization
dtype (torch.dtype): Data type for embeddings
Returns:
- Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths
+ tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths
"""
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -377,9 +377,9 @@ def _encode_prompt_qwen(
def _encode_prompt_clip(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
"""
Encode prompt using CLIP text encoder.
@@ -388,7 +388,7 @@ def _encode_prompt_clip(
semantic information.
Args:
- prompt (Union[str, List[str]]): Input prompt or list of prompts
+ prompt (str | list[str]): Input prompt or list of prompts
device (torch.device): Device to run encoding on
dtype (torch.dtype): Data type for embeddings
@@ -453,11 +453,11 @@ def normalize_first_frame(self, latents, reference_frames=5, clump_values=False)
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_videos_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes a single prompt (positive or negative) into text encoder hidden states.
@@ -466,7 +466,7 @@ def encode_prompt(
representations for video generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
Prompt to be encoded.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos to generate per prompt.
@@ -478,7 +478,7 @@ def encode_prompt(
Torch dtype.
Returns:
- Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- Qwen text embeddings of shape (batch_size * num_videos_per_prompt, sequence_length, embedding_dim)
- CLIP pooled embeddings of shape (batch_size * num_videos_per_prompt, clip_embedding_dim)
- Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size *
@@ -641,10 +641,10 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Prepare initial latent variables for image-to-video generation.
@@ -750,28 +750,26 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 768,
num_frames: int = 121,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds_qwen: Optional[torch.Tensor] = None,
- prompt_embeds_clip: Optional[torch.Tensor] = None,
- negative_prompt_embeds_qwen: Optional[torch.Tensor] = None,
- negative_prompt_embeds_clip: Optional[torch.Tensor] = None,
- prompt_cu_seqlens: Optional[torch.Tensor] = None,
- negative_prompt_cu_seqlens: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds_qwen: torch.Tensor | None = None,
+ prompt_embeds_clip: torch.Tensor | None = None,
+ negative_prompt_embeds_qwen: torch.Tensor | None = None,
+ negative_prompt_embeds_clip: torch.Tensor | None = None,
+ prompt_cu_seqlens: torch.Tensor | None = None,
+ negative_prompt_cu_seqlens: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
@@ -780,9 +778,9 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, pass `prompt_embeds` instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to avoid during video generation. If not defined, pass `negative_prompt_embeds`
instead. Ignored when not using guidance (`guidance_scale` < `1`).
height (`int`, defaults to `512`):
@@ -797,7 +795,7 @@ def __call__(
Guidance scale as defined in classifier-free guidance.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A torch generator to make generation deterministic.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents.
diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py
index bb5c40327b4e..2a58d4bed33a 100644
--- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py
+++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py
@@ -13,7 +13,7 @@
# limitations under the License.
import html
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import regex as re
@@ -184,10 +184,10 @@ def __init__(
def _encode_prompt_qwen(
self,
- prompt: List[str],
- device: Optional[torch.device] = None,
+ prompt: list[str],
+ device: torch.device | None = None,
max_sequence_length: int = 512,
- dtype: Optional[torch.dtype] = None,
+ dtype: torch.dtype | None = None,
):
"""
Encode prompt using Qwen2.5-VL text encoder.
@@ -196,13 +196,13 @@ def _encode_prompt_qwen(
image generation.
Args:
- prompt List[str]: Input list of prompts
+ prompt list[str]: Input list of prompts
device (torch.device): Device to run encoding on
max_sequence_length (int): Maximum sequence length for tokenization
dtype (torch.dtype): Data type for embeddings
Returns:
- Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths
+ tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths
"""
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -252,9 +252,9 @@ def _encode_prompt_qwen(
def _encode_prompt_clip(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
"""
Encode prompt using CLIP text encoder.
@@ -263,7 +263,7 @@ def _encode_prompt_clip(
semantic information.
Args:
- prompt (Union[str, List[str]]): Input prompt or list of prompts
+ prompt (str | list[str]): Input prompt or list of prompts
device (torch.device): Device to run encoding on
dtype (torch.dtype): Data type for embeddings
@@ -288,11 +288,11 @@ def _encode_prompt_clip(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes a single prompt (positive or negative) into text encoder hidden states.
@@ -301,7 +301,7 @@ def encode_prompt(
representations for image generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
Prompt to be encoded.
num_images_per_prompt (`int`, *optional*, defaults to 1):
Number of images to generate per prompt.
@@ -313,7 +313,7 @@ def encode_prompt(
Torch dtype.
Returns:
- Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- Qwen text embeddings of shape (batch_size * num_images_per_prompt, sequence_length, embedding_dim)
- CLIP pooled embeddings of shape (batch_size * num_images_per_prompt, clip_embedding_dim)
- Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size *
@@ -472,10 +472,10 @@ def prepare_latents(
num_channels_latents: int = 16,
height: int = 1024,
width: int = 1024,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Prepare initial latent variables for text-to-image generation.
@@ -535,36 +535,34 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 1024,
width: int = 1024,
num_inference_steps: int = 50,
guidance_scale: float = 3.5,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds_qwen: Optional[torch.Tensor] = None,
- prompt_embeds_clip: Optional[torch.Tensor] = None,
- negative_prompt_embeds_qwen: Optional[torch.Tensor] = None,
- negative_prompt_embeds_clip: Optional[torch.Tensor] = None,
- prompt_cu_seqlens: Optional[torch.Tensor] = None,
- negative_prompt_cu_seqlens: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds_qwen: torch.Tensor | None = None,
+ prompt_embeds_clip: torch.Tensor | None = None,
+ negative_prompt_embeds_qwen: torch.Tensor | None = None,
+ negative_prompt_embeds_clip: torch.Tensor | None = None,
+ prompt_cu_seqlens: torch.Tensor | None = None,
+ negative_prompt_cu_seqlens: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for text-to-image generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds`
instead. Ignored when not using guidance (`guidance_scale` < `1`).
height (`int`, defaults to `1024`):
@@ -577,7 +575,7 @@ def __call__(
Guidance scale as defined in classifier-free guidance.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A torch generator to make generation deterministic.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents.
diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_output.py b/src/diffusers/pipelines/kandinsky5/pipeline_output.py
index 2172ddff7e22..9b41d2f14424 100644
--- a/src/diffusers/pipelines/kandinsky5/pipeline_output.py
+++ b/src/diffusers/pipelines/kandinsky5/pipeline_output.py
@@ -11,8 +11,8 @@ class KandinskyPipelineOutput(BaseOutput):
Output class for kandinsky video pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
@@ -26,7 +26,7 @@ class KandinskyImagePipelineOutput(BaseOutput):
Output class for kandinsky image pipelines.
Args:
- image (`torch.Tensor`, `np.ndarray`, or List[PIL.Image.Image]):
+ image (`torch.Tensor`, `np.ndarray`, or list[PIL.Image.Image]):
List of image outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image. It can also be a NumPy array or Torch tensor of shape `(batch_size, channels, height,
width)`.
diff --git a/src/diffusers/pipelines/kolors/pipeline_kolors.py b/src/diffusers/pipelines/kolors/pipeline_kolors.py
index 3c7442afcaae..1e11faf8b9b6 100644
--- a/src/diffusers/pipelines/kolors/pipeline_kolors.py
+++ b/src/diffusers/pipelines/kolors/pipeline_kolors.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
@@ -63,10 +63,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -81,15 +81,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -199,21 +199,21 @@ def __init__(
def encode_prompt(
self,
prompt,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
negative_prompt=None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 256,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -221,7 +221,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -291,7 +291,7 @@ def encode_prompt(
if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
negative_prompt_embeds = torch.zeros_like(prompt_embeds)
elif do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -648,45 +648,43 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -702,11 +700,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -723,7 +721,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -732,7 +730,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -754,7 +752,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -768,31 +766,31 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -802,7 +800,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py b/src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py
index 8c3adf33b845..d9b519267216 100644
--- a/src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py
+++ b/src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -67,7 +67,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -82,10 +82,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -100,15 +100,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -219,21 +219,21 @@ def __init__(
def encode_prompt(
self,
prompt,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
negative_prompt=None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 256,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -241,7 +241,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -311,7 +311,7 @@ def encode_prompt(
if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
negative_prompt_embeds = torch.zeros_like(prompt_embeds)
elif do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -780,51 +780,49 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
strength: float = 0.3,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_start: Optional[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_start: float | None = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`):
+ image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[np.ndarray]`):
The image(s) to modify with the pipeline.
strength (`float`, *optional*, defaults to 0.3):
Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`
@@ -846,11 +844,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -874,7 +872,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -883,7 +881,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -905,7 +903,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -919,31 +917,31 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -953,7 +951,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/kolors/pipeline_output.py b/src/diffusers/pipelines/kolors/pipeline_output.py
index 310ee7e8a89b..0418191cc6e5 100644
--- a/src/diffusers/pipelines/kolors/pipeline_output.py
+++ b/src/diffusers/pipelines/kolors/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class KolorsPipelineOutput(BaseOutput):
Output class for Kolors pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/kolors/text_encoder.py b/src/diffusers/pipelines/kolors/text_encoder.py
index 6fd17156a116..434f4fed6fbb 100644
--- a/src/diffusers/pipelines/kolors/text_encoder.py
+++ b/src/diffusers/pipelines/kolors/text_encoder.py
@@ -13,7 +13,6 @@
# limitations under the License.
import math
-from typing import List, Optional, Tuple
import torch
import torch.nn.functional as F
@@ -234,7 +233,7 @@ def split_tensor_along_last_dim(
tensor: torch.Tensor,
num_partitions: int,
contiguous_split_chunks: bool = False,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
"""Split a tensor along its last dimension.
Arguments:
@@ -566,8 +565,8 @@ def forward(
attention_mask,
rotary_pos_emb,
kv_caches=None,
- use_cache: Optional[bool] = True,
- output_hidden_states: Optional[bool] = False,
+ use_cache: bool | None = True,
+ output_hidden_states: bool | None = False,
):
if not kv_caches:
kv_caches = [None for _ in range(self.num_layers)]
@@ -782,6 +781,9 @@ def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
self.prefix_encoder = PrefixEncoder(config)
self.dropout = torch.nn.Dropout(0.1)
+ if hasattr(self, "post_init"):
+ self.post_init()
+
def get_input_embeddings(self):
return self.embedding.word_embeddings
@@ -799,19 +801,19 @@ def get_prompt(self, batch_size, device, dtype=torch.half):
def forward(
self,
input_ids,
- position_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.BoolTensor] = None,
- full_attention_mask: Optional[torch.BoolTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
+ position_ids: torch.Tensor | None = None,
+ attention_mask: torch.BoolTensor | None = None,
+ full_attention_mask: torch.BoolTensor | None = None,
+ past_key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
):
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
- use_cache = use_cache if use_cache is not None else self.config.use_cache
+ use_cache = use_cache if use_cache is not None else getattr(self.config, "use_cache", None)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
batch_size, seq_length = input_ids.shape
diff --git a/src/diffusers/pipelines/kolors/tokenizer.py b/src/diffusers/pipelines/kolors/tokenizer.py
index b824ba12e079..c0a1831bdeee 100644
--- a/src/diffusers/pipelines/kolors/tokenizer.py
+++ b/src/diffusers/pipelines/kolors/tokenizer.py
@@ -15,7 +15,6 @@
import json
import os
import re
-from typing import Dict, List, Optional, Union
from sentencepiece import SentencePieceProcessor
from transformers import PreTrainedTokenizer
@@ -61,7 +60,7 @@ def tokenize(self, s: str, encode_special_tokens=False):
else:
return self.sp_model.EncodeAsPieces(s)
- def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
+ def encode(self, s: str, bos: bool = False, eos: bool = False) -> list[int]:
assert isinstance(s, str)
t = self.sp_model.encode(s)
if bos:
@@ -70,7 +69,7 @@ def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
t = t + [self.eos_id]
return t
- def decode(self, t: List[int]) -> str:
+ def decode(self, t: list[int]) -> str:
text, buffer = "", []
for token in t:
if token in self.index_special_tokens:
@@ -84,7 +83,7 @@ def decode(self, t: List[int]) -> str:
text += self.sp_model.decode(buffer)
return text
- def decode_tokens(self, tokens: List[str]) -> str:
+ def decode_tokens(self, tokens: list[str]) -> str:
text = self.sp_model.DecodePieces(tokens)
return text
@@ -192,7 +191,7 @@ def _convert_id_to_token(self, index):
"""Converts an index (integer) in a token (str) using the vocab."""
return self.tokenizer.convert_id_to_token(index)
- def convert_tokens_to_string(self, tokens: List[str]) -> str:
+ def convert_tokens_to_string(self, tokens: list[str]) -> str:
return self.tokenizer.decode_tokens(tokens)
def save_vocabulary(self, save_directory, filename_prefix=None):
@@ -206,7 +205,7 @@ def save_vocabulary(self, save_directory, filename_prefix=None):
An optional prefix to add to the named of the saved files.
Returns:
- `Tuple(str)`: Paths to the files saved.
+ `tuple(str)`: Paths to the files saved.
"""
if os.path.isdir(save_directory):
vocab_file = os.path.join(save_directory, self.vocab_files_names["vocab_file"])
@@ -246,8 +245,8 @@ def build_chat_input(self, query, history=None, role="user"):
return self.batch_encode_plus([input_ids], return_tensors="pt", is_split_into_words=True)
def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
+ self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+ ) -> list[int]:
"""
Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. A BERT sequence has the following format:
@@ -256,13 +255,13 @@ def build_inputs_with_special_tokens(
- pair of sequences: `[CLS] A [SEP] B [SEP]`
Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
+ token_ids_0 (`list[int]`):
+ list of IDs to which the special tokens will be added.
+ token_ids_1 (`list[int]`, *optional*):
Optional second list of IDs for sequence pairs.
Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+ `list[int]`: list of [input IDs](../glossary#input-ids) with the appropriate special tokens.
"""
prefix_tokens = self.get_prefix_tokens()
token_ids_0 = prefix_tokens + token_ids_0
@@ -272,19 +271,19 @@ def build_inputs_with_special_tokens(
def _pad(
self,
- encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
- max_length: Optional[int] = None,
+ encoded_inputs: dict[str, EncodedInput] | BatchEncoding,
+ max_length: int | None = None,
padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
- pad_to_multiple_of: Optional[int] = None,
- return_attention_mask: Optional[bool] = None,
- padding_side: Optional[bool] = None,
+ pad_to_multiple_of: int | None = None,
+ return_attention_mask: bool | None = None,
+ padding_side: bool | None = None,
) -> dict:
"""
Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
Args:
encoded_inputs:
- Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
+ Dictionary of tokenized inputs (`list[int]`) or batch of tokenized inputs (`list[list[int]]`).
max_length: maximum length of the returned list and optionally padding length (see below).
Will truncate by taking into account the special tokens.
padding_strategy: PaddingStrategy to use for padding.
diff --git a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py
index c28e358c51b6..6861438e4c63 100644
--- a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py
+++ b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py
@@ -16,7 +16,7 @@
# and https://github.com/hojonathanho/diffusion
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -53,7 +53,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -68,10 +68,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -86,15 +86,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -208,7 +208,7 @@ def __init__(
scheduler: LCMScheduler,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -245,16 +245,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -262,7 +262,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -361,7 +361,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -635,10 +635,10 @@ def get_timesteps(self, num_inference_steps, strength, device):
def check_inputs(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
strength: float,
callback_steps: int,
- prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
ip_adapter_image=None,
ip_adapter_image_embeds=None,
callback_on_step_end_tensor_inputs=None,
@@ -710,32 +710,32 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
num_inference_steps: int = 4,
strength: float = 0.8,
original_inference_steps: int = None,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 8.5,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -749,7 +749,7 @@ def __call__(
we will draw `num_inference_steps` evenly spaced timesteps from as our final timestep schedule,
following the Skipping-Step method in the paper (see Section 4.3). If not set this will default to the
scheduler's `original_inference_steps` attribute.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps on the original LCM training/distillation timestep schedule are used. Must be in descending
order.
@@ -761,7 +761,7 @@ def __call__(
0`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -773,7 +773,7 @@ def __call__(
provided, text embeddings are generated from the `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -794,7 +794,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py
index bc71d7bd171a..60f59ec7f9d3 100644
--- a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py
+++ b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py
@@ -16,7 +16,7 @@
# and https://github.com/hojonathanho/diffusion
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
@@ -73,10 +73,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -91,15 +91,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -187,7 +187,7 @@ def __init__(
scheduler: LCMScheduler,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -230,16 +230,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -247,7 +247,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -346,7 +346,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -565,11 +565,11 @@ def prepare_extra_step_kwargs(self, generator, eta):
# Currently StableDiffusionPipeline.check_inputs with negative prompt stuff removed
def check_inputs(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
height: int,
width: int,
callback_steps: int,
- prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
ip_adapter_image=None,
ip_adapter_image_embeds=None,
callback_on_step_end_tensor_inputs=None,
@@ -641,32 +641,32 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 4,
original_inference_steps: int = None,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 8.5,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -680,7 +680,7 @@ def __call__(
we will draw `num_inference_steps` evenly spaced timesteps from as our final timestep schedule,
following the Skipping-Step method in the paper (see Section 4.3). If not set this will default to the
scheduler's `original_inference_steps` attribute.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps on the original LCM training/distillation timestep schedule are used. Must be in descending
order.
@@ -692,7 +692,7 @@ def __call__(
0`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -704,7 +704,7 @@ def __call__(
provided, text embeddings are generated from the `prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -725,7 +725,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
index f1bf4701e31f..ec43988f9389 100644
--- a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
+++ b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Tuple, Union
import torch
import torch.nn as nn
@@ -62,11 +61,11 @@ class LDMTextToImagePipeline(DiffusionPipeline):
def __init__(
self,
- vqvae: Union[VQModel, AutoencoderKL],
+ vqvae: VQModel | AutoencoderKL,
bert: PreTrainedModel,
tokenizer: PreTrainedTokenizer,
- unet: Union[UNet2DModel, UNet2DConditionModel],
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
+ unet: UNet2DModel | UNet2DConditionModel,
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler,
):
super().__init__()
self.register_modules(vqvae=vqvae, bert=bert, tokenizer=tokenizer, unet=unet, scheduler=scheduler)
@@ -75,23 +74,23 @@ def __init__(
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
- height: Optional[int] = None,
- width: Optional[int] = None,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 1.0,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ prompt: str | list[str],
+ height: int | None = None,
+ width: int | None = None,
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 1.0,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
**kwargs,
- ) -> Union[Tuple, ImagePipelineOutput]:
+ ) -> tuple | ImagePipelineOutput:
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -236,7 +235,7 @@ def __call__(
logger = logging.get_logger(__name__)
-LDMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = [
+LDMBERT_PRETRAINED_MODEL_ARCHIVE_list = [
"ldm-bert",
# See all LDMBert models at https://huggingface.co/models?filter=ldmbert
]
@@ -297,7 +296,7 @@ def __init__(
super().__init__(pad_token_id=pad_token_id, **kwargs)
-def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
"""
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
"""
@@ -345,12 +344,12 @@ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
def forward(
self,
hidden_states: torch.Tensor,
- key_value_states: Optional[torch.Tensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- layer_head_mask: Optional[torch.Tensor] = None,
+ key_value_states: torch.Tensor | None = None,
+ past_key_value: tuple[torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ layer_head_mask: torch.Tensor | None = None,
output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
"""Input shape: Batch x Time x Channel"""
# if key_value_states are provided this layer is used as a cross-attention layer
@@ -382,10 +381,10 @@ def forward(
value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
if self.is_decoder:
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+ # if cross_attention save tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
# Further calls to cross_attention layer can then reuse all cross-attention
# key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+ # if uni-directional self-attention (decoder) save tuple(torch.Tensor, torch.Tensor) of
# all previous decoder key/value_states. Further calls to uni-directional self-attention
# can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
# if encoder bi-directional self-attention `past_key_value` is always `None`
@@ -479,8 +478,8 @@ def forward(
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
layer_head_mask: torch.Tensor,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
"""
Args:
hidden_states (`torch.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)`
@@ -591,14 +590,14 @@ def set_input_embeddings(self, value):
def forward(
self,
input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ head_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | BaseModelOutput:
r"""
Args:
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
diff --git a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py
index 631539e5c667..18cb8274f9b5 100644
--- a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py
+++ b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py
@@ -1,5 +1,4 @@
import inspect
-from typing import List, Optional, Tuple, Union
import numpy as np
import PIL.Image
@@ -59,14 +58,12 @@ def __init__(
self,
vqvae: VQModel,
unet: UNet2DModel,
- scheduler: Union[
- DDIMScheduler,
- PNDMScheduler,
- LMSDiscreteScheduler,
- EulerDiscreteScheduler,
- EulerAncestralDiscreteScheduler,
- DPMSolverMultistepScheduler,
- ],
+ scheduler: DDIMScheduler
+ | PNDMScheduler
+ | LMSDiscreteScheduler
+ | EulerDiscreteScheduler
+ | EulerAncestralDiscreteScheduler
+ | DPMSolverMultistepScheduler,
):
super().__init__()
self.register_modules(vqvae=vqvae, unet=unet, scheduler=scheduler)
@@ -74,14 +71,14 @@ def __init__(
@torch.no_grad()
def __call__(
self,
- image: Union[torch.Tensor, PIL.Image.Image] = None,
- batch_size: Optional[int] = 1,
- num_inference_steps: Optional[int] = 100,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ image: torch.Tensor | PIL.Image.Image = None,
+ batch_size: int | None = 1,
+ num_inference_steps: int | None = 100,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- ) -> Union[Tuple, ImagePipelineOutput]:
+ ) -> tuple | ImagePipelineOutput:
r"""
The call function to the pipeline for generation.
@@ -96,7 +93,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
output_type (`str`, *optional*, defaults to `"pil"`):
diff --git a/src/diffusers/pipelines/latte/pipeline_latte.py b/src/diffusers/pipelines/latte/pipeline_latte.py
index 7fde18e4fbbb..eed7762cebf1 100644
--- a/src/diffusers/pipelines/latte/pipeline_latte.py
+++ b/src/diffusers/pipelines/latte/pipeline_latte.py
@@ -18,7 +18,7 @@
import re
import urllib.parse as ul
from dataclasses import dataclass
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import T5EncoderModel, T5Tokenizer
@@ -80,10 +80,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -98,15 +98,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -205,13 +205,13 @@ def mask_text_embeddings(self, emb, mask):
# Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
clean_caption: bool = False,
mask_feature: bool = True,
dtype=None,
@@ -220,9 +220,9 @@ def encode_prompt(
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
Latte, this should be "".
@@ -614,46 +614,44 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
guidance_scale: float = 7.5,
num_images_per_prompt: int = 1,
video_length: int = 16,
height: int = 512,
width: int = 512,
eta: float = 0.0,
- 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,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
output_type: str = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
clean_caption: bool = True,
mask_feature: bool = True,
enable_temporal_attentions: bool = True,
decode_chunk_size: int = 14,
- ) -> Union[LattePipelineOutput, Tuple]:
+ ) -> LattePipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the video generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality video at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 7.0):
@@ -673,7 +671,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -691,9 +689,9 @@ def __call__(
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
- callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+ callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
A callback function or a list of callback functions to be called at the end of each denoising step.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
A list of tensor inputs that should be passed to the callback function. If not defined, all tensor
inputs will be passed.
clean_caption (`bool`, *optional*, defaults to `True`):
diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py
index fbf4dc23d043..864f9feeb5aa 100644
--- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py
+++ b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py
@@ -1,7 +1,7 @@
import inspect
import math
from itertools import repeat
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
import torch.nn.functional as F
@@ -115,7 +115,7 @@ def get_attention(self, step: int):
return attention
def aggregate_attention(
- self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int
+ self, attention_maps, prompts, res: int | tuple[int], from_where: list[str], is_cross: bool, select: int
):
out = [[] for x in range(self.batch_size)]
if isinstance(res, int):
@@ -309,7 +309,7 @@ def __init__(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
- scheduler: Union[DDIMScheduler, DPMSolverMultistepScheduler],
+ scheduler: DDIMScheduler | DPMSolverMultistepScheduler,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = True,
@@ -525,10 +525,10 @@ def encode_prompt(
enable_edit_guidance,
negative_prompt=None,
editing_prompt=None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- editing_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ editing_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
@@ -540,11 +540,11 @@ def encode_prompt(
number of images that should be generated per prompt
enable_edit_guidance (`bool`):
whether to perform any editing or reconstruct the input image instead
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- editing_prompt (`str` or `List[str]`, *optional*):
+ editing_prompt (`str` or `list[str]`, *optional*):
Editing prompt(s) to be encoded. If not defined, one has to pass `editing_prompt_embeds` instead.
editing_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
@@ -574,7 +574,7 @@ def encode_prompt(
num_edit_tokens = None
if negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif isinstance(negative_prompt, str):
@@ -774,29 +774,29 @@ def disable_vae_tiling(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- editing_prompt: Optional[Union[str, List[str]]] = None,
- editing_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- reverse_editing_direction: Optional[Union[bool, List[bool]]] = False,
- edit_guidance_scale: Optional[Union[float, List[float]]] = 5,
- edit_warmup_steps: Optional[Union[int, List[int]]] = 0,
- edit_cooldown_steps: Optional[Union[int, List[int]]] = None,
- edit_threshold: Optional[Union[float, List[float]]] = 0.9,
- user_mask: Optional[torch.Tensor] = None,
- sem_guidance: Optional[List[torch.Tensor]] = None,
+ editing_prompt: str | list[str] | None = None,
+ editing_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ reverse_editing_direction: bool | list[bool] | None = False,
+ edit_guidance_scale: float | list[float] | None = 5,
+ edit_warmup_steps: int | list[int] | None = 0,
+ edit_cooldown_steps: int | list[int] | None = None,
+ edit_threshold: float | list[float] | None = 0.9,
+ user_mask: torch.Tensor | None = None,
+ sem_guidance: list[torch.Tensor] | None = None,
use_cross_attn_mask: bool = False,
use_intersect_mask: bool = True,
- attn_store_steps: Optional[List[int]] = [],
+ attn_store_steps: list[int] | None = [],
store_averaged_over_steps: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
@@ -805,7 +805,7 @@ def __call__(
always be performed for the last inverted image(s).
Args:
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
generator (`torch.Generator`, *optional*):
@@ -817,7 +817,7 @@ def __call__(
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] instead of a plain
tuple.
- editing_prompt (`str` or `List[str]`, *optional*):
+ editing_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. The image is reconstructed by setting
`editing_prompt = None`. Guidance direction of prompt should be specified via
`reverse_editing_direction`.
@@ -827,25 +827,25 @@ def __call__(
negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
- reverse_editing_direction (`bool` or `List[bool]`, *optional*, defaults to `False`):
+ reverse_editing_direction (`bool` or `list[bool]`, *optional*, defaults to `False`):
Whether the corresponding prompt in `editing_prompt` should be increased or decreased.
- edit_guidance_scale (`float` or `List[float]`, *optional*, defaults to 5):
+ edit_guidance_scale (`float` or `list[float]`, *optional*, defaults to 5):
Guidance scale for guiding the image generation. If provided as list values should correspond to
`editing_prompt`. `edit_guidance_scale` is defined as `s_e` of equation 12 of [LEDITS++
Paper](https://huggingface.co/papers/2301.12247).
- edit_warmup_steps (`float` or `List[float]`, *optional*, defaults to 10):
+ edit_warmup_steps (`float` or `list[float]`, *optional*, defaults to 10):
Number of diffusion steps (for each prompt) for which guidance will not be applied.
- edit_cooldown_steps (`float` or `List[float]`, *optional*, defaults to `None`):
+ edit_cooldown_steps (`float` or `list[float]`, *optional*, defaults to `None`):
Number of diffusion steps (for each prompt) after which guidance will no longer be applied.
- edit_threshold (`float` or `List[float]`, *optional*, defaults to 0.9):
+ edit_threshold (`float` or `list[float]`, *optional*, defaults to 0.9):
Masking threshold of guidance. Threshold should be proportional to the image region that is modified.
'edit_threshold' is defined as 'λ' of equation 12 of [LEDITS++
Paper](https://huggingface.co/papers/2301.12247).
user_mask (`torch.Tensor`, *optional*):
User-provided mask for even better control over the editing process. This is helpful when LEDITS++'s
implicit masks do not meet user preferences.
- sem_guidance (`List[torch.Tensor]`, *optional*):
- List of pre-generated guidance vectors to be applied at generation. Length of the list has to
+ sem_guidance (`list[torch.Tensor]`, *optional*):
+ list of pre-generated guidance vectors to be applied at generation. Length of the list has to
correspond to `num_inference_steps`.
use_cross_attn_mask (`bool`, defaults to `False`):
Whether cross-attention masks are used. Cross-attention masks are always used when use_intersect_mask
@@ -855,7 +855,7 @@ def __call__(
Whether the masking term is calculated as intersection of cross-attention masks and masks derived from
the noise estimate. Cross-attention mask are defined as 'M^1' and masks derived from the noise estimate
are defined as 'M^2' of equation 12 of [LEDITS++ paper](https://huggingface.co/papers/2311.16711).
- attn_store_steps (`List[int]`, *optional*):
+ attn_store_steps (`list[int]`, *optional*):
Steps for which the attention maps are stored in the AttentionStore. Just for visualization purposes.
store_averaged_over_steps (`bool`, defaults to `True`):
Whether the attention maps for the 'attn_store_steps' are stored averaged over the diffusion steps. If
@@ -875,7 +875,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
@@ -1282,13 +1282,13 @@ def invert(
source_guidance_scale: float = 3.5,
num_inversion_steps: int = 30,
skip: float = 0.15,
- generator: Optional[torch.Generator] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- resize_mode: Optional[str] = "default",
- crops_coords: Optional[Tuple[int, int, int, int]] = None,
+ generator: torch.Generator | None = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ resize_mode: str | None = "default",
+ crops_coords: tuple[int, int, int, int] | None = None,
):
r"""
The function to the pipeline for image inversion as described by the [LEDITS++
@@ -1331,7 +1331,7 @@ def invert(
image to fit within the specified width and height, maintaining the aspect ratio, and then center the
image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
supported for PIL image input.
- crops_coords (`List[Tuple[int, int, int, int]]`, *optional*, defaults to `None`):
+ crops_coords (`list[tuple[int, int, int, int]]`, *optional*, defaults to `None`):
The crop coordinates for each image in the batch. If `None`, will not crop the image.
Returns:
diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py
index 993957a052fc..a136770b9f26 100644
--- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py
+++ b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
import torch.nn.functional as F
@@ -145,7 +145,7 @@ def get_attention(self, step: int):
return attention
def aggregate_attention(
- self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int
+ self, attention_maps, prompts, res: int | tuple[int], from_where: list[str], is_cross: bool, select: int
):
out = [[] for x in range(self.batch_size)]
if isinstance(res, int):
@@ -351,11 +351,11 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- scheduler: Union[DPMSolverMultistepScheduler, DDIMScheduler],
+ scheduler: DPMSolverMultistepScheduler | DDIMScheduler,
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
):
super().__init__()
@@ -399,18 +399,18 @@ def __init__(
def encode_prompt(
self,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
enable_edit_guidance: bool = True,
- editing_prompt: Optional[str] = None,
- editing_prompt_embeds: Optional[torch.Tensor] = None,
- editing_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ editing_prompt: str | None = None,
+ editing_prompt_embeds: torch.Tensor | None = None,
+ editing_pooled_prompt_embeds: torch.Tensor | None = None,
) -> object:
r"""
Encodes the prompt into text encoder hidden states.
@@ -420,10 +420,10 @@ def encode_prompt(
torch device
num_images_per_prompt (`int`):
number of images that should be generated per prompt
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead.
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
negative_prompt_embeds (`torch.Tensor`, *optional*):
@@ -441,7 +441,7 @@ def encode_prompt(
the output of the pre-final layer will be used for computing the prompt embeddings.
enable_edit_guidance (`bool`):
Whether to guide towards an editing prompt or not.
- editing_prompt (`str` or `List[str]`, *optional*):
+ editing_prompt (`str` or `list[str]`, *optional*):
Editing prompt(s) to be encoded. If not defined and 'enable_edit_guidance' is True, one has to pass
`editing_prompt_embeds` instead.
editing_prompt_embeds (`torch.Tensor`, *optional*):
@@ -495,7 +495,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if batch_size != len(negative_prompt):
raise ValueError(
@@ -837,35 +837,35 @@ def prepare_unet(self, attention_store, PnP: bool = False):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- denoising_end: Optional[float] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ denoising_end: float | None = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- editing_prompt: Optional[Union[str, List[str]]] = None,
- editing_prompt_embeddings: Optional[torch.Tensor] = None,
- editing_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- reverse_editing_direction: Optional[Union[bool, List[bool]]] = False,
- edit_guidance_scale: Optional[Union[float, List[float]]] = 5,
- edit_warmup_steps: Optional[Union[int, List[int]]] = 0,
- edit_cooldown_steps: Optional[Union[int, List[int]]] = None,
- edit_threshold: Optional[Union[float, List[float]]] = 0.9,
- sem_guidance: Optional[List[torch.Tensor]] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
+ editing_prompt: str | list[str] | None = None,
+ editing_prompt_embeddings: torch.Tensor | None = None,
+ editing_pooled_prompt_embeds: torch.Tensor | None = None,
+ reverse_editing_direction: bool | list[bool] | None = False,
+ edit_guidance_scale: float | list[float] | None = 5,
+ edit_warmup_steps: int | list[int] | None = 0,
+ edit_cooldown_steps: int | list[int] | None = None,
+ edit_threshold: float | list[float] | None = 0.9,
+ sem_guidance: list[torch.Tensor] | None = None,
use_cross_attn_mask: bool = False,
use_intersect_mask: bool = False,
- user_mask: Optional[torch.Tensor] = None,
- attn_store_steps: Optional[List[int]] = [],
+ user_mask: torch.Tensor | None = None,
+ attn_store_steps: list[int] | None = [],
store_averaged_over_steps: bool = True,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
@@ -880,11 +880,11 @@ def __call__(
still retain a substantial amount of noise as determined by the discrete timesteps selected by the
scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
"Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
negative_prompt_embeds (`torch.Tensor`, *optional*):
@@ -919,16 +919,16 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- editing_prompt (`str` or `List[str]`, *optional*):
+ editing_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. The image is reconstructed by setting
`editing_prompt = None`. Guidance direction of prompt should be specified via
`reverse_editing_direction`.
@@ -939,22 +939,22 @@ def __call__(
Pre-generated pooled edit text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, editing_prompt_embeddings will be generated from `editing_prompt` input
argument.
- reverse_editing_direction (`bool` or `List[bool]`, *optional*, defaults to `False`):
+ reverse_editing_direction (`bool` or `list[bool]`, *optional*, defaults to `False`):
Whether the corresponding prompt in `editing_prompt` should be increased or decreased.
- edit_guidance_scale (`float` or `List[float]`, *optional*, defaults to 5):
+ edit_guidance_scale (`float` or `list[float]`, *optional*, defaults to 5):
Guidance scale for guiding the image generation. If provided as list values should correspond to
`editing_prompt`. `edit_guidance_scale` is defined as `s_e` of equation 12 of [LEDITS++
Paper](https://huggingface.co/papers/2301.12247).
- edit_warmup_steps (`float` or `List[float]`, *optional*, defaults to 10):
+ edit_warmup_steps (`float` or `list[float]`, *optional*, defaults to 10):
Number of diffusion steps (for each prompt) for which guidance is not applied.
- edit_cooldown_steps (`float` or `List[float]`, *optional*, defaults to `None`):
+ edit_cooldown_steps (`float` or `list[float]`, *optional*, defaults to `None`):
Number of diffusion steps (for each prompt) after which guidance is no longer applied.
- edit_threshold (`float` or `List[float]`, *optional*, defaults to 0.9):
+ edit_threshold (`float` or `list[float]`, *optional*, defaults to 0.9):
Masking threshold of guidance. Threshold should be proportional to the image region that is modified.
'edit_threshold' is defined as 'λ' of equation 12 of [LEDITS++
Paper](https://huggingface.co/papers/2301.12247).
- sem_guidance (`List[torch.Tensor]`, *optional*):
- List of pre-generated guidance vectors to be applied at generation. Length of the list has to
+ sem_guidance (`list[torch.Tensor]`, *optional*):
+ list of pre-generated guidance vectors to be applied at generation. Length of the list has to
correspond to `num_inference_steps`.
use_cross_attn_mask:
Whether cross-attention masks are used. Cross-attention masks are always used when use_intersect_mask
@@ -980,7 +980,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
@@ -1482,14 +1482,14 @@ def invert(
negative_prompt_2: str = None,
num_inversion_steps: int = 50,
skip: float = 0.15,
- generator: Optional[torch.Generator] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
+ generator: torch.Generator | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
num_zero_noise_steps: int = 3,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- resize_mode: Optional[str] = "default",
- crops_coords: Optional[Tuple[int, int, int, int]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ resize_mode: str | None = "default",
+ crops_coords: tuple[int, int, int, int] | None = None,
):
r"""
The function to the pipeline for image inversion as described by the [LEDITS++
@@ -1505,11 +1505,11 @@ def invert(
if the `source_prompt` is `""`.
source_guidance_scale (`float`, defaults to `3.5`):
Strength of guidance during inversion.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_inversion_steps (`int`, defaults to `50`):
@@ -1520,7 +1520,7 @@ def invert(
generator (`torch.Generator`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make inversion
deterministic.
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_output.py b/src/diffusers/pipelines/ledits_pp/pipeline_output.py
index 756be82b0069..4a69e7442bdc 100644
--- a/src/diffusers/pipelines/ledits_pp/pipeline_output.py
+++ b/src/diffusers/pipelines/ledits_pp/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -13,16 +12,16 @@ class LEditsPPDiffusionPipelineOutput(BaseOutput):
Output class for LEdits++ Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
- nsfw_content_detected (`List[bool]`)
- List indicating whether the corresponding generated image contains “not-safe-for-work” (nsfw) content or
+ nsfw_content_detected (`list[bool]`)
+ list indicating whether the corresponding generated image contains “not-safe-for-work” (nsfw) content or
`None` if safety checking could not be performed.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_content_detected: Optional[List[bool]]
+ images: list[PIL.Image.Image] | np.ndarray
+ nsfw_content_detected: list[bool] | None
@dataclass
@@ -31,13 +30,13 @@ class LEditsPPInversionPipelineOutput(BaseOutput):
Output class for LEdits++ Diffusion pipelines.
Args:
- input_images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of the cropped and resized input images as PIL images of length `batch_size` or NumPy array of shape `
+ input_images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of the cropped and resized input images as PIL images of length `batch_size` or NumPy array of shape `
(batch_size, height, width, num_channels)`.
- vae_reconstruction_images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of VAE reconstruction of all input images as PIL images of length `batch_size` or NumPy array of shape
+ vae_reconstruction_images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of VAE reconstruction of all input images as PIL images of length `batch_size` or NumPy array of shape
` (batch_size, height, width, num_channels)`.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- vae_reconstruction_images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
+ vae_reconstruction_images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py b/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py
index ca28422f9ca0..19720d7bbab8 100644
--- a/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py
+++ b/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
import re
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
import numpy as np
import torch
@@ -145,10 +145,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -163,15 +163,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -271,7 +271,7 @@ def rewire_prompt(self, prompt, device):
rewrite_prompt = output_text
return rewrite_prompt
- def _encode_prompt(self, prompt: List[str]):
+ def _encode_prompt(self, prompt: list[str]):
batch_all_tokens = []
for each_prompt in prompt:
@@ -334,9 +334,9 @@ def _encode_prompt(self, prompt: List[str]):
def encode_prompt(
self,
- prompt: Union[str, List[str]] = None,
- num_images_per_prompt: Optional[int] = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
+ prompt: str | list[str] = None,
+ num_images_per_prompt: int | None = 1,
+ prompt_embeds: torch.Tensor | None = None,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
batch_size = len(prompt)
@@ -472,24 +472,24 @@ def check_inputs(
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 4.5,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- enable_cfg_renorm: Optional[bool] = True,
- cfg_renorm_min: Optional[float] = 0.0,
- enable_prompt_rewrite: Optional[bool] = True,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ enable_cfg_renorm: bool | None = True,
+ cfg_renorm_min: float | None = 0.0,
+ enable_prompt_rewrite: bool | None = True,
):
r"""
Function invoked when calling the pipeline for generation.
diff --git a/src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py b/src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py
index e55a2a47f343..69d5d82f18ec 100644
--- a/src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py
+++ b/src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py
@@ -14,7 +14,7 @@
import inspect
import math
import re
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
import numpy as np
import PIL
@@ -144,10 +144,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -162,15 +162,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -203,7 +203,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -343,10 +343,10 @@ def _encode_prompt(self, prompt, image):
def encode_prompt(
self,
- prompt: List[str] = None,
- image: Optional[torch.Tensor] = None,
- num_images_per_prompt: Optional[int] = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
+ prompt: list[str] = None,
+ image: torch.Tensor | None = None,
+ num_images_per_prompt: int | None = 1,
+ prompt_embeds: torch.Tensor | None = None,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
batch_size = len(prompt)
@@ -532,20 +532,20 @@ def check_inputs(
@torch.no_grad()
def __call__(
self,
- image: Optional[PIL.Image.Image] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ image: PIL.Image.Image | None = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 4.5,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
diff --git a/src/diffusers/pipelines/longcat_image/pipeline_output.py b/src/diffusers/pipelines/longcat_image/pipeline_output.py
index e3c25f1cbfa7..290178eb1a88 100644
--- a/src/diffusers/pipelines/longcat_image/pipeline_output.py
+++ b/src/diffusers/pipelines/longcat_image/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class LongCatImagePipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image, np.ndarray]
diff --git a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py
index 6dce792a2b43..f579cf00dbe7 100644
--- a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py
+++ b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional
-
import torch
from ...configuration_utils import ConfigMixin, register_to_config
@@ -21,7 +19,7 @@
class ResBlock(torch.nn.Module):
- def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3):
+ def __init__(self, channels: int, mid_channels: int | None = None, dims: int = 3):
super().__init__()
if mid_channels is None:
mid_channels = channels
diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx.py b/src/diffusers/pipelines/ltx/pipeline_ltx.py
index 3c90da1c7051..e2514c3bca24 100644
--- a/src/diffusers/pipelines/ltx/pipeline_ltx.py
+++ b/src/diffusers/pipelines/ltx/pipeline_ltx.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -83,10 +83,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -101,15 +101,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -233,11 +233,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -282,25 +282,25 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.mochi.pipeline_mochi.MochiPipeline.encode_prompt with 256->128
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -478,10 +478,10 @@ def prepare_latents(
height: int = 512,
width: int = 704,
num_frames: int = 161,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -536,37 +536,37 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 704,
num_frames: int = 161,
frame_rate: int = 25,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 3,
guidance_rescale: float = 0.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- decode_timestep: Union[float, List[float]] = 0.0,
- decode_noise_scale: Optional[Union[float, List[float]]] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ decode_timestep: float | list[float] = 0.0,
+ decode_noise_scale: float | list[float] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 128,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, defaults to `512`):
@@ -578,7 +578,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -596,7 +596,7 @@ def __call__(
using zero terminal SNR.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -631,7 +631,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_condition.py b/src/diffusers/pipelines/ltx/pipeline_ltx_condition.py
index 10c9432a7f46..539a28f56e67 100644
--- a/src/diffusers/pipelines/ltx/pipeline_ltx_condition.py
+++ b/src/diffusers/pipelines/ltx/pipeline_ltx_condition.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -100,7 +100,7 @@ class LTXVideoCondition:
Attributes:
image (`PIL.Image.Image`):
The image to condition the video on.
- video (`List[PIL.Image.Image]`):
+ video (`list[PIL.Image.Image]`):
The video to condition the video on.
frame_index (`int`):
The frame index at which the image or video will conditionally effect the video generation.
@@ -108,8 +108,8 @@ class LTXVideoCondition:
The strength of the conditioning effect. A value of `1.0` means the conditioning effect is fully applied.
"""
- image: Optional[PIL.Image.Image] = None
- video: Optional[List[PIL.Image.Image]] = None
+ image: PIL.Image.Image | None = None
+ video: list[PIL.Image.Image] | None = None
frame_index: int = 0
strength: float = 1.0
@@ -151,10 +151,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -169,15 +169,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -210,7 +210,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -319,11 +319,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -368,25 +368,25 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.mochi.pipeline_mochi.MochiPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -671,21 +671,21 @@ def add_noise_to_image_conditioning_latents(
def prepare_latents(
self,
- conditions: Optional[List[torch.Tensor]] = None,
- condition_strength: Optional[List[float]] = None,
- condition_frame_index: Optional[List[int]] = None,
+ conditions: list[torch.Tensor] | None = None,
+ condition_strength: list[float] | None = None,
+ condition_frame_index: list[int] | None = None,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 704,
num_frames: int = 161,
num_prefix_latent_frames: int = 2,
- sigma: Optional[torch.Tensor] = None,
- latents: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+ sigma: torch.Tensor | None = None,
+ latents: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
num_latent_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
latent_height = height // self.vae_spatial_compression_ratio
latent_width = width // self.vae_spatial_compression_ratio
@@ -849,61 +849,61 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- conditions: Union[LTXVideoCondition, List[LTXVideoCondition]] = None,
- image: Union[PipelineImageInput, List[PipelineImageInput]] = None,
- video: List[PipelineImageInput] = None,
- frame_index: Union[int, List[int]] = 0,
- strength: Union[float, List[float]] = 1.0,
+ conditions: LTXVideoCondition | list[LTXVideoCondition] = None,
+ image: PipelineImageInput | list[PipelineImageInput] = None,
+ video: list[PipelineImageInput] = None,
+ frame_index: int | list[int] = 0,
+ strength: float | list[float] = 1.0,
denoise_strength: float = 1.0,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 704,
num_frames: int = 161,
frame_rate: int = 25,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 3,
guidance_rescale: float = 0.0,
image_cond_noise_scale: float = 0.15,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- decode_timestep: Union[float, List[float]] = 0.0,
- decode_noise_scale: Optional[Union[float, List[float]]] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ decode_timestep: float | list[float] = 0.0,
+ decode_noise_scale: float | list[float] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- conditions (`List[LTXVideoCondition], *optional*`):
+ conditions (`list[LTXVideoCondition], *optional*`):
The list of frame-conditioning items for the video generation.If not provided, conditions will be
created using `image`, `video`, `frame_index` and `strength`.
- image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*):
+ image (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*):
The image or images to condition the video generation. If not provided, one has to pass `video` or
`conditions`.
- video (`List[PipelineImageInput]`, *optional*):
+ video (`list[PipelineImageInput]`, *optional*):
The video to condition the video generation. If not provided, one has to pass `image` or `conditions`.
- frame_index (`int` or `List[int]`, *optional*):
+ frame_index (`int` or `list[int]`, *optional*):
The frame index or frame indices at which the image or video will conditionally effect the video
generation. If not provided, one has to pass `conditions`.
- strength (`float` or `List[float]`, *optional*):
+ strength (`float` or `list[float]`, *optional*):
The strength or strengths of the conditioning effect. If not provided, one has to pass `conditions`.
denoise_strength (`float`, defaults to `1.0`):
The strength of the noise added to the latents for editing. Higher strength leads to more noise added
to the latents, therefore leading to more differences between original video and generated video. This
is useful for video-to-video editing.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, defaults to `512`):
@@ -915,7 +915,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -933,7 +933,7 @@ def __call__(
using zero terminal SNR.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -968,7 +968,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_i2v_long_multi_prompt.py b/src/diffusers/pipelines/ltx/pipeline_ltx_i2v_long_multi_prompt.py
index 7965bd3b4b87..838d5afc5c5a 100644
--- a/src/diffusers/pipelines/ltx/pipeline_ltx_i2v_long_multi_prompt.py
+++ b/src/diffusers/pipelines/ltx/pipeline_ltx_i2v_long_multi_prompt.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import copy
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -84,7 +86,7 @@ def get_latent_coords(
device: torch.device for the resulting tensor.
rope_interpolation_scale:
tuple[int|float, int|float, int|float]. Scale per (t, y, x) latent step to pixel coords.
- latent_idx: Optional[int]. When not None, shifts the time coordinate to align segments:
+ latent_idx: int | None. When not None, shifts the time coordinate to align segments:
- <= 0 uses step multiples of rope_interpolation_scale[0]
- > 0 starts at 1 then increments by rope_interpolation_scale[0]
@@ -141,7 +143,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
def adain_normalize_latents(
- curr_latents: torch.Tensor, ref_latents: Optional[torch.Tensor], factor: float
+ curr_latents: torch.Tensor, ref_latents: torch.Tensor | None, factor: float
) -> torch.Tensor:
"""
Optional AdaIN normalization: channel-wise mean/variance matching of curr_latents to ref_latents, controlled by
@@ -150,7 +152,7 @@ def adain_normalize_latents(
Args:
curr_latents: Tensor [B, C, T, H, W]. Current window latents.
ref_latents:
- Optional[Tensor] [B, C, T_ref, H, W]. Reference latents (e.g., first window) used to compute target stats.
+ Tensor | None [B, C, T_ref, H, W]. Reference latents (e.g., first window) used to compute target stats.
factor: float in [0, 1]. 0 keeps current stats; 1 matches reference stats.
Returns:
@@ -180,7 +182,7 @@ def adain_normalize_latents(
def split_into_temporal_windows(
latent_len: int, temporal_tile_size: int, temporal_overlap: int, compression: int
-) -> List[Tuple[int, int]]:
+) -> list[tuple[int, int]]:
"""
Split latent frames into sliding windows.
@@ -231,26 +233,26 @@ def linear_overlap_fuse(prev: torch.Tensor, new: torch.Tensor, overlap: int) ->
def inject_prev_tail_latents(
window_latents: torch.Tensor,
- prev_tail_latents: Optional[torch.Tensor],
+ prev_tail_latents: torch.Tensor | None,
window_cond_mask_5d: torch.Tensor,
overlap_lat: int,
- strength: Optional[float],
+ strength: float | None,
prev_overlap_len: int,
-) -> Tuple[torch.Tensor, torch.Tensor, int]:
+) -> tuple[torch.Tensor, torch.Tensor, int]:
"""
Inject the tail latents from the previous window at the beginning of the current window (first k frames), where k =
min(overlap_lat, T_curr, T_prev_tail).
Args:
window_latents: Tensor [B, C, T, H, W]. Current window latents.
- prev_tail_latents: Optional[Tensor] [B, C, T_prev, H, W]. Tail segment from the previous window.
+ prev_tail_latents: Tensor | None [B, C, T_prev, H, W]. Tail segment from the previous window.
window_cond_mask_5d: Tensor [B, 1, T, H, W]. Per-token conditioning mask (1 = free, 0 = hard condition).
overlap_lat: int. Number of latent frames to inject from the previous tail.
- strength: Optional[float] in [0, 1]. Blend strength; 1.0 replaces, 0.0 keeps original.
+ strength: float | None in [0, 1]. Blend strength; 1.0 replaces, 0.0 keeps original.
prev_overlap_len: int. Accumulated overlap length so far (used for trimming later).
Returns:
- Tuple[Tensor, Tensor, int]: (updated_window_latents, updated_cond_mask, updated_prev_overlap_len)
+ tuple[Tensor, Tensor, int]: (updated_window_latents, updated_cond_mask, updated_prev_overlap_len)
"""
if prev_tail_latents is None or overlap_lat <= 0 or strength is None or strength <= 0:
return window_latents, window_cond_mask_5d, prev_overlap_len
@@ -316,7 +318,7 @@ def build_video_coords_for_window(
return fractional_coords
-def parse_prompt_segments(prompt: Union[str, List[str]], prompt_segments: Optional[List[Dict[str, Any]]]) -> List[str]:
+def parse_prompt_segments(prompt: str | list[str], prompt_segments: list[dict[str, Any]] | None) -> list[str]:
"""
Return a list of positive prompts per window index.
@@ -368,7 +370,7 @@ def batch_normalize(latents, reference, factor):
reference: dict containing "samples" used to compute target stats
factor: float in [0, 1]; 0 = no change, 1 = full match to reference
Returns:
- Tuple[dict]: a single-element tuple with the updated latents dict.
+ tuple[dict]: a single-element tuple with the updated latents dict.
"""
latents_copy = copy.deepcopy(latents)
t = latents_copy["samples"] # B x C x F x H x W
@@ -500,11 +502,11 @@ def interrupt(self):
# Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -549,25 +551,25 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -697,13 +699,13 @@ def prepare_latents(
width: int,
num_frames: int,
device: torch.device,
- generator: Optional[torch.Generator],
+ generator: torch.Generator | None,
dtype: torch.dtype = torch.float32,
- latents: Optional[torch.Tensor] = None,
- cond_latents: Optional[torch.Tensor] = None,
+ latents: torch.Tensor | None = None,
+ cond_latents: torch.Tensor | None = None,
cond_strength: float = 0.0,
- negative_index_latents: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], int, int, int]:
+ negative_index_latents: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor | None, int, int, int]:
"""
Prepare base latents and optionally inject first-frame conditioning latents.
@@ -737,18 +739,18 @@ def prepare_latents(
def vae_decode_tiled(
self,
latents: torch.Tensor,
- decode_timestep: Optional[float] = None,
- decode_noise_scale: Optional[float] = None,
+ decode_timestep: float | None = None,
+ decode_noise_scale: float | None = None,
horizontal_tiles: int = 4,
vertical_tiles: int = 4,
overlap: int = 3,
last_frame_fix: bool = True,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
output_type: str = "pt",
auto_denormalize: bool = True,
compute_dtype: torch.dtype = torch.float32,
enable_vae_tiling: bool = False,
- ) -> Union[torch.Tensor, np.ndarray, List[PIL.Image.Image]]:
+ ) -> torch.Tensor | np.ndarray | list[PIL.Image.Image]:
"""
VAE-based spatial tiled decoding (ComfyUI parity) implemented in Diffusers style.
- Linearly feather and blend overlapping tiles to avoid seams.
@@ -837,8 +839,8 @@ def vae_decode_tiled(
base_tile_h = (h_lat + (vertical_tiles - 1) * overlap) // vertical_tiles
base_tile_w = (w_lat + (horizontal_tiles - 1) * overlap) // horizontal_tiles
- output: Optional[torch.Tensor] = None # [B, C_img, F, H, W], fused using compute_dtype
- weights: Optional[torch.Tensor] = None # [B, 1, F, H, W], fused using compute_dtype
+ output: torch.Tensor | None = None # [B, C_img, F, H, W], fused using compute_dtype
+ weights: torch.Tensor | None = None # [B, 1, F, H, W], fused using compute_dtype
# Iterate tiles in latent space (no temporal tiling)
for v in range(vertical_tiles):
@@ -934,52 +936,52 @@ def vae_decode_tiled(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_segments: Optional[List[Dict[str, Any]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_segments: list[dict[str, Any]] | None = None,
height: int = 512,
width: int = 704,
num_frames: int = 161,
frame_rate: float = 25,
guidance_scale: float = 1.0,
guidance_rescale: float = 0.0,
- num_inference_steps: Optional[int] = 8,
- sigmas: Optional[Union[List[float], torch.Tensor]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- seed: Optional[int] = 0,
- cond_image: Optional[Union["PIL.Image.Image", torch.Tensor]] = None,
+ num_inference_steps: int | None = 8,
+ sigmas: list[float, torch.Tensor] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ seed: int | None = 0,
+ cond_image: "PIL.Image.Image" | torch.Tensor | None = None,
cond_strength: float = 0.5,
- latents: Optional[torch.Tensor] = None,
+ latents: torch.Tensor | None = None,
temporal_tile_size: int = 80,
temporal_overlap: int = 24,
temporal_overlap_cond_strength: float = 0.5,
adain_factor: float = 0.25,
- guidance_latents: Optional[torch.Tensor] = None,
+ guidance_latents: torch.Tensor | None = None,
guiding_strength: float = 1.0,
- negative_index_latents: Optional[torch.Tensor] = None,
+ negative_index_latents: torch.Tensor | None = None,
negative_index_strength: float = 1.0,
- skip_steps_sigma_threshold: Optional[float] = 1,
- decode_timestep: Optional[float] = 0.05,
- decode_noise_scale: Optional[float] = 0.025,
+ skip_steps_sigma_threshold: float | None = 1,
+ decode_timestep: float | None = 0.05,
+ decode_noise_scale: float | None = 0.025,
decode_horizontal_tiles: int = 4,
decode_vertical_tiles: int = 4,
decode_overlap: int = 3,
- output_type: Optional[str] = "latent", # "latent" | "pt" | "np" | "pil"
+ output_type: str | None = "latent", # "latent" | "pt" | "np" | "pil"
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 128,
):
r"""
Generate an image-to-video sequence via temporal sliding windows and multi-prompt scheduling.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
Positive text prompt(s) per window. If a single string contains '|', parts are split by bars.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
Negative prompt(s) to suppress undesired content.
- prompt_segments (`List[dict]`, *optional*):
+ prompt_segments (`list[dict]`, *optional*):
Segment mapping with {"start_window", "end_window", "text"} to override prompts per window.
height (`int`, defaults to `512`):
Output image height in pixels; must be divisible by 32.
@@ -995,9 +997,9 @@ def __call__(
Optional rescale to mitigate overexposure under CFG (see `rescale_noise_cfg`).
num_inference_steps (`int`, *optional*, defaults to `8`):
Denoising steps per window. Ignored if `sigmas` is provided.
- sigmas (`List[float]` or `torch.Tensor`, *optional*):
+ sigmas (`list[float]` or `torch.Tensor`, *optional*):
Explicit sigma schedule per window; if set, overrides `num_inference_steps`.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
Controls stochasticity; list accepted but first element is used (batch=1).
seed (`int`, *optional*, defaults to `0`):
If provided, seeds the shared generator for global latents and derives a window-local generator with
@@ -1045,7 +1047,7 @@ def __call__(
Extra attention parameters forwarded to the transformer.
callback_on_step_end (`PipelineCallback` or `MultiPipelineCallbacks`, *optional*):
Per-step callback hook.
- callback_on_step_end_tensor_inputs (`List[str]`, defaults to `["latents"]`):
+ callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`):
Keys from locals() to pass into the callback.
max_sequence_length (`int`, defaults to `128`):
Tokenizer max length for prompt encoding.
@@ -1060,7 +1062,7 @@ def __call__(
- "latent"/"pt": `torch.Tensor` [B, C, F, H, W]; "latent" is in normalized latent space, "pt" is VAE
output space.
- "np": `np.ndarray` post-processed.
- - "pil": `List[PIL.Image.Image]` list of PIL images.
+ - "pil": `list[PIL.Image.Image]` list of PIL images.
Shapes:
Latent sizes (when auto-generated):
diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py b/src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py
index 3226b045cccb..497f505c4dd8 100644
--- a/src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py
+++ b/src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -88,10 +88,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -106,15 +106,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -147,7 +147,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -256,11 +256,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -305,25 +305,25 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.mochi.pipeline_mochi.MochiPipeline.encode_prompt with 256->128
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 128,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -501,16 +501,16 @@ def _denormalize_latents(
def prepare_latents(
self,
- image: Optional[torch.Tensor] = None,
+ image: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 704,
num_frames: int = 161,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
height = height // self.vae_spatial_compression_ratio
width = width // self.vae_spatial_compression_ratio
@@ -598,30 +598,30 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 704,
num_frames: int = 161,
frame_rate: int = 25,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 3,
guidance_rescale: float = 0.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- decode_timestep: Union[float, List[float]] = 0.0,
- decode_noise_scale: Optional[Union[float, List[float]]] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ decode_timestep: float | list[float] = 0.0,
+ decode_noise_scale: float | list[float] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 128,
):
r"""
@@ -630,7 +630,7 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, defaults to `512`):
@@ -642,7 +642,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -660,7 +660,7 @@ def __call__(
using zero terminal SNR.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -695,7 +695,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py
index 9acff105e56d..17d4e1d8fc57 100644
--- a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py
+++ b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Union
-
import torch
from ...image_processor import PipelineImageInput
@@ -31,7 +29,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -65,12 +63,12 @@ def __init__(
def prepare_latents(
self,
- video: Optional[torch.Tensor] = None,
+ video: torch.Tensor | None = None,
batch_size: int = 1,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -243,16 +241,16 @@ def check_inputs(self, video, height, width, latents, tone_map_compression_ratio
@torch.no_grad()
def __call__(
self,
- video: Optional[List[PipelineImageInput]] = None,
+ video: list[PipelineImageInput] | None = None,
height: int = 512,
width: int = 704,
- latents: Optional[torch.Tensor] = None,
- decode_timestep: Union[float, List[float]] = 0.0,
- decode_noise_scale: Optional[Union[float, List[float]]] = None,
+ latents: torch.Tensor | None = None,
+ decode_timestep: float | list[float] = 0.0,
+ decode_noise_scale: float | list[float] | None = None,
adain_factor: float = 0.0,
tone_map_compression_ratio: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
):
self.check_inputs(
diff --git a/src/diffusers/pipelines/ltx/pipeline_output.py b/src/diffusers/pipelines/ltx/pipeline_output.py
index 36ec3ea884a2..f5cb34aa508d 100644
--- a/src/diffusers/pipelines/ltx/pipeline_output.py
+++ b/src/diffusers/pipelines/ltx/pipeline_output.py
@@ -11,8 +11,8 @@ class LTXPipelineOutput(BaseOutput):
Output class for LTX pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index 22ca42d37902..4b2a81a9dc2c 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -1,5 +1,3 @@
-from typing import Optional, Tuple, Union
-
import torch
import torch.nn as nn
import torch.nn.functional as F
@@ -40,8 +38,8 @@ def forward(
self,
batch_size: int,
pos: int,
- device: Union[str, torch.device],
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ device: str | torch.device,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# 1. Get 1D position ids
grid_1d = torch.arange(pos, dtype=torch.float32, device=device)
# Get fractional indices relative to self.base_seq_len
@@ -127,8 +125,8 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- rotary_emb: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ rotary_emb: torch.Tensor | None = None,
) -> torch.Tensor:
norm_hidden_states = self.norm1(hidden_states)
attn_hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, query_rotary_emb=rotary_emb)
@@ -202,9 +200,9 @@ def __init__(
def forward(
self,
hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
attn_mask_binarize_threshold: float = -9000.0,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
# hidden_states shape: [batch_size, seq_len, hidden_dim]
# attention_mask shape: [batch_size, seq_len] or [batch_size, 1, 1, seq_len]
batch_size, seq_len, _ = hidden_states.shape
diff --git a/src/diffusers/pipelines/ltx2/export_utils.py b/src/diffusers/pipelines/ltx2/export_utils.py
index 347601422c83..f0287506b8db 100644
--- a/src/diffusers/pipelines/ltx2/export_utils.py
+++ b/src/diffusers/pipelines/ltx2/export_utils.py
@@ -16,7 +16,6 @@
from collections.abc import Iterator
from fractions import Fraction
from itertools import chain
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -109,10 +108,10 @@ def _write_audio(
def encode_video(
- video: Union[List[PIL.Image.Image], np.ndarray, torch.Tensor, Iterator[torch.Tensor]],
+ video: list[PIL.Image.Image] | np.ndarray | torch.Tensor | Iterator[torch.Tensor],
fps: int,
- audio: Optional[torch.Tensor],
- audio_sample_rate: Optional[int],
+ audio: torch.Tensor,
+ audio_sample_rate: int,
output_path: str,
video_chunks_number: int = 1,
) -> None:
diff --git a/src/diffusers/pipelines/ltx2/latent_upsampler.py b/src/diffusers/pipelines/ltx2/latent_upsampler.py
index 69a9b1d9193f..f6c589a70ab6 100644
--- a/src/diffusers/pipelines/ltx2/latent_upsampler.py
+++ b/src/diffusers/pipelines/ltx2/latent_upsampler.py
@@ -13,7 +13,6 @@
# limitations under the License.
import math
-from typing import Optional
import torch
import torch.nn.functional as F
@@ -32,7 +31,7 @@
# Copied from diffusers.pipelines.ltx.modeling_latent_upsampler.ResBlock
class ResBlock(torch.nn.Module):
- def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3):
+ def __init__(self, channels: int, mid_channels: int | None = None, dims: int = 3):
super().__init__()
if mid_channels is None:
mid_channels = channels
@@ -196,7 +195,7 @@ def __init__(
dims: int = 3,
spatial_upsample: bool = True,
temporal_upsample: bool = False,
- rational_spatial_scale: Optional[float] = 2.0,
+ rational_spatial_scale: float | None = 2.0,
):
super().__init__()
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index cb01159a81a7..037840360137 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -14,7 +14,7 @@
import copy
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -98,10 +98,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -116,15 +116,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -218,7 +218,7 @@ def __init__(
vae: AutoencoderKLLTX2Video,
audio_vae: AutoencoderKLLTX2Audio,
text_encoder: Gemma3ForConditionalGeneration,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
connectors: LTX2TextConnectors,
transformer: LTX2VideoTransformer3DModel,
vocoder: LTX2Vocoder,
@@ -272,7 +272,7 @@ def __init__(
def _pack_text_embeds(
text_hidden_states: torch.Tensor,
sequence_lengths: torch.Tensor,
- device: Union[str, torch.device],
+ device: str | torch.device,
padding_side: str = "left",
scale_factor: int = 8,
eps: float = 1e-6,
@@ -337,18 +337,18 @@ def _pack_text_embeds(
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_videos_per_prompt: int = 1,
max_sequence_length: int = 1024,
scale_factor: int = 8,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`str` or `torch.device`):
torch device to place the resulting embeddings on
@@ -410,26 +410,26 @@ def _get_gemma_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
scale_factor: int = 8,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -617,7 +617,7 @@ def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor
@staticmethod
def _create_noised_state(
- latents: torch.Tensor, noise_scale: Union[float, torch.Tensor], generator: Optional[torch.Generator] = None
+ latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None
):
noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype)
noised_latents = noise_scale * noise + (1 - noise_scale) * latents
@@ -625,7 +625,7 @@ def _create_noised_state(
@staticmethod
def _pack_audio_latents(
- latents: torch.Tensor, patch_size: Optional[int] = None, patch_size_t: Optional[int] = None
+ latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None
) -> torch.Tensor:
# Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins
if patch_size is not None and patch_size_t is not None:
@@ -649,8 +649,8 @@ def _unpack_audio_latents(
latents: torch.Tensor,
latent_length: int,
num_mel_bins: int,
- patch_size: Optional[int] = None,
- patch_size_t: Optional[int] = None,
+ patch_size: int | None = None,
+ patch_size_t: int | None = None,
) -> torch.Tensor:
# Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M],
# where L is the latent audio length and M is the number of mel bins.
@@ -671,10 +671,10 @@ def prepare_latents(
width: int = 768,
num_frames: int = 121,
noise_scale: float = 0.0,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
if latents.ndim == 5:
@@ -717,10 +717,10 @@ def prepare_audio_latents(
audio_latent_length: int = 1, # 1 is just a dummy value
num_mel_bins: int = 64,
noise_scale: float = 0.0,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
if latents.ndim == 4:
@@ -781,40 +781,40 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 768,
num_frames: int = 121,
frame_rate: float = 24.0,
num_inference_steps: int = 40,
- sigmas: Optional[List[float]] = None,
- timesteps: List[int] = None,
+ sigmas: list[float] | None = None,
+ timesteps: list[int] = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- audio_latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- decode_timestep: Union[float, List[float]] = 0.0,
- decode_noise_scale: Optional[Union[float, List[float]]] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ audio_latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ decode_timestep: float | list[float] = 0.0,
+ decode_noise_scale: float | list[float] | None = None,
+ output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 1024,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to `512`):
@@ -832,7 +832,7 @@ def __call__(
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -853,7 +853,7 @@ def __call__(
the `latents` and `audio_latents` before continue denoising.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1081,6 +1081,10 @@ def __call__(
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
audio_latents.shape[0], audio_num_frames, audio_latents.device
)
+ # Duplicate the positional ids as well if using CFG
+ if self.do_classifier_free_guidance:
+ video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim
+ audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
# 7. Denoising loop
with self.progress_bar(total=num_inference_steps) as progress_bar:
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index c120e1f010e9..98d323efd477 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -14,7 +14,7 @@
import copy
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -22,7 +22,7 @@
from ...callbacks import MultiPipelineCallbacks, PipelineCallback
from ...image_processor import PipelineImageInput
-from ...loaders import FromSingleFileMixin, LTXVideoLoraLoaderMixin
+from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin
from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video
from ...models.transformers import LTX2VideoTransformer3DModel
from ...schedulers import FlowMatchEulerDiscreteScheduler
@@ -48,7 +48,7 @@
Examples:
```py
>>> import torch
- >>> from diffusers import LTX2Pipeline
+ >>> from diffusers import LTX2ImageToVideoPipeline
>>> from diffusers.pipelines.ltx2.export_utils import encode_video
>>> from diffusers.utils import load_image
@@ -62,7 +62,7 @@
>>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
>>> frame_rate = 24.0
- >>> video = pipe(
+ >>> video, audio = pipe(
... image=image,
... prompt=prompt,
... negative_prompt=negative_prompt,
@@ -89,7 +89,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -118,10 +118,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -136,15 +136,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -202,7 +202,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
return noise_cfg
-class LTX2ImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTXVideoLoraLoaderMixin):
+class LTX2ImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin):
r"""
Pipeline for image-to-video generation.
@@ -221,7 +221,7 @@ def __init__(
vae: AutoencoderKLLTX2Video,
audio_vae: AutoencoderKLLTX2Audio,
text_encoder: Gemma3ForConditionalGeneration,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
connectors: LTX2TextConnectors,
transformer: LTX2VideoTransformer3DModel,
vocoder: LTX2Vocoder,
@@ -276,7 +276,7 @@ def __init__(
def _pack_text_embeds(
text_hidden_states: torch.Tensor,
sequence_lengths: torch.Tensor,
- device: Union[str, torch.device],
+ device: str | torch.device,
padding_side: str = "left",
scale_factor: int = 8,
eps: float = 1e-6,
@@ -342,18 +342,18 @@ def _pack_text_embeds(
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._get_gemma_prompt_embeds
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_videos_per_prompt: int = 1,
max_sequence_length: int = 1024,
scale_factor: int = 8,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`str` or `torch.device`):
torch device to place the resulting embeddings on
@@ -416,26 +416,26 @@ def _get_gemma_prompt_embeds(
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
scale_factor: int = 8,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -615,7 +615,7 @@ def _denormalize_latents(
@staticmethod
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._create_noised_state
def _create_noised_state(
- latents: torch.Tensor, noise_scale: Union[float, torch.Tensor], generator: Optional[torch.Generator] = None
+ latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None
):
noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype)
noised_latents = noise_scale * noise + (1 - noise_scale) * latents
@@ -624,7 +624,7 @@ def _create_noised_state(
@staticmethod
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_audio_latents
def _pack_audio_latents(
- latents: torch.Tensor, patch_size: Optional[int] = None, patch_size_t: Optional[int] = None
+ latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None
) -> torch.Tensor:
# Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins
if patch_size is not None and patch_size_t is not None:
@@ -649,8 +649,8 @@ def _unpack_audio_latents(
latents: torch.Tensor,
latent_length: int,
num_mel_bins: int,
- patch_size: Optional[int] = None,
- patch_size_t: Optional[int] = None,
+ patch_size: int | None = None,
+ patch_size_t: int | None = None,
) -> torch.Tensor:
# Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M],
# where L is the latent audio length and M is the number of mel bins.
@@ -679,17 +679,17 @@ def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor
def prepare_latents(
self,
- image: Optional[torch.Tensor] = None,
+ image: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 704,
num_frames: int = 161,
noise_scale: float = 0.0,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
height = height // self.vae_spatial_compression_ratio
width = width // self.vae_spatial_compression_ratio
@@ -764,10 +764,10 @@ def prepare_audio_latents(
audio_latent_length: int = 1, # 1 is just a dummy value
num_mel_bins: int = 64,
noise_scale: float = 0.0,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
if latents.ndim == 4:
@@ -829,33 +829,33 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
height: int = 512,
width: int = 768,
num_frames: int = 121,
frame_rate: float = 24.0,
num_inference_steps: int = 40,
- sigmas: Optional[List[float]] = None,
- timesteps: List[int] = None,
+ sigmas: list[float] | None = None,
+ timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- audio_latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- decode_timestep: Union[float, List[float]] = 0.0,
- decode_noise_scale: Optional[Union[float, List[float]]] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ audio_latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ decode_timestep: float | list[float] = 0.0,
+ decode_noise_scale: float | list[float] | None = None,
+ output_type: str = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 1024,
):
r"""
@@ -864,7 +864,7 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to `512`):
@@ -903,7 +903,7 @@ def __call__(
the `latents` and `audio_latents` before continue denoising.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1139,6 +1139,10 @@ def __call__(
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
audio_latents.shape[0], audio_num_frames, audio_latents.device
)
+ # Duplicate the positional ids as well if using CFG
+ if self.do_classifier_free_guidance:
+ video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim
+ audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
# 7. Denoising loop
with self.progress_bar(total=num_inference_steps) as progress_bar:
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py
index b0db1bdee317..00d81dfd11c3 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import List, Optional, Union
import torch
@@ -90,7 +89,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -124,17 +123,17 @@ def __init__(
def prepare_latents(
self,
- video: Optional[torch.Tensor] = None,
+ video: torch.Tensor | None = None,
batch_size: int = 1,
num_frames: int = 121,
height: int = 512,
width: int = 768,
spatial_patch_size: int = 1,
temporal_patch_size: int = 1,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
if latents.ndim == 3:
@@ -266,27 +265,27 @@ def check_inputs(self, video, height, width, latents, tone_map_compression_ratio
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- video: Optional[List[PipelineImageInput]] = None,
+ video: list[PipelineImageInput] | None = None,
height: int = 512,
width: int = 768,
num_frames: int = 121,
spatial_patch_size: int = 1,
temporal_patch_size: int = 1,
- latents: Optional[torch.Tensor] = None,
+ latents: torch.Tensor | None = None,
latents_normalized: bool = False,
- decode_timestep: Union[float, List[float]] = 0.0,
- decode_noise_scale: Optional[Union[float, List[float]]] = None,
+ decode_timestep: float | list[float] = 0.0,
+ decode_noise_scale: float | list[float] | None = None,
adain_factor: float = 0.0,
tone_map_compression_ratio: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- video (`List[PipelineImageInput]`, *optional*)
+ video (`list[PipelineImageInput]`, *optional*)
The video to be upsampled (such as a LTX 2.0 first stage output). If not supplied, `latents` should be
supplied.
height (`int`, *optional*, defaults to `512`):
@@ -319,7 +318,7 @@ def __call__(
This is useful for regularizing high-variance latents or for conditioning outputs during generation.
Should be in [0, 1], where 0.0 (the default) means tone mapping is not applied and 1.0 corresponds to
the full compression effect.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
output_type (`str`, *optional*, defaults to `"pil"`):
diff --git a/src/diffusers/pipelines/ltx2/pipeline_output.py b/src/diffusers/pipelines/ltx2/pipeline_output.py
index eacd571125b0..02891219ee67 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_output.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_output.py
@@ -11,7 +11,7 @@ class LTX2PipelineOutput(BaseOutput):
Output class for LTX pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py
index 217c68103e39..551c3ac5980f 100644
--- a/src/diffusers/pipelines/ltx2/vocoder.py
+++ b/src/diffusers/pipelines/ltx2/vocoder.py
@@ -1,5 +1,4 @@
import math
-from typing import List, Tuple
import torch
import torch.nn as nn
@@ -15,7 +14,7 @@ def __init__(
channels: int,
kernel_size: int = 3,
stride: int = 1,
- dilations: Tuple[int, ...] = (1, 3, 5),
+ dilations: tuple[int, ...] = (1, 3, 5),
leaky_relu_negative_slope: float = 0.1,
padding_mode: str = "same",
):
@@ -58,10 +57,10 @@ def __init__(
in_channels: int = 128,
hidden_channels: int = 1024,
out_channels: int = 2,
- upsample_kernel_sizes: List[int] = [16, 15, 8, 4, 4],
- upsample_factors: List[int] = [6, 5, 2, 2, 2],
- resnet_kernel_sizes: List[int] = [3, 7, 11],
- resnet_dilations: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ upsample_kernel_sizes: list[int] = [16, 15, 8, 4, 4],
+ upsample_factors: list[int] = [6, 5, 2, 2, 2],
+ resnet_kernel_sizes: list[int] = [3, 7, 11],
+ resnet_dilations: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
leaky_relu_negative_slope: float = 0.1,
output_sampling_rate: int = 24000,
):
diff --git a/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py b/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py
index 8065a17b7889..392af492b702 100644
--- a/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py
+++ b/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py
@@ -17,7 +17,7 @@
# - Based on pipeline_wan.py, but with supports receiving a condition video appended to the channel dimension.
import html
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import regex as re
import torch
@@ -51,7 +51,7 @@
EXAMPLE_DOC_STRING = """
Examples:
```python
- >>> from typing import List
+ >>> from typing import list
>>> import torch
>>> from PIL import Image
@@ -69,7 +69,7 @@
>>> # Load video
- >>> def convert_video(video: List[Image.Image]) -> List[Image.Image]:
+ >>> def convert_video(video: list[Image.Image]) -> list[Image.Image]:
... video = load_video(url)[:num_frames]
... video = [video[i].resize((width, height)) for i in range(num_frames)]
... return video
@@ -119,7 +119,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -172,9 +172,9 @@ def __init__(
text_encoder: UMT5EncoderModel,
vae: AutoencoderKLWan,
scheduler: FlowMatchEulerDiscreteScheduler,
- transformer: Optional[WanTransformer3DModel] = None,
- transformer_2: Optional[WanTransformer3DModel] = None,
- boundary_ratio: Optional[float] = None,
+ transformer: WanTransformer3DModel | None = None,
+ transformer_2: WanTransformer3DModel | None = None,
+ boundary_ratio: float | None = None,
expand_timesteps: bool = False, # Wan2.2 ti2v
):
super().__init__()
@@ -196,11 +196,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -238,23 +238,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -368,16 +368,16 @@ def check_inputs(
def prepare_latents(
self,
- video: Optional[torch.Tensor] = None,
+ video: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 16,
height: int = 480,
width: int = 832,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
@@ -451,38 +451,36 @@ def attention_kwargs(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- video: List[Image.Image],
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ video: list[Image.Image],
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 480,
width: int = 832,
num_frames: int = 81,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- guidance_scale_2: Optional[float] = None,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ guidance_scale_2: float | None = None,
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- video (`List[Image.Image]`):
+ video (`list[Image.Image]`):
The video to use as the condition for the video generation.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds`
instead. Ignored when not using guidance (`guidance_scale` < `1`).
height (`int`, defaults to `480`):
@@ -506,7 +504,7 @@ def __call__(
and the pipeline's `boundary_ratio` are not None.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -529,7 +527,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/lucy/pipeline_output.py b/src/diffusers/pipelines/lucy/pipeline_output.py
index cf9ea91fd106..197ce194f475 100644
--- a/src/diffusers/pipelines/lucy/pipeline_output.py
+++ b/src/diffusers/pipelines/lucy/pipeline_output.py
@@ -11,8 +11,8 @@ class LucyPipelineOutput(BaseOutput):
Output class for Lucy pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
diff --git a/src/diffusers/pipelines/lumina/pipeline_lumina.py b/src/diffusers/pipelines/lumina/pipeline_lumina.py
index f4711cf9d9d8..cc123218f4ee 100644
--- a/src/diffusers/pipelines/lumina/pipeline_lumina.py
+++ b/src/diffusers/pipelines/lumina/pipeline_lumina.py
@@ -17,7 +17,7 @@
import math
import re
import urllib.parse as ul
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import GemmaPreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
@@ -76,10 +76,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -94,15 +94,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -182,7 +182,7 @@ def __init__(
scheduler: FlowMatchEulerDiscreteScheduler,
vae: AutoencoderKL,
text_encoder: GemmaPreTrainedModel,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
):
super().__init__()
@@ -205,11 +205,11 @@ def __init__(
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- clean_caption: Optional[bool] = False,
- max_length: Optional[int] = None,
+ device: torch.device | None = None,
+ clean_caption: bool | None = False,
+ max_length: int | None = None,
):
device = device or self._execution_device
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -261,15 +261,15 @@ def _get_gemma_prompt_embeds(
# Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
- negative_prompt: Union[str, List[str]] = None,
+ negative_prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
**kwargs,
):
@@ -277,9 +277,9 @@ def encode_prompt(
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
Lumina-T2I, this should be "".
@@ -633,46 +633,44 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- width: Optional[int] = None,
- height: Optional[int] = None,
+ prompt: str | list[str] = None,
+ width: int | None = None,
+ height: int | None = None,
num_inference_steps: int = 30,
guidance_scale: float = 4.0,
- negative_prompt: Union[str, List[str]] = None,
- sigmas: List[float] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] = None,
+ sigmas: list[float] = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
clean_caption: bool = True,
max_sequence_length: int = 256,
- scaling_watershed: Optional[float] = 1.0,
- proportional_attn: Optional[bool] = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- ) -> Union[ImagePipelineOutput, Tuple]:
+ scaling_watershed: float | None = 1.0,
+ proportional_attn: bool | None = True,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ ) -> ImagePipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 30):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -691,7 +689,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -723,7 +721,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
@@ -946,7 +944,7 @@ def __init__(
scheduler: FlowMatchEulerDiscreteScheduler,
vae: AutoencoderKL,
text_encoder: GemmaPreTrainedModel,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
):
deprecation_message = "`LuminaText2ImgPipeline` has been renamed to `LuminaPipeline` and will be removed in a future version. Please use `LuminaPipeline` instead."
deprecate("diffusers.pipelines.lumina.pipeline_lumina.LuminaText2ImgPipeline", "0.34", deprecation_message)
diff --git a/src/diffusers/pipelines/lumina2/pipeline_lumina2.py b/src/diffusers/pipelines/lumina2/pipeline_lumina2.py
index 8151b29b25fd..576d3e8d9486 100644
--- a/src/diffusers/pipelines/lumina2/pipeline_lumina2.py
+++ b/src/diffusers/pipelines/lumina2/pipeline_lumina2.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -77,10 +77,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -95,15 +95,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -164,7 +164,7 @@ def __init__(
scheduler: FlowMatchEulerDiscreteScheduler,
vae: AutoencoderKL,
text_encoder: Gemma2PreTrainedModel,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
):
super().__init__()
@@ -191,10 +191,10 @@ def __init__(
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
max_sequence_length: int = 256,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
device = device or self._execution_device
prompt = [prompt] if isinstance(prompt, str) else prompt
text_inputs = self.tokenizer(
@@ -237,25 +237,25 @@ def _get_gemma_prompt_embeds(
# Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
- negative_prompt: Union[str, List[str]] = None,
+ negative_prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- system_prompt: Optional[str] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ system_prompt: str | None = None,
max_sequence_length: int = 256,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
Lumina-T2I, this should be "".
@@ -525,45 +525,45 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- width: Optional[int] = None,
- height: Optional[int] = None,
+ prompt: str | list[str] = None,
+ width: int | None = None,
+ height: int | None = None,
num_inference_steps: int = 30,
guidance_scale: float = 4.0,
- negative_prompt: Union[str, List[str]] = None,
- sigmas: List[float] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] = None,
+ sigmas: list[float] = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- system_prompt: Optional[str] = None,
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ system_prompt: str | None = None,
cfg_trunc_ratio: float = 1.0,
cfg_normalization: bool = True,
max_sequence_length: int = 256,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 30):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -582,7 +582,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -612,7 +612,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
@@ -805,7 +805,7 @@ def __init__(
scheduler: FlowMatchEulerDiscreteScheduler,
vae: AutoencoderKL,
text_encoder: Gemma2PreTrainedModel,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
):
deprecation_message = "`Lumina2Text2ImgPipeline` has been renamed to `Lumina2Pipeline` and will be removed in a future version. Please use `Lumina2Pipeline` instead."
deprecate("diffusers.pipelines.lumina2.pipeline_lumina2.Lumina2Text2ImgPipeline", "0.34", deprecation_message)
diff --git a/src/diffusers/pipelines/marigold/marigold_image_processing.py b/src/diffusers/pipelines/marigold/marigold_image_processing.py
index 5130a876606a..4a76ef9892e4 100644
--- a/src/diffusers/pipelines/marigold/marigold_image_processing.py
+++ b/src/diffusers/pipelines/marigold/marigold_image_processing.py
@@ -16,7 +16,7 @@
# More information and citation instructions are available on the
# Marigold project website: https://marigoldcomputervision.github.io
# --------------------------------------------------------------------------
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import PIL
@@ -47,7 +47,7 @@ def __init__(
super().__init__()
@staticmethod
- def expand_tensor_or_array(images: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]:
+ def expand_tensor_or_array(images: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
"""
Expand a tensor or array to a specified number of images.
"""
@@ -90,7 +90,7 @@ def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
@staticmethod
def resize_antialias(
- image: torch.Tensor, size: Tuple[int, int], mode: str, is_aa: Optional[bool] = None
+ image: torch.Tensor, size: tuple[int, int], mode: str, is_aa: bool | None = None
) -> torch.Tensor:
if not torch.is_tensor(image):
raise ValueError(f"Invalid input type={type(image)}.")
@@ -126,7 +126,7 @@ def resize_to_max_edge(image: torch.Tensor, max_edge_sz: int, mode: str) -> torc
return image
@staticmethod
- def pad_image(image: torch.Tensor, align: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
+ def pad_image(image: torch.Tensor, align: int) -> tuple[torch.Tensor, tuple[int, int]]:
if not torch.is_tensor(image):
raise ValueError(f"Invalid input type={type(image)}.")
if not torch.is_floating_point(image):
@@ -142,7 +142,7 @@ def pad_image(image: torch.Tensor, align: int) -> Tuple[torch.Tensor, Tuple[int,
return image, (ph, pw)
@staticmethod
- def unpad_image(image: torch.Tensor, padding: Tuple[int, int]) -> torch.Tensor:
+ def unpad_image(image: torch.Tensor, padding: tuple[int, int]) -> torch.Tensor:
if not torch.is_tensor(image):
raise ValueError(f"Invalid input type={type(image)}.")
if not torch.is_floating_point(image):
@@ -160,10 +160,10 @@ def unpad_image(image: torch.Tensor, padding: Tuple[int, int]) -> torch.Tensor:
@staticmethod
def load_image_canonical(
- image: Union[torch.Tensor, np.ndarray, Image.Image],
+ image: torch.Tensor | np.ndarray | Image.Image,
device: torch.device = torch.device("cpu"),
dtype: torch.dtype = torch.float32,
- ) -> Tuple[torch.Tensor, int]:
+ ) -> tuple[torch.Tensor, int]:
if isinstance(image, Image.Image):
image = np.array(image)
@@ -216,7 +216,7 @@ def check_image_values_range(image: torch.Tensor) -> None:
def preprocess(
self,
image: PipelineImageInput,
- processing_resolution: Optional[int] = None,
+ processing_resolution: int | None = None,
resample_method_input: str = "bilinear",
device: torch.device = torch.device("cpu"),
dtype: torch.dtype = torch.float32,
@@ -256,11 +256,11 @@ def preprocess(
@staticmethod
def colormap(
- image: Union[np.ndarray, torch.Tensor],
+ image: np.ndarray | torch.Tensor,
cmap: str = "Spectral",
bytes: bool = False,
- _force_method: Optional[str] = None,
- ) -> Union[np.ndarray, torch.Tensor]:
+ _force_method: str | None = None,
+ ) -> np.ndarray | torch.Tensor:
"""
Converts a monochrome image into an RGB image by applying the specified colormap. This function mimics the
behavior of matplotlib.colormaps, but allows the user to use the most discriminative color maps ("Spectral",
@@ -386,30 +386,28 @@ def method_custom(image, cmap, bytes=False):
@staticmethod
def visualize_depth(
- depth: Union[
- PIL.Image.Image,
- np.ndarray,
- torch.Tensor,
- List[PIL.Image.Image],
- List[np.ndarray],
- List[torch.Tensor],
- ],
+ depth: PIL.Image.Image
+ | np.ndarray
+ | torch.Tensor
+ | list[PIL.Image.Image]
+ | list[np.ndarray]
+ | list[torch.Tensor],
val_min: float = 0.0,
val_max: float = 1.0,
color_map: str = "Spectral",
- ) -> List[PIL.Image.Image]:
+ ) -> list[PIL.Image.Image]:
"""
Visualizes depth maps, such as predictions of the `MarigoldDepthPipeline`.
Args:
- depth (`Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image], List[np.ndarray],
- List[torch.Tensor]]`): Depth maps.
+ depth (`PIL.Image.Image | np.ndarray | torch.Tensor | list[PIL.Image.Image, list[np.ndarray],
+ list[torch.Tensor]]`): Depth maps.
val_min (`float`, *optional*, defaults to `0.0`): Minimum value of the visualized depth range.
val_max (`float`, *optional*, defaults to `1.0`): Maximum value of the visualized depth range.
color_map (`str`, *optional*, defaults to `"Spectral"`): Color map used to convert a single-channel
depth prediction into colored representation.
- Returns: `List[PIL.Image.Image]` with depth maps visualization.
+ Returns: `list[PIL.Image.Image]` with depth maps visualization.
"""
if val_max <= val_min:
raise ValueError(f"Invalid values range: [{val_min}, {val_max}].")
@@ -451,10 +449,10 @@ def visualize_depth_one(img, idx=None):
@staticmethod
def export_depth_to_16bit_png(
- depth: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]],
+ depth: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor],
val_min: float = 0.0,
val_max: float = 1.0,
- ) -> List[PIL.Image.Image]:
+ ) -> list[PIL.Image.Image]:
def export_depth_to_16bit_png_one(img, idx=None):
prefix = "Depth" + (f"[{idx}]" if idx else "")
if not isinstance(img, np.ndarray) and not torch.is_tensor(img):
@@ -487,21 +485,16 @@ def export_depth_to_16bit_png_one(img, idx=None):
@staticmethod
def visualize_normals(
- normals: Union[
- np.ndarray,
- torch.Tensor,
- List[np.ndarray],
- List[torch.Tensor],
- ],
+ normals: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor],
flip_x: bool = False,
flip_y: bool = False,
flip_z: bool = False,
- ) -> List[PIL.Image.Image]:
+ ) -> list[PIL.Image.Image]:
"""
Visualizes surface normals, such as predictions of the `MarigoldNormalsPipeline`.
Args:
- normals (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
+ normals (`np.ndarray | torch.Tensor | list[np.ndarray, list[torch.Tensor]]`):
Surface normals.
flip_x (`bool`, *optional*, defaults to `False`): Flips the X axis of the normals frame of reference.
Default direction is right.
@@ -510,7 +503,7 @@ def visualize_normals(
flip_z (`bool`, *optional*, defaults to `False`): Flips the Z axis of the normals frame of reference.
Default direction is facing the observer.
- Returns: `List[PIL.Image.Image]` with surface normals visualization.
+ Returns: `list[PIL.Image.Image]` with surface normals visualization.
"""
flip_vec = None
if any((flip_x, flip_y, flip_z)):
@@ -548,30 +541,25 @@ def visualize_normals_one(img, idx=None):
@staticmethod
def visualize_intrinsics(
- prediction: Union[
- np.ndarray,
- torch.Tensor,
- List[np.ndarray],
- List[torch.Tensor],
- ],
- target_properties: Dict[str, Any],
- color_map: Union[str, Dict[str, str]] = "binary",
- ) -> List[Dict[str, PIL.Image.Image]]:
+ prediction: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor],
+ target_properties: dict[str, Any],
+ color_map: str | dict[str, str] = "binary",
+ ) -> list[dict[str, PIL.Image.Image]]:
"""
Visualizes intrinsic image decomposition, such as predictions of the `MarigoldIntrinsicsPipeline`.
Args:
- prediction (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
+ prediction (`np.ndarray | torch.Tensor | list[np.ndarray, list[torch.Tensor]]`):
Intrinsic image decomposition.
- target_properties (`Dict[str, Any]`):
- Decomposition properties. Expected entries: `target_names: List[str]` and a dictionary with keys
- `prediction_space: str`, `sub_target_names: List[Union[str, Null]]` (must have 3 entries, null for
- missing modalities), `up_to_scale: bool`, one for each target and sub-target.
- color_map (`Union[str, Dict[str, str]]`, *optional*, defaults to `"Spectral"`):
+ target_properties (`dict[str, Any]`):
+ Decomposition properties. Expected entries: `target_names: list[str]` and a dictionary with keys
+ `prediction_space: str`, `sub_target_names: list[str | Null]` (must have 3 entries, null for missing
+ modalities), `up_to_scale: bool`, one for each target and sub-target.
+ color_map (`str | dict[str, str]`, *optional*, defaults to `"Spectral"`):
Color map used to convert a single-channel predictions into colored representations. When a dictionary
is passed, each modality can be colored with its own color map.
- Returns: `List[Dict[str, PIL.Image.Image]]` with intrinsic image decomposition visualization.
+ Returns: `list[dict[str, PIL.Image.Image]]` with intrinsic image decomposition visualization.
"""
if "target_names" not in target_properties:
raise ValueError("Missing `target_names` in target_properties")
@@ -641,25 +629,20 @@ def visualize_targets_one(images, idx=None):
@staticmethod
def visualize_uncertainty(
- uncertainty: Union[
- np.ndarray,
- torch.Tensor,
- List[np.ndarray],
- List[torch.Tensor],
- ],
+ uncertainty: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor],
saturation_percentile=95,
- ) -> List[PIL.Image.Image]:
+ ) -> list[PIL.Image.Image]:
"""
Visualizes dense uncertainties, such as produced by `MarigoldDepthPipeline`, `MarigoldNormalsPipeline`, or
`MarigoldIntrinsicsPipeline`.
Args:
- uncertainty (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
+ uncertainty (`np.ndarray | torch.Tensor | list[np.ndarray, list[torch.Tensor]]`):
Uncertainty maps.
saturation_percentile (`int`, *optional*, defaults to `95`):
Specifies the percentile uncertainty value visualized with maximum intensity.
- Returns: `List[PIL.Image.Image]` with uncertainty visualization.
+ Returns: `list[PIL.Image.Image]` with uncertainty visualization.
"""
def visualize_uncertainty_one(img, idx=None):
diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py b/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py
index 92ec16fd455b..a81d1c51742c 100644
--- a/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py
+++ b/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py
@@ -18,7 +18,7 @@
# --------------------------------------------------------------------------
from dataclasses import dataclass
from functools import partial
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -96,9 +96,9 @@ class MarigoldDepthOutput(BaseOutput):
The shape is `numimages * numensemble × 4 × latentheight × latentwidth`.
"""
- prediction: Union[np.ndarray, torch.Tensor]
- uncertainty: Union[None, np.ndarray, torch.Tensor]
- latent: Union[None, torch.Tensor]
+ prediction: np.ndarray | torch.Tensor
+ uncertainty: None | np.ndarray | torch.Tensor
+ latent: None | torch.Tensor
class MarigoldDepthPipeline(DiffusionPipeline):
@@ -150,14 +150,14 @@ def __init__(
self,
unet: UNet2DConditionModel,
vae: AutoencoderKL,
- scheduler: Union[DDIMScheduler, LCMScheduler],
+ scheduler: DDIMScheduler | LCMScheduler,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- prediction_type: Optional[str] = None,
- scale_invariant: Optional[bool] = True,
- shift_invariant: Optional[bool] = True,
- default_denoising_steps: Optional[int] = None,
- default_processing_resolution: Optional[int] = None,
+ prediction_type: str | None = None,
+ scale_invariant: bool | None = True,
+ shift_invariant: bool | None = True,
+ default_denoising_steps: int | None = None,
+ default_processing_resolution: int | None = None,
):
super().__init__()
@@ -202,9 +202,9 @@ def check_inputs(
resample_method_input: str,
resample_method_output: str,
batch_size: int,
- ensembling_kwargs: Optional[Dict[str, Any]],
- latents: Optional[torch.Tensor],
- generator: Optional[Union[torch.Generator, List[torch.Generator]]],
+ ensembling_kwargs: dict[str, Any] | None,
+ latents: torch.Tensor | None,
+ generator: torch.Generator | list[torch.Generator] | None,
output_type: str,
output_uncertainty: bool,
) -> int:
@@ -349,16 +349,16 @@ def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
def __call__(
self,
image: PipelineImageInput,
- num_inference_steps: Optional[int] = None,
+ num_inference_steps: int | None = None,
ensemble_size: int = 1,
- processing_resolution: Optional[int] = None,
+ processing_resolution: int | None = None,
match_input_resolution: bool = True,
resample_method_input: str = "bilinear",
resample_method_output: str = "bilinear",
batch_size: int = 1,
- ensembling_kwargs: Optional[Dict[str, Any]] = None,
- latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ ensembling_kwargs: dict[str, Any] | None = None,
+ latents: torch.Tensor | list[torch.Tensor] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
output_type: str = "np",
output_uncertainty: bool = False,
output_latent: bool = False,
@@ -368,8 +368,8 @@ def __call__(
Function invoked when calling the pipeline.
Args:
- image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
- `List[torch.Tensor]`: An input image or images used as an input for the depth estimation task. For
+ image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`),
+ `list[torch.Tensor]`: An input image or images used as an input for the depth estimation task. For
arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
@@ -406,10 +406,10 @@ def __call__(
tolerance is reached.
- max_res (`int`, *optional*, defaults to `None`): Resolution at which the alignment is performed;
`None` matches the `processing_resolution`.
- latents (`torch.Tensor`, or `List[torch.Tensor]`, *optional*, defaults to `None`):
+ latents (`torch.Tensor`, or `list[torch.Tensor]`, *optional*, defaults to `None`):
Latent noise tensors to replace the random initialization. These can be taken from the previous
function call's output.
- generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`):
+ generator (`torch.Generator`, or `list[torch.Generator]`, *optional*, defaults to `None`):
Random number generator object to ensure reproducibility.
output_type (`str`, *optional*, defaults to `"np"`):
Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
@@ -621,11 +621,11 @@ def __call__(
def prepare_latents(
self,
image: torch.Tensor,
- latents: Optional[torch.Tensor],
- generator: Optional[torch.Generator],
+ latents: torch.Tensor | None,
+ generator: torch.Generator | None,
ensemble_size: int,
batch_size: int,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
def retrieve_latents(encoder_output):
if hasattr(encoder_output, "latent_dist"):
return encoder_output.latent_dist.mode()
@@ -680,7 +680,7 @@ def ensemble_depth(
max_iter: int = 2,
tol: float = 1e-3,
max_res: int = 1024,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
"""
Ensembles the depth maps represented by the `depth` tensor with expected shape `(B, 1, H, W)`, where B is the
number of ensemble members for a given prediction of size `(H x W)`. Even though the function is designed for
@@ -754,7 +754,7 @@ def align(depth: torch.Tensor, param: np.ndarray) -> torch.Tensor:
def ensemble(
depth_aligned: torch.Tensor, return_uncertainty: bool = False
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
uncertainty = None
if reduction == "mean":
prediction = torch.mean(depth_aligned, dim=0, keepdim=True)
diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py b/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py
index bef9ca77c708..9488d8f5c9b8 100644
--- a/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py
+++ b/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py
@@ -17,7 +17,7 @@
# Marigold project website: https://marigoldcomputervision.github.io
# --------------------------------------------------------------------------
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -112,9 +112,9 @@ class MarigoldIntrinsicsOutput(BaseOutput):
The shape is `(numimages * numensemble) × (numtargets * 4) × latentheight × latentwidth`.
"""
- prediction: Union[np.ndarray, torch.Tensor]
- uncertainty: Union[None, np.ndarray, torch.Tensor]
- latent: Union[None, torch.Tensor]
+ prediction: np.ndarray | torch.Tensor
+ uncertainty: None | np.ndarray | torch.Tensor
+ latent: None | torch.Tensor
class MarigoldIntrinsicsPipeline(DiffusionPipeline):
@@ -139,8 +139,8 @@ class MarigoldIntrinsicsPipeline(DiffusionPipeline):
CLIP tokenizer.
prediction_type (`str`, *optional*):
Type of predictions made by the model.
- target_properties (`Dict[str, Any]`, *optional*):
- Properties of the predicted modalities, such as `target_names`, a `List[str]` used to define the number,
+ target_properties (`dict[str, Any]`, *optional*):
+ Properties of the predicted modalities, such as `target_names`, a `list[str]` used to define the number,
order and names of the predicted modalities, and any other metadata that may be required to interpret the
predictions.
default_denoising_steps (`int`, *optional*):
@@ -163,13 +163,13 @@ def __init__(
self,
unet: UNet2DConditionModel,
vae: AutoencoderKL,
- scheduler: Union[DDIMScheduler, LCMScheduler],
+ scheduler: DDIMScheduler | LCMScheduler,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- prediction_type: Optional[str] = None,
- target_properties: Optional[Dict[str, Any]] = None,
- default_denoising_steps: Optional[int] = None,
- default_processing_resolution: Optional[int] = None,
+ prediction_type: str | None = None,
+ target_properties: dict[str, Any] | None = None,
+ default_denoising_steps: int | None = None,
+ default_processing_resolution: int | None = None,
):
super().__init__()
@@ -216,9 +216,9 @@ def check_inputs(
resample_method_input: str,
resample_method_output: str,
batch_size: int,
- ensembling_kwargs: Optional[Dict[str, Any]],
- latents: Optional[torch.Tensor],
- generator: Optional[Union[torch.Generator, List[torch.Generator]]],
+ ensembling_kwargs: dict[str, Any] | None,
+ latents: torch.Tensor | None,
+ generator: torch.Generator | list[torch.Generator] | None,
output_type: str,
output_uncertainty: bool,
) -> int:
@@ -361,16 +361,16 @@ def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
def __call__(
self,
image: PipelineImageInput,
- num_inference_steps: Optional[int] = None,
+ num_inference_steps: int | None = None,
ensemble_size: int = 1,
- processing_resolution: Optional[int] = None,
+ processing_resolution: int | None = None,
match_input_resolution: bool = True,
resample_method_input: str = "bilinear",
resample_method_output: str = "bilinear",
batch_size: int = 1,
- ensembling_kwargs: Optional[Dict[str, Any]] = None,
- latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ ensembling_kwargs: dict[str, Any] | None = None,
+ latents: torch.Tensor | list[torch.Tensor] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
output_type: str = "np",
output_uncertainty: bool = False,
output_latent: bool = False,
@@ -380,8 +380,8 @@ def __call__(
Function invoked when calling the pipeline.
Args:
- image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
- `List[torch.Tensor]`: An input image or images used as an input for the intrinsic decomposition task.
+ image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`),
+ `list[torch.Tensor]`: An input image or images used as an input for the intrinsic decomposition task.
For arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is
possible by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
@@ -413,7 +413,7 @@ def __call__(
latents (`torch.Tensor`, *optional*, defaults to `None`):
Latent noise tensors to replace the random initialization. These can be taken from the previous
function call's output.
- generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`):
+ generator (`torch.Generator`, or `list[torch.Generator]`, *optional*, defaults to `None`):
Random number generator object to ensure reproducibility.
output_type (`str`, *optional*, defaults to `"np"`):
Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
@@ -627,11 +627,11 @@ def __call__(
def prepare_latents(
self,
image: torch.Tensor,
- latents: Optional[torch.Tensor],
- generator: Optional[torch.Generator],
+ latents: torch.Tensor | None,
+ generator: torch.Generator | None,
ensemble_size: int,
batch_size: int,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
def retrieve_latents(encoder_output):
if hasattr(encoder_output, "latent_dist"):
return encoder_output.latent_dist.mode()
@@ -680,7 +680,7 @@ def ensemble_intrinsics(
targets: torch.Tensor,
output_uncertainty: bool = False,
reduction: str = "median",
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
"""
Ensembles the intrinsic decomposition represented by the `targets` tensor with expected shape `(B, T, 3, H,
W)`, where B is the number of ensemble members for a given prediction of size `(H x W)`, and T is the number of
diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py b/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py
index 485a39c995ec..3f94ce441232 100644
--- a/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py
+++ b/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py
@@ -17,7 +17,7 @@
# Marigold project website: https://marigoldcomputervision.github.io
# --------------------------------------------------------------------------
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
import numpy as np
import torch
@@ -91,9 +91,9 @@ class MarigoldNormalsOutput(BaseOutput):
The shape is `numimages * numensemble × 4 × latentheight × latentwidth`.
"""
- prediction: Union[np.ndarray, torch.Tensor]
- uncertainty: Union[None, np.ndarray, torch.Tensor]
- latent: Union[None, torch.Tensor]
+ prediction: np.ndarray | torch.Tensor
+ uncertainty: None | np.ndarray | torch.Tensor
+ latent: None | torch.Tensor
class MarigoldNormalsPipeline(DiffusionPipeline):
@@ -140,13 +140,13 @@ def __init__(
self,
unet: UNet2DConditionModel,
vae: AutoencoderKL,
- scheduler: Union[DDIMScheduler, LCMScheduler],
+ scheduler: DDIMScheduler | LCMScheduler,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- prediction_type: Optional[str] = None,
- use_full_z_range: Optional[bool] = True,
- default_denoising_steps: Optional[int] = None,
- default_processing_resolution: Optional[int] = None,
+ prediction_type: str | None = None,
+ use_full_z_range: bool | None = True,
+ default_denoising_steps: int | None = None,
+ default_processing_resolution: int | None = None,
):
super().__init__()
@@ -189,9 +189,9 @@ def check_inputs(
resample_method_input: str,
resample_method_output: str,
batch_size: int,
- ensembling_kwargs: Optional[Dict[str, Any]],
- latents: Optional[torch.Tensor],
- generator: Optional[Union[torch.Generator, List[torch.Generator]]],
+ ensembling_kwargs: dict[str, Any] | None,
+ latents: torch.Tensor | None,
+ generator: torch.Generator | list[torch.Generator] | None,
output_type: str,
output_uncertainty: bool,
) -> int:
@@ -334,16 +334,16 @@ def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
def __call__(
self,
image: PipelineImageInput,
- num_inference_steps: Optional[int] = None,
+ num_inference_steps: int | None = None,
ensemble_size: int = 1,
- processing_resolution: Optional[int] = None,
+ processing_resolution: int | None = None,
match_input_resolution: bool = True,
resample_method_input: str = "bilinear",
resample_method_output: str = "bilinear",
batch_size: int = 1,
- ensembling_kwargs: Optional[Dict[str, Any]] = None,
- latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ ensembling_kwargs: dict[str, Any] | None = None,
+ latents: torch.Tensor | list[torch.Tensor] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
output_type: str = "np",
output_uncertainty: bool = False,
output_latent: bool = False,
@@ -353,8 +353,8 @@ def __call__(
Function invoked when calling the pipeline.
Args:
- image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
- `List[torch.Tensor]`: An input image or images used as an input for the normals estimation task. For
+ image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`),
+ `list[torch.Tensor]`: An input image or images used as an input for the normals estimation task. For
arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
@@ -386,7 +386,7 @@ def __call__(
latents (`torch.Tensor`, *optional*, defaults to `None`):
Latent noise tensors to replace the random initialization. These can be taken from the previous
function call's output.
- generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`):
+ generator (`torch.Generator`, or `list[torch.Generator]`, *optional*, defaults to `None`):
Random number generator object to ensure reproducibility.
output_type (`str`, *optional*, defaults to `"np"`):
Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
@@ -595,11 +595,11 @@ def __call__(
def prepare_latents(
self,
image: torch.Tensor,
- latents: Optional[torch.Tensor],
- generator: Optional[torch.Generator],
+ latents: torch.Tensor | None,
+ generator: torch.Generator | None,
ensemble_size: int,
batch_size: int,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> tuple[torch.Tensor, torch.Tensor]:
def retrieve_latents(encoder_output):
if hasattr(encoder_output, "latent_dist"):
return encoder_output.latent_dist.mode()
@@ -660,7 +660,7 @@ def normalize_normals(normals: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
@staticmethod
def ensemble_normals(
normals: torch.Tensor, output_uncertainty: bool, reduction: str = "closest"
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
"""
Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is
the number of ensemble members for a given prediction of size `(H x W)`.
diff --git a/src/diffusers/pipelines/mochi/pipeline_mochi.py b/src/diffusers/pipelines/mochi/pipeline_mochi.py
index 19a36c73f9ed..e8acc0a75e4d 100644
--- a/src/diffusers/pipelines/mochi/pipeline_mochi.py
+++ b/src/diffusers/pipelines/mochi/pipeline_mochi.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -78,10 +78,10 @@ def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -96,15 +96,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -196,11 +196,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -253,25 +253,25 @@ def _get_t5_prompt_embeds(
# Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -498,33 +498,33 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_frames: int = 19,
num_inference_steps: int = 64,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 4.5,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to `self.default_height`):
@@ -536,7 +536,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -548,7 +548,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -579,7 +579,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/mochi/pipeline_output.py b/src/diffusers/pipelines/mochi/pipeline_output.py
index d15827bc0084..5068cf930aaa 100644
--- a/src/diffusers/pipelines/mochi/pipeline_output.py
+++ b/src/diffusers/pipelines/mochi/pipeline_output.py
@@ -11,8 +11,8 @@ class MochiPipelineOutput(BaseOutput):
Output class for Mochi pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
diff --git a/src/diffusers/pipelines/musicldm/pipeline_musicldm.py b/src/diffusers/pipelines/musicldm/pipeline_musicldm.py
index c909e5eb0d26..e7747a4f8c3d 100644
--- a/src/diffusers/pipelines/musicldm/pipeline_musicldm.py
+++ b/src/diffusers/pipelines/musicldm/pipeline_musicldm.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -106,9 +106,9 @@ class MusicLDMPipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusi
def __init__(
self,
vae: AutoencoderKL,
- text_encoder: Union[ClapTextModelWithProjection, ClapModel],
- tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast],
- feature_extractor: Optional[ClapFeatureExtractor],
+ text_encoder: ClapTextModelWithProjection | ClapModel,
+ tokenizer: RobertaTokenizer | RobertaTokenizerFast,
+ feature_extractor: ClapFeatureExtractor | None,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
vocoder: SpeechT5HifiGan,
@@ -133,14 +133,14 @@ def _encode_prompt(
num_waveforms_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device (`torch.device`):
torch device
@@ -148,7 +148,7 @@ def _encode_prompt(
number of waveforms that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the audio generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -207,7 +207,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -434,28 +434,28 @@ def enable_model_cpu_offload(self, gpu_id=0):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- audio_length_in_s: Optional[float] = None,
+ prompt: str | list[str] = None,
+ audio_length_in_s: float | None = None,
num_inference_steps: int = 200,
guidance_scale: float = 2.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_waveforms_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_waveforms_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- output_type: Optional[str] = "np",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ output_type: str | None = "np",
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide audio generation. If not defined, you need to pass `prompt_embeds`.
audio_length_in_s (`int`, *optional*, defaults to 10.24):
The length of the generated audio sample in seconds.
@@ -465,7 +465,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 2.0):
A higher guidance scale value encourages the model to generate audio that is closely linked to the text
`prompt` at the expense of lower sound quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in audio generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_waveforms_per_prompt (`int`, *optional*, defaults to 1):
@@ -477,7 +477,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/omnigen/pipeline_omnigen.py b/src/diffusers/pipelines/omnigen/pipeline_omnigen.py
index 96c209813f54..bdf4e30c6619 100644
--- a/src/diffusers/pipelines/omnigen/pipeline_omnigen.py
+++ b/src/diffusers/pipelines/omnigen/pipeline_omnigen.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -59,10 +59,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -77,15 +77,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -170,9 +170,9 @@ def __init__(
def encode_input_images(
self,
- input_pixel_values: List[torch.Tensor],
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ input_pixel_values: list[torch.Tensor],
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
"""
get the continue embedding of input images by VAE
@@ -331,32 +331,32 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- input_images: Union[PipelineImageInput, List[PipelineImageInput]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str],
+ input_images: PipelineImageInput | list[PipelineImageInput] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
max_input_image_size: int = 1024,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 2.5,
img_guidance_scale: float = 1.6,
use_input_image_size_as_output: bool = False,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If the input includes images, need to add
placeholders `![]()
<|image_i|>` in the prompt to indicate the position of the i-th images.
- input_images (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*):
+ input_images (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*):
The list of input images. We will replace the "<|image_i|>" in prompt with the i-th image in list.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The height in pixels of the generated image. This is set to 1024 by default for the best results.
@@ -367,7 +367,7 @@ def __call__(
expense of slower inference.
max_input_image_size (`int`, *optional*, defaults to 1024):
the maximum size of input image, which will be used to crop the input image to the maximum size
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -384,7 +384,7 @@ def __call__(
e.g., image editing task
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -401,7 +401,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/omnigen/processor_omnigen.py b/src/diffusers/pipelines/omnigen/processor_omnigen.py
index 7ed11871bb2a..8a7f53a77136 100644
--- a/src/diffusers/pipelines/omnigen/processor_omnigen.py
+++ b/src/diffusers/pipelines/omnigen/processor_omnigen.py
@@ -13,7 +13,6 @@
# limitations under the License.
import re
-from typing import Dict, List
import numpy as np
import torch
@@ -132,8 +131,8 @@ def add_prefix_instruction(self, prompt):
def __call__(
self,
- instructions: List[str],
- input_images: List[List[str]] = None,
+ instructions: list[str],
+ input_images: list[list[str]] = None,
height: int = 1024,
width: int = 1024,
negative_prompt: str = "low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers.",
@@ -141,7 +140,7 @@ def __call__(
separate_cfg_input: bool = False,
use_input_image_size_as_output: bool = False,
num_images_per_prompt: int = 1,
- ) -> Dict:
+ ) -> dict:
if isinstance(instructions, str):
instructions = [instructions]
input_images = [input_images]
diff --git a/src/diffusers/pipelines/onnx_utils.py b/src/diffusers/pipelines/onnx_utils.py
index 74e9f0b97800..657750f90add 100644
--- a/src/diffusers/pipelines/onnx_utils.py
+++ b/src/diffusers/pipelines/onnx_utils.py
@@ -14,11 +14,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
import os
import shutil
from pathlib import Path
-from typing import Optional, Union
import numpy as np
from huggingface_hub import hf_hub_download
@@ -61,7 +61,7 @@ def __call__(self, **kwargs):
return self.model.run(None, inputs)
@staticmethod
- def load_model(path: Union[str, Path], provider=None, sess_options=None, provider_options=None):
+ def load_model(path: str | Path, provider=None, sess_options=None, provider_options=None):
"""
Loads an ONNX Inference session with an ExecutionProvider. Default provider is `CPUExecutionProvider`
@@ -84,7 +84,7 @@ def load_model(path: Union[str, Path], provider=None, sess_options=None, provide
path, providers=[provider], sess_options=sess_options, provider_options=provider_options
)
- def _save_pretrained(self, save_directory: Union[str, Path], file_name: Optional[str] = None, **kwargs):
+ def _save_pretrained(self, save_directory: str | Path, file_name: str | None = None, **kwargs):
"""
Save a model and its configuration file to a directory, so that it can be re-loaded using the
[`~optimum.onnxruntime.modeling_ort.ORTModel.from_pretrained`] class method. It will always save the
@@ -117,7 +117,7 @@ def _save_pretrained(self, save_directory: Union[str, Path], file_name: Optional
def save_pretrained(
self,
- save_directory: Union[str, os.PathLike],
+ save_directory: str | os.PathLike,
**kwargs,
):
"""
@@ -141,14 +141,14 @@ def save_pretrained(
@validate_hf_hub_args
def _from_pretrained(
cls,
- model_id: Union[str, Path],
- token: Optional[Union[bool, str, None]] = None,
- revision: Optional[Union[str, None]] = None,
+ model_id: str | Path,
+ token: bool | str | None | None = None,
+ revision: str | None | None = None,
force_download: bool = False,
- cache_dir: Optional[str] = None,
- file_name: Optional[str] = None,
- provider: Optional[str] = None,
- sess_options: Optional["ort.SessionOptions"] = None,
+ cache_dir: str | None = None,
+ file_name: str | None = None,
+ provider: str | None = None,
+ sess_options: "ort.SessionOptions" | None = None,
**kwargs,
):
"""
@@ -161,7 +161,7 @@ def _from_pretrained(
Is needed to load models from a private or gated repository
revision (`str`):
Revision is the specific model version to use. It can be a branch name, a tag name, or a commit id
- cache_dir (`Union[str, Path]`, *optional*):
+ cache_dir (`str | Path`, *optional*):
Path to a directory in which a downloaded pretrained model configuration should be cached if the
standard cache should not be used.
force_download (`bool`, *optional*, defaults to `False`):
@@ -210,10 +210,10 @@ def _from_pretrained(
@validate_hf_hub_args
def from_pretrained(
cls,
- model_id: Union[str, Path],
+ model_id: str | Path,
force_download: bool = True,
- token: Optional[str] = None,
- cache_dir: Optional[str] = None,
+ token: str | None = None,
+ cache_dir: str | None = None,
**model_kwargs,
):
revision = None
diff --git a/src/diffusers/pipelines/ovis_image/pipeline_output.py b/src/diffusers/pipelines/ovis_image/pipeline_output.py
index 160c5b73a917..4daecf3b2197 100644
--- a/src/diffusers/pipelines/ovis_image/pipeline_output.py
+++ b/src/diffusers/pipelines/ovis_image/pipeline_output.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -27,9 +26,9 @@ class OvisImagePipelineOutput(BaseOutput):
Output class for Ovis-Image pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image, np.ndarray]
diff --git a/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py b/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py
index 94d6cee93d7e..c8ff8227f27e 100644
--- a/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py
+++ b/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -69,10 +69,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -87,15 +87,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -181,7 +181,7 @@ def __init__(
def _get_messages(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
messages = []
@@ -200,10 +200,10 @@ def _get_messages(
def _get_ovis_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -239,10 +239,10 @@ def _get_ovis_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
):
r"""
@@ -413,33 +413,33 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = "",
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = "",
guidance_scale: float = 5.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- output_type: Optional[str] = "pil",
+ sigmas: list[float] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
not greater than `1`).
@@ -453,13 +453,13 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
diff --git a/src/diffusers/pipelines/pag/pag_utils.py b/src/diffusers/pipelines/pag/pag_utils.py
index 8a56961f321c..41395ece7421 100644
--- a/src/diffusers/pipelines/pag/pag_utils.py
+++ b/src/diffusers/pipelines/pag/pag_utils.py
@@ -13,7 +13,6 @@
# limitations under the License.
import re
-from typing import Dict, List, Tuple, Union
import torch
import torch.nn as nn
@@ -112,7 +111,7 @@ def _apply_perturbed_attention_guidance(
return_pred_text (bool): Whether to return the text noise prediction.
Returns:
- Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: The updated noise prediction tensor after applying
+ torch.Tensor | tuple[torch.Tensor, torch.Tensor]: The updated noise prediction tensor after applying
perturbed attention guidance and the text noise prediction.
"""
pag_scale = self._get_pag_scale(t)
@@ -151,8 +150,8 @@ def _prepare_perturbed_attention_guidance(self, cond, uncond, do_classifier_free
def set_pag_applied_layers(
self,
- pag_applied_layers: Union[str, List[str]],
- pag_attn_processors: Tuple[AttentionProcessor, AttentionProcessor] = (
+ pag_applied_layers: str | list[str],
+ pag_attn_processors: tuple[AttentionProcessor, AttentionProcessor] = (
PAGCFGIdentitySelfAttnProcessor2_0(),
PAGIdentitySelfAttnProcessor2_0(),
),
@@ -161,7 +160,7 @@ def set_pag_applied_layers(
Set the self-attention layers to apply PAG. Raise ValueError if the input is invalid.
Args:
- pag_applied_layers (`str` or `List[str]`):
+ pag_applied_layers (`str` or `list[str]`):
One or more strings identifying the layer names, or a simple regex for matching multiple layers, where
PAG is to be applied. A few ways of expected usage are as follows:
- Single layers specified as - "blocks.{layer_index}"
@@ -169,7 +168,7 @@ def set_pag_applied_layers(
- Multiple layers as a block name - "mid"
- Multiple layers as regex - "blocks.({layer_index_1}|{layer_index_2})"
pag_attn_processors:
- (`Tuple[AttentionProcessor, AttentionProcessor]`, defaults to `(PAGCFGIdentitySelfAttnProcessor2_0(),
+ (`tuple[AttentionProcessor, AttentionProcessor]`, defaults to `(PAGCFGIdentitySelfAttnProcessor2_0(),
PAGIdentitySelfAttnProcessor2_0())`): A tuple of two attention processors. The first attention
processor is for PAG with Classifier-free guidance enabled (conditional and unconditional). The second
attention processor is for PAG with CFG disabled (unconditional only).
@@ -214,7 +213,7 @@ def do_perturbed_attention_guidance(self) -> bool:
return self._pag_scale > 0 and len(self.pag_applied_layers) > 0
@property
- def pag_attn_processors(self) -> Dict[str, AttentionProcessor]:
+ def pag_attn_processors(self) -> dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of PAG attention processors: A dictionary contains all PAG attention processors used in the model
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py
index 389927aafcbc..807c42d21bb4 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py
@@ -14,7 +14,7 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -108,10 +108,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -126,15 +126,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -196,7 +196,7 @@ class StableDiffusionControlNetPAGPipeline(
A `CLIPTokenizer` to tokenize text.
unet ([`UNet2DConditionModel`]):
A `UNet2DConditionModel` to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -222,13 +222,13 @@ def __init__(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
requires_safety_checker: bool = True,
- pag_applied_layers: Union[str, List[str]] = "mid",
+ pag_applied_layers: str | list[str] = "mid",
):
super().__init__()
@@ -279,16 +279,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -296,7 +296,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -395,7 +395,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -864,35 +864,33 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -900,10 +898,10 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -919,18 +917,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -938,7 +936,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -952,7 +950,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -965,16 +963,16 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
The ControlNet encoder tries to recognize the content of the input image even if you remove all
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
@@ -984,7 +982,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py
index 2781af789018..ebc2e882868c 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py
@@ -15,7 +15,7 @@
# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -119,7 +119,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -170,7 +170,7 @@ class StableDiffusionControlNetPAGInpaintPipeline(
A `CLIPTokenizer` to tokenize text.
unet ([`UNet2DConditionModel`]):
A `UNet2DConditionModel` to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -196,13 +196,13 @@ def __init__(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
requires_safety_checker: bool = True,
- pag_applied_layers: Union[str, List[str]] = "mid",
+ pag_applied_layers: str | list[str] = "mid",
):
super().__init__()
@@ -255,16 +255,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -272,7 +272,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -371,7 +371,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -973,36 +973,34 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 1.0,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 0.5,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 0.5,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -1010,25 +1008,25 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`,
- `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`,
+ `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, NumPy array or tensor representing an image batch to be used as the starting point. For both
NumPy array and PyTorch tensor, the expected value range is between `[0, 1]`. If it's a tensor or a
list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a NumPy array or
a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`. It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`,
- `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`,
+ `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, NumPy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a NumPy array or PyTorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for PyTorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for NumPy array, it would be for `(B, H, W, 1)`, `(B, H, W)`, `(H,
W, 1)`, or `(H, W)`.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`,
- `List[List[torch.Tensor]]`, or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`,
+ `list[list[torch.Tensor]]`, or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -1058,7 +1056,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1066,7 +1064,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1080,7 +1078,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1093,13 +1091,13 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 0.5):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 0.5):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
@@ -1109,7 +1107,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py
index 8b7df89f039c..b0e2c03faed7 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py
@@ -14,7 +14,7 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -125,10 +125,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -143,15 +143,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -218,7 +218,7 @@ class StableDiffusionXLControlNetPAGPipeline(
A `CLIPTokenizer` to tokenize text.
unet ([`UNet2DConditionModel`]):
A `UNet2DConditionModel` to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the `unet` during the denoising process. If you set multiple
ControlNets as a list, the outputs from each ControlNet are added together to create one combined
additional conditioning.
@@ -262,13 +262,13 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
- pag_applied_layers: Union[str, List[str]] = "mid", # ["down.block_2", "up.block_1.attentions_0"], "mid"
+ pag_applied_layers: str | list[str] = "mid", # ["down.block_2", "up.block_1.attentions_0"], "mid"
):
super().__init__()
@@ -306,26 +306,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -334,11 +334,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -456,7 +456,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -1002,45 +1002,43 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -1048,13 +1046,13 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
@@ -1072,11 +1070,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1090,10 +1088,10 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 5.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2`
and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders.
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1101,7 +1099,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1122,7 +1120,7 @@ def __call__(
weighting). If not provided, pooled `negative_prompt_embeds` are generated from `negative_prompt` input
argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1135,39 +1133,39 @@ def __call__(
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1180,7 +1178,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py
index 5a6b8d5e9f37..8967e50251b9 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py
@@ -14,7 +14,7 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -149,7 +149,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -202,7 +202,7 @@ class StableDiffusionXLControlNetPAGImg2ImgPipeline(
Second Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
- controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
+ controlnet ([`ControlNetModel`] or `list[ControlNetModel]`):
Provides additional conditioning to the unet during the denoising process. If you set multiple ControlNets
as a list, the outputs from each ControlNet are added together to create one combined additional
conditioning.
@@ -250,14 +250,14 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
+ controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel,
scheduler: KarrasDiffusionSchedulers,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
- pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"]
+ pag_applied_layers: str | list[str] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"]
):
super().__init__()
@@ -297,26 +297,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -325,11 +325,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -447,7 +447,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -1079,47 +1079,45 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.8,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ controlnet_conditioning_scale: float | list[float] = 0.8,
guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -1127,18 +1125,18 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The initial image will be used as the starting point for the image generation process. Can also accept
image latents as `image`, if passing latents directly, it will not be encoded again.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If
the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also
be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height
@@ -1168,11 +1166,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1180,7 +1178,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1202,7 +1200,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1217,42 +1215,42 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original unet. If multiple ControlNets are specified in init, you can set the
corresponding scale as a list.
guess_mode (`bool`, *optional*, defaults to `False`):
In this mode, the ControlNet encoder will try best to recognize the content of the input image even if
you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+ control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the controlnet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+ control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the controlnet stops applying.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1273,7 +1271,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py b/src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py
index 6704924b2512..15ac665acd2b 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import numpy as np
import torch
@@ -164,10 +164,10 @@ class HunyuanDiTPAGPipeline(DiffusionPipeline, PAGMixin):
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. We use
`sdxl-vae-fp16-fix`.
- text_encoder (Optional[`~transformers.BertModel`, `~transformers.CLIPTextModel`]):
+ text_encoder (`~transformers.BertModel`, `~transformers.CLIPTextModel` | None):
Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
HunyuanDiT uses a fine-tuned [bilingual CLIP].
- tokenizer (Optional[`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer`]):
+ tokenizer (`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer` | None):
A `BertTokenizer` or `CLIPTokenizer` to tokenize text.
transformer ([`HunyuanDiT2DModel`]):
The HunyuanDiT model designed by Tencent Hunyuan.
@@ -204,12 +204,12 @@ def __init__(
tokenizer: BertTokenizer,
transformer: HunyuanDiT2DModel,
scheduler: DDPMScheduler,
- safety_checker: Optional[StableDiffusionSafetyChecker] = None,
- feature_extractor: Optional[CLIPImageProcessor] = None,
+ safety_checker: StableDiffusionSafetyChecker | None = None,
+ feature_extractor: CLIPImageProcessor | None = None,
requires_safety_checker: bool = True,
- text_encoder_2: Optional[T5EncoderModel] = None,
- tokenizer_2: Optional[T5Tokenizer] = None,
- pag_applied_layers: Union[str, List[str]] = "blocks.1", # "blocks.16.attn1", "blocks.16", "16", 16
+ text_encoder_2: T5EncoderModel | None = None,
+ tokenizer_2: T5Tokenizer | None = None,
+ pag_applied_layers: str | list[str] = "blocks.1", # "blocks.16.attn1", "blocks.16", "16", 16
):
super().__init__()
@@ -262,19 +262,19 @@ def encode_prompt(
dtype: torch.dtype = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- max_sequence_length: Optional[int] = None,
+ negative_prompt: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ max_sequence_length: int | None = None,
text_encoder_index: int = 0,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -284,7 +284,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -373,7 +373,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -580,34 +580,35 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_2: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_2: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- prompt_attention_mask_2: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask_2: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
+ num_inference_steps: int = 50,
+ guidance_scale: float = 5.0,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int = 1,
+ eta: float = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_2: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_2: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ prompt_attention_mask_2: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask_2: torch.Tensor | None = None,
+ output_type: str = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int, dict], None]
+ | PipelineCallback
+ | MultiPipelineCallbacks
+ | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
guidance_rescale: float = 0.0,
- original_size: Optional[Tuple[int, int]] = (1024, 1024),
- target_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
+ original_size: tuple[int, int] = (1024, 1024),
+ target_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
use_resolution_binning: bool = True,
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
@@ -616,7 +617,7 @@ def __call__(
The call function to the pipeline for generation with HunyuanDiT.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`):
The height in pixels of the generated image.
@@ -628,7 +629,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -636,7 +637,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -664,19 +665,19 @@ def __call__(
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
plain tuple.
- callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+ callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
A callback function or a list of callback functions to be called at the end of each denoising step.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
A list of tensor inputs that should be passed to the callback function. If not defined, all tensor
inputs will be passed.
guidance_rescale (`float`, *optional*, defaults to 0.0):
Rescale the noise_cfg according to `guidance_rescale`. Based on findings of [Common Diffusion Noise
Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). See Section 3.4
- original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
+ original_size (`tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
The original size of the image. Used to calculate the time ids.
- target_size (`Tuple[int, int]`, *optional*):
+ target_size (`tuple[int, int]`, *optional*):
The target size of the image. Used to calculate the time ids.
- crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`):
+ crops_coords_top_left (`tuple[int, int]`, *optional*, defaults to `(0, 0)`):
The top left coordinates of the crop. Used to calculate the time ids.
use_resolution_binning (`bool`, *optional*, defaults to `True`):
Whether to use resolution binning or not. If `True`, the input resolution will be mapped to the closest
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_kolors.py b/src/diffusers/pipelines/pag/pipeline_pag_kolors.py
index 5b82d546445b..4f138d91d9c6 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_kolors.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_kolors.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
@@ -68,10 +68,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -86,15 +86,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -154,7 +154,7 @@ class KolorsPAGPipeline(
force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"False"`):
Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of
`Kwai-Kolors/Kolors-diffusers`.
- pag_applied_layers (`str` or `List[str]``, *optional*, defaults to `"mid"`):
+ pag_applied_layers (`str` or `list[str]``, *optional*, defaults to `"mid"`):
Set the transformer attention layers where to apply the perturbed attention guidance. Can be a string or a
list of strings with "down", "mid", "up", a whole transformer block or specific transformer block attention
layers, e.g.:
@@ -187,7 +187,7 @@ def __init__(
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
force_zeros_for_empty_prompt: bool = False,
- pag_applied_layers: Union[str, List[str]] = "mid",
+ pag_applied_layers: str | list[str] = "mid",
):
super().__init__()
@@ -216,21 +216,21 @@ def __init__(
def encode_prompt(
self,
prompt,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
negative_prompt=None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 256,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -238,7 +238,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -308,7 +308,7 @@ def encode_prompt(
if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
negative_prompt_embeds = torch.zeros_like(prompt_embeds)
elif do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -666,38 +666,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
max_sequence_length: int = 256,
@@ -706,7 +704,7 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -722,11 +720,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -743,7 +741,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -752,7 +750,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -774,7 +772,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -788,31 +786,31 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -822,7 +820,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py b/src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py
index 283862989c71..2e7e8ab9814f 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py
@@ -16,7 +16,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import T5EncoderModel, T5Tokenizer
@@ -84,10 +84,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -102,15 +102,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -173,7 +173,7 @@ def __init__(
vae: AutoencoderKL,
transformer: PixArtTransformer2DModel,
scheduler: KarrasDiffusionSchedulers,
- pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block
+ pag_applied_layers: str | list[str] = "blocks.1", # 1st transformer block
):
super().__init__()
@@ -189,15 +189,15 @@ def __init__(
# Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->300
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
**kwargs,
@@ -206,9 +206,9 @@ def encode_prompt(
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -575,51 +575,51 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 20,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 4.5,
- num_images_per_prompt: Optional[int] = 1,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ num_images_per_prompt: int | None = 1,
+ height: int | None = None,
+ width: int | None = None,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
clean_caption: bool = True,
use_resolution_binning: bool = True,
max_sequence_length: int = 300,
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -638,7 +638,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sana.py b/src/diffusers/pipelines/pag/pipeline_pag_sana.py
index 466996889417..71861f366477 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sana.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sana.py
@@ -17,7 +17,7 @@
import re
import urllib.parse as ul
import warnings
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
@@ -88,10 +88,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -106,15 +106,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -161,12 +161,12 @@ class SanaPAGPipeline(DiffusionPipeline, PAGMixin):
def __init__(
self,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
text_encoder: Gemma2PreTrainedModel,
vae: AutoencoderDC,
transformer: SanaTransformer2DModel,
scheduler: FlowMatchEulerDiscreteScheduler,
- pag_applied_layers: Union[str, List[str]] = "transformer_blocks.0",
+ pag_applied_layers: str | list[str] = "transformer_blocks.0",
):
super().__init__()
@@ -241,26 +241,26 @@ def disable_vae_tiling(self):
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
+ complex_human_instruction: list[str] | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -649,30 +649,30 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 20,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 4.5,
- num_images_per_prompt: Optional[int] = 1,
+ num_images_per_prompt: int | None = 1,
height: int = 1024,
width: int = 1024,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
clean_caption: bool = False,
use_resolution_binning: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 300,
- complex_human_instruction: List[str] = [
+ complex_human_instruction: list[str] = [
"Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:",
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
@@ -684,26 +684,26 @@ def __call__(
],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 20):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -722,7 +722,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -756,12 +756,12 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to 300): Maximum sequence length to use with the `prompt`.
- complex_human_instruction (`List[str]`, *optional*):
+ complex_human_instruction (`list[str]`, *optional*):
Instructions for complex human attention:
https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.
pag_scale (`float`, *optional*, defaults to 3.0):
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd.py b/src/diffusers/pipelines/pag/pipeline_pag_sd.py
index 67676fb28798..26ea717556c5 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from packaging import version
@@ -97,10 +97,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -115,15 +115,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -212,7 +212,7 @@ def __init__(
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
requires_safety_checker: bool = True,
- pag_applied_layers: Union[str, List[str]] = "mid",
+ pag_applied_layers: str | list[str] = "mid",
):
super().__init__()
@@ -308,16 +308,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -325,7 +325,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -424,7 +424,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -746,29 +746,29 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -776,7 +776,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -785,18 +785,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -804,7 +804,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -818,7 +818,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -843,7 +843,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py
index 303a0a2f0b2e..f0fbef29b699 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import (
@@ -76,10 +76,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -94,15 +94,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -185,7 +185,7 @@ def __init__(
tokenizer_2: CLIPTokenizer,
text_encoder_3: T5EncoderModel,
tokenizer_3: T5TokenizerFast,
- pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block
+ pag_applied_layers: str | list[str] = "blocks.1", # 1st transformer block
):
super().__init__()
@@ -221,11 +221,11 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -278,10 +278,10 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- clip_skip: Optional[int] = None,
+ device: torch.device | None = None,
+ clip_skip: int | None = None,
clip_model_index: int = 0,
):
device = device or self._execution_device
@@ -334,32 +334,32 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]],
- prompt_3: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ prompt_3: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ clip_skip: int | None = None,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -368,14 +368,14 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
prompt_embeds (`torch.FloatTensor`, *optional*):
@@ -685,30 +685,30 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
@@ -717,13 +717,13 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
will be used instead
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -733,7 +733,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -743,19 +743,19 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used instead
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used instead
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -791,7 +791,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py
index 2005c865c22b..84b727dc0613 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -77,7 +77,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -92,10 +92,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -110,15 +110,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -201,7 +201,7 @@ def __init__(
tokenizer_2: CLIPTokenizer,
text_encoder_3: T5EncoderModel,
tokenizer_3: T5TokenizerFast,
- pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block
+ pag_applied_layers: str | list[str] = "blocks.1", # 1st transformer block
):
super().__init__()
@@ -237,11 +237,11 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -294,10 +294,10 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- clip_skip: Optional[int] = None,
+ device: torch.device | None = None,
+ clip_skip: int | None = None,
clip_model_index: int = 0,
):
device = device or self._execution_device
@@ -350,32 +350,32 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]],
- prompt_3: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ prompt_3: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ clip_skip: int | None = None,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -384,14 +384,14 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
prompt_embeds (`torch.FloatTensor`, *optional*):
@@ -736,32 +736,32 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
image: PipelineImageInput = None,
strength: float = 0.6,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
@@ -770,16 +770,16 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
will be used instead
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -794,7 +794,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -804,19 +804,19 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used instead
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used instead
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -852,7 +852,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py
index de13be9c4d22..62d1c912283f 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
@@ -135,12 +135,12 @@ def __init__(
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetMotionModel],
+ unet: UNet2DConditionModel | UNetMotionModel,
motion_adapter: MotionAdapter,
scheduler: KarrasDiffusionSchedulers,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
- pag_applied_layers: Union[str, List[str]] = "mid_block.*attn1", # ["mid"], ["down_blocks.1"]
+ pag_applied_layers: str | list[str] = "mid_block.*attn1", # ["mid"], ["down_blocks.1"]
):
super().__init__()
if isinstance(unet, UNet2DConditionModel):
@@ -169,16 +169,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -186,7 +186,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -285,7 +285,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -576,27 +576,27 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- num_frames: Optional[int] = 16,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ num_frames: int | None = 16,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
decode_chunk_size: int = 16,
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
@@ -605,7 +605,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated video.
@@ -620,13 +620,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -642,7 +642,7 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -663,7 +663,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py
index 42b5db0fa762..822483eca995 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -77,7 +77,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -92,10 +92,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -110,15 +110,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -207,7 +207,7 @@ def __init__(
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
requires_safety_checker: bool = True,
- pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"]
+ pag_applied_layers: str | list[str] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"]
):
super().__init__()
@@ -303,16 +303,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -320,7 +320,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -419,7 +419,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -783,29 +783,27 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
strength: float = 0.8,
- num_inference_steps: Optional[int] = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- guidance_scale: Optional[float] = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ num_inference_steps: int | None = 50,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ guidance_scale: float | None = 7.5,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
clip_skip: int = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -813,9 +811,9 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -830,18 +828,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference. This parameter is modulated by `strength`.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -849,7 +847,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -859,7 +857,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -880,7 +878,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py
index cf8c4972762f..0f6fbbd9ae16 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -82,7 +82,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -124,10 +124,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -142,15 +142,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -239,7 +239,7 @@ def __init__(
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
requires_safety_checker: bool = True,
- pag_applied_layers: Union[str, List[str]] = "mid",
+ pag_applied_layers: str | list[str] = "mid",
):
super().__init__()
@@ -338,16 +338,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -355,7 +355,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -454,7 +454,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -911,34 +911,34 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: torch.Tensor = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.9999,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -946,7 +946,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -955,18 +955,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -974,7 +974,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -988,7 +988,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1013,7 +1013,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py
index 0613ec23f740..2987c90626ef 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import (
@@ -112,10 +112,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -130,15 +130,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -254,8 +254,8 @@ def __init__(
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
- pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"],["down.block_1"],["up.block_0.attentions_0"]
+ add_watermarker: bool | None = None,
+ pag_applied_layers: str | list[str] = "mid", # ["mid"],["down.block_1"],["up.block_0.attentions_0"]
):
super().__init__()
@@ -293,26 +293,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -321,11 +321,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -443,7 +443,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -835,40 +835,40 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -876,10 +876,10 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -895,11 +895,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -916,11 +916,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -928,7 +928,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -950,7 +950,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -971,31 +971,31 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1005,7 +1005,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py
index 1081993f46e6..433b9edc69b7 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -116,7 +116,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -131,10 +131,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -149,15 +149,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -277,8 +277,8 @@ def __init__(
feature_extractor: CLIPImageProcessor = None,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
- pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"]
+ add_watermarker: bool | None = None,
+ pag_applied_layers: str | list[str] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"]
):
super().__init__()
@@ -311,26 +311,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -339,11 +339,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -461,7 +461,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -988,45 +988,43 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
strength: float = 0.3,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_start: Optional[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_start: float | None = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -1034,13 +1032,13 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
- image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`):
+ image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[np.ndarray]`):
The image(s) to modify with the pipeline.
strength (`float`, *optional*, defaults to 0.3):
Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`
@@ -1052,11 +1050,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1081,11 +1079,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1093,7 +1091,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1115,7 +1113,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1136,31 +1134,31 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1181,7 +1179,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py
index f6c4982c1c6c..9caf50e5e333 100644
--- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py
+++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -129,7 +129,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -144,10 +144,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -162,15 +162,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -293,8 +293,8 @@ def __init__(
feature_extractor: CLIPImageProcessor = None,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
- pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"]
+ add_watermarker: bool | None = None,
+ pag_applied_layers: str | list[str] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"]
):
super().__init__()
@@ -401,26 +401,26 @@ def prepare_ip_adapter_image_embeds(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -429,11 +429,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -551,7 +551,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -1079,50 +1079,48 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: torch.Tensor = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.9999,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_start: Optional[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_start: float | None = None,
+ denoising_end: float | None = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
pag_scale: float = 3.0,
pag_adaptive_scale: float = 0.0,
):
@@ -1130,10 +1128,10 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
image (`PIL.Image.Image`):
@@ -1172,11 +1170,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1201,11 +1199,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -1223,7 +1221,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1250,31 +1248,31 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1295,7 +1293,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/paint_by_example/__init__.py b/src/diffusers/pipelines/paint_by_example/__init__.py
index aaa775f690c3..d2906b540c6e 100644
--- a/src/diffusers/pipelines/paint_by_example/__init__.py
+++ b/src/diffusers/pipelines/paint_by_example/__init__.py
@@ -1,5 +1,5 @@
from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING
import numpy as np
import PIL
diff --git a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py
index c09992befbcb..aa7dbaa720e5 100644
--- a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py
+++ b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -43,7 +43,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -65,7 +65,7 @@ def prepare_mask_and_masked_image(image, mask):
binarized (``mask > 0.5``) and cast to ``torch.float32`` too.
Args:
- image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint.
+ image (np.array | PIL.Image | torch.Tensor): The image to inpaint.
It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width``
``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``.
mask (_type_): The mask to apply to the image, i.e. regions to inpaint.
@@ -198,7 +198,7 @@ def __init__(
vae: AutoencoderKL,
image_encoder: PaintByExampleImageEncoder,
unet: UNet2DConditionModel,
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = False,
@@ -270,7 +270,7 @@ def check_inputs(self, image, height, width, callback_steps):
and not isinstance(image, list)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+ "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
f" {type(image)}"
)
@@ -397,33 +397,33 @@ def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free
@torch.no_grad()
def __call__(
self,
- example_image: Union[torch.Tensor, PIL.Image.Image],
- image: Union[torch.Tensor, PIL.Image.Image],
- mask_image: Union[torch.Tensor, PIL.Image.Image],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ example_image: torch.Tensor | PIL.Image.Image,
+ image: torch.Tensor | PIL.Image.Image,
+ mask_image: torch.Tensor | PIL.Image.Image,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
):
r"""
The call function to the pipeline for generation.
Args:
- example_image (`torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]`):
+ example_image (`torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]`):
An example image to guide image generation.
- image (`torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]`):
+ image (`torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]`):
`Image` or tensor representing an image batch to be inpainted (parts of the image are masked out with
`mask_image` and repainted according to `prompt`).
- mask_image (`torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]`):
+ mask_image (`torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]`):
`Image` or tensor representing an image batch to mask `image`. White pixels in the mask are repainted,
while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel
(luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3, so the
@@ -438,7 +438,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -446,7 +446,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/pia/pipeline_pia.py b/src/diffusers/pipelines/pia/pipeline_pia.py
index dfc6e83fbd7c..d108deb9c5da 100644
--- a/src/diffusers/pipelines/pia/pipeline_pia.py
+++ b/src/diffusers/pipelines/pia/pipeline_pia.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -85,7 +85,7 @@
```
"""
-RANGE_LIST = [
+RANGE_list = [
[1.0, 0.9, 0.85, 0.85, 0.85, 0.8], # 0 Small Motion
[1.0, 0.8, 0.8, 0.8, 0.79, 0.78, 0.75], # Moderate Motion
[1.0, 0.8, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.6, 0.5, 0.5], # Large Motion
@@ -103,7 +103,7 @@ def prepare_mask_coef_by_statistics(num_frames: int, cond_frame: int, motion_sca
assert num_frames > cond_frame, "video_length should be greater than cond_frame"
- range_list = RANGE_LIST
+ range_list = RANGE_list
assert motion_scale < len(range_list), f"motion_scale type{motion_scale} not implemented"
@@ -122,13 +122,13 @@ class PIAPipelineOutput(BaseOutput):
Output class for PIAPipeline.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
Nested list of length `batch_size` with denoised PIL image sequences of length `num_frames`, NumPy array of
shape `(batch_size, num_frames, channels, height, width, Torch tensor of shape `(batch_size, num_frames,
channels, height, width)`.
"""
- frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]]
+ frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]
class PIAPipeline(
@@ -179,16 +179,14 @@ def __init__(
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
- unet: Union[UNet2DConditionModel, UNetMotionModel],
- scheduler: Union[
- DDIMScheduler,
- PNDMScheduler,
- LMSDiscreteScheduler,
- EulerDiscreteScheduler,
- EulerAncestralDiscreteScheduler,
- DPMSolverMultistepScheduler,
- ],
- motion_adapter: Optional[MotionAdapter] = None,
+ unet: UNet2DConditionModel | UNetMotionModel,
+ scheduler: DDIMScheduler
+ | PNDMScheduler
+ | LMSDiscreteScheduler
+ | EulerDiscreteScheduler
+ | EulerAncestralDiscreteScheduler
+ | DPMSolverMultistepScheduler,
+ motion_adapter: MotionAdapter | None = None,
feature_extractor: CLIPImageProcessor = None,
image_encoder: CLIPVisionModelWithProjection = None,
):
@@ -217,16 +215,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -234,7 +232,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -333,7 +331,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -674,29 +672,29 @@ def num_timesteps(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
strength: float = 1.0,
- num_frames: Optional[int] = 16,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ num_frames: int | None = 16,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
motion_scale: int = 0,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
r"""
The call function to the pipeline for generation.
@@ -704,7 +702,7 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to be used for video generation.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
strength (`float`, *optional*, defaults to 1.0):
Indicates extent to transform the reference `image`. Must be between 0 and 1.
@@ -721,13 +719,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -743,7 +741,7 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -769,7 +767,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/pipeline_flax_utils.py b/src/diffusers/pipelines/pipeline_flax_utils.py
index 2724c764c771..51dcf9a2fecf 100644
--- a/src/diffusers/pipelines/pipeline_flax_utils.py
+++ b/src/diffusers/pipelines/pipeline_flax_utils.py
@@ -17,7 +17,7 @@
import importlib
import inspect
import os
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
import flax
import numpy as np
@@ -90,12 +90,12 @@ class FlaxImagePipelineOutput(BaseOutput):
Output class for image pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
class FlaxDiffusionPipeline(ConfigMixin, PushToHubMixin):
@@ -150,8 +150,8 @@ def register_modules(self, **kwargs):
def save_pretrained(
self,
- save_directory: Union[str, os.PathLike],
- params: Union[Dict, FrozenDict],
+ save_directory: str | os.PathLike,
+ params: dict | FrozenDict,
push_to_hub: bool = False,
**kwargs,
):
@@ -168,7 +168,7 @@ class implements both a save and loading method. The pipeline is easily reloaded
Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
namespace).
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
"""
self.save_config(save_directory)
@@ -228,7 +228,7 @@ class implements both a save and loading method. The pipeline is easily reloaded
@classmethod
@validate_hf_hub_args
- def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs):
r"""
Instantiate a Flax-based diffusion pipeline from pretrained pipeline weights.
@@ -254,7 +254,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
@@ -544,7 +544,7 @@ def _get_signature_keys(cls, obj):
return expected_modules, optional_parameters
@property
- def components(self) -> Dict[str, Any]:
+ def components(self) -> dict[str, Any]:
r"""
The `self.components` property can be useful to run different pipelines with the same weights and
diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py
index 57d4eaa8f89e..b2564d25505e 100644
--- a/src/diffusers/pipelines/pipeline_loading_utils.py
+++ b/src/diffusers/pipelines/pipeline_loading_utils.py
@@ -17,7 +17,7 @@
import re
import warnings
from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import httpx
import requests
@@ -210,7 +210,7 @@ def filter_with_regex(filenames, pattern_re):
return {f for f in filenames if pattern_re.match(f.split("/")[-1]) is not None}
-def variant_compatible_siblings(filenames, variant=None, ignore_patterns=None) -> Union[List[os.PathLike], str]:
+def variant_compatible_siblings(filenames, variant=None, ignore_patterns=None) -> list[os.PathLike] | str:
weight_names = [
WEIGHTS_NAME,
SAFETENSORS_WEIGHTS_NAME,
@@ -525,12 +525,12 @@ def _get_pipeline_class(
def _load_empty_model(
library_name: str,
class_name: str,
- importable_classes: List[Any],
+ importable_classes: list[Any],
pipelines: Any,
is_pipeline_module: bool,
name: str,
- torch_dtype: Union[str, torch.dtype],
- cached_folder: Union[str, os.PathLike],
+ torch_dtype: str | torch.dtype,
+ cached_folder: str | os.PathLike,
**kwargs,
):
# retrieve class objects.
@@ -607,7 +607,7 @@ def _load_empty_model(
def _assign_components_to_devices(
- module_sizes: Dict[str, float], device_memory: Dict[str, float], device_mapping_strategy: str = "balanced"
+ module_sizes: dict[str, float], device_memory: dict[str, float], device_mapping_strategy: str = "balanced"
):
device_ids = list(device_memory.keys())
device_cycle = device_ids + device_ids[::-1]
@@ -738,28 +738,28 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic
def load_sub_model(
library_name: str,
class_name: str,
- importable_classes: List[Any],
+ importable_classes: list[Any],
pipelines: Any,
is_pipeline_module: bool,
pipeline_class: Any,
torch_dtype: torch.dtype,
provider: Any,
sess_options: Any,
- device_map: Optional[Union[Dict[str, torch.device], str]],
- max_memory: Optional[Dict[Union[int, str], Union[int, str]]],
- offload_folder: Optional[Union[str, os.PathLike]],
+ device_map: dict[str, torch.device] | str | None,
+ max_memory: dict[int | str, int | str] | None,
+ offload_folder: str | os.PathLike | None,
offload_state_dict: bool,
- model_variants: Dict[str, str],
+ model_variants: dict[str, str],
name: str,
from_flax: bool,
variant: str,
low_cpu_mem_usage: bool,
- cached_folder: Union[str, os.PathLike],
+ cached_folder: str | os.PathLike,
use_safetensors: bool,
- dduf_entries: Optional[Dict[str, DDUFEntry]],
+ dduf_entries: dict[str, DDUFEntry] | None,
provider_options: Any,
disable_mmap: bool,
- quantization_config: Optional[Any] = None,
+ quantization_config: Any | None = None,
):
"""Helper method to load the module `name` from `library_name` and `class_name`"""
from ..quantizers import PipelineQuantizationConfig
@@ -1050,10 +1050,10 @@ def get_connected_passed_kwargs(prefix):
def _get_custom_components_and_folders(
pretrained_model_name: str,
- config_dict: Dict[str, Any],
- filenames: Optional[List[str]] = None,
- variant_filenames: Optional[List[str]] = None,
- variant: Optional[str] = None,
+ config_dict: dict[str, Any],
+ filenames: list[str] | None = None,
+ variant_filenames: list[str] | None = None,
+ variant: str | None = None,
):
config_dict = config_dict.copy()
@@ -1086,15 +1086,15 @@ def _get_custom_components_and_folders(
def _get_ignore_patterns(
passed_components,
- model_folder_names: List[str],
- model_filenames: List[str],
+ model_folder_names: list[str],
+ model_filenames: list[str],
use_safetensors: bool,
from_flax: bool,
allow_pickle: bool,
use_onnx: bool,
is_onnx: bool,
- variant: Optional[str] = None,
-) -> List[str]:
+ variant: str | None = None,
+) -> list[str]:
if (
use_safetensors
and not allow_pickle
diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py
index b96305c74131..d675f1de04a7 100644
--- a/src/diffusers/pipelines/pipeline_utils.py
+++ b/src/diffusers/pipelines/pipeline_utils.py
@@ -19,9 +19,10 @@
import os
import re
import sys
+import types
from dataclasses import dataclass
from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union, get_args, get_origin
+from typing import Any, Callable, Dict, List, Union, get_args, get_origin, get_type_hints
import httpx
import numpy as np
@@ -111,7 +112,7 @@
for library in LOADABLE_CLASSES:
LIBRARIES.append(library)
-SUPPORTED_DEVICE_MAP = ["balanced"] + [get_device()]
+SUPPORTED_DEVICE_MAP = ["balanced"] + [get_device(), "cpu"]
logger = logging.get_logger(__name__)
@@ -127,7 +128,7 @@ class ImagePipelineOutput(BaseOutput):
num_channels)`.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
@dataclass
@@ -238,10 +239,10 @@ def __setattr__(self, name: str, value: Any):
def save_pretrained(
self,
- save_directory: Union[str, os.PathLike],
+ save_directory: str | os.PathLike,
safe_serialization: bool = True,
- variant: Optional[str] = None,
- max_shard_size: Optional[Union[int, str]] = None,
+ variant: str | None = None,
+ max_shard_size: int | str | None = None,
push_to_hub: bool = False,
**kwargs,
):
@@ -340,6 +341,7 @@ def is_saveable_module(name, value):
save_method_accept_safe = "safe_serialization" in save_method_signature.parameters
save_method_accept_variant = "variant" in save_method_signature.parameters
save_method_accept_max_shard_size = "max_shard_size" in save_method_signature.parameters
+ save_method_accept_peft_format = "save_peft_format" in save_method_signature.parameters
save_kwargs = {}
if save_method_accept_safe:
@@ -349,6 +351,11 @@ def is_saveable_module(name, value):
if save_method_accept_max_shard_size and max_shard_size is not None:
# max_shard_size is expected to not be None in ModelMixin
save_kwargs["max_shard_size"] = max_shard_size
+ if save_method_accept_peft_format:
+ # Set save_peft_format=False for transformers>=5.0.0 compatibility
+ # In transformers 5.0.0+, the default save_peft_format=True adds "base_model.model" prefix
+ # to adapter keys, but from_pretrained expects keys without this prefix
+ save_kwargs["save_peft_format"] = False
save_method(os.path.join(save_directory, pipeline_component_name), **save_kwargs)
@@ -467,8 +474,7 @@ def module_is_offloaded(module):
pipeline_is_sequentially_offloaded = any(
module_is_sequentially_offloaded(module) for _, module in self.components.items()
)
-
- is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1
+ is_pipeline_device_mapped = self._is_pipeline_device_mapped()
if is_pipeline_device_mapped:
raise ValueError(
"It seems like you have activated a device mapping strategy on the pipeline which doesn't allow explicit device placement using `to()`. You can call `reset_device_map()` to remove the existing device map from the pipeline."
@@ -605,7 +611,7 @@ def dtype(self) -> torch.dtype:
@classmethod
@validate_hf_hub_args
- def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs) -> Self:
+ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwargs) -> Self:
r"""
Instantiate a PyTorch diffusion pipeline from pretrained pipeline weights.
@@ -1170,7 +1176,7 @@ def remove_all_hooks(self):
accelerate.hooks.remove_hook_from_module(model, recurse=True)
self._all_hooks = []
- def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the accelerator when its
@@ -1187,7 +1193,7 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t
"""
self._maybe_raise_error_if_group_offload_active(raise_error=True)
- is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1
+ is_pipeline_device_mapped = self._is_pipeline_device_mapped()
if is_pipeline_device_mapped:
raise ValueError(
"It seems like you have activated a device mapping strategy on the pipeline so calling `enable_model_cpu_offload() isn't allowed. You can call `reset_device_map()` first and then call `enable_model_cpu_offload()`."
@@ -1288,7 +1294,7 @@ def maybe_free_model_hooks(self):
# make sure the model is in the same state as before calling it
self.enable_model_cpu_offload(device=getattr(self, "_offload_device", "cuda"))
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using 🤗 Accelerate, significantly reducing memory usage. When called, the state
dicts of all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are saved to CPU
@@ -1311,7 +1317,7 @@ def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Un
raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")
self.remove_all_hooks()
- is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1
+ is_pipeline_device_mapped = self._is_pipeline_device_mapped()
if is_pipeline_device_mapped:
raise ValueError(
"It seems like you have activated a device mapping strategy on the pipeline so calling `enable_sequential_cpu_offload() isn't allowed. You can call `reset_device_map()` first and then call `enable_sequential_cpu_offload()`."
@@ -1360,13 +1366,13 @@ def enable_group_offload(
onload_device: torch.device,
offload_device: torch.device = torch.device("cpu"),
offload_type: str = "block_level",
- num_blocks_per_group: Optional[int] = None,
+ num_blocks_per_group: int | None = None,
non_blocking: bool = False,
use_stream: bool = False,
record_stream: bool = False,
low_cpu_mem_usage=False,
- offload_to_disk_path: Optional[str] = None,
- exclude_modules: Optional[Union[str, List[str]]] = None,
+ offload_to_disk_path: str | None = None,
+ exclude_modules: str | list[str] | None = None,
) -> None:
r"""
Applies group offloading to the internal layers of a torch.nn.Module. To understand what group offloading is,
@@ -1497,7 +1503,7 @@ def reset_device_map(self):
@classmethod
@validate_hf_hub_args
- def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
+ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike:
r"""
Download and cache a PyTorch diffusion pipeline from pretrained pipeline weights.
@@ -1593,7 +1599,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
use_onnx = kwargs.pop("use_onnx", None)
load_connected_pipeline = kwargs.pop("load_connected_pipeline", False)
trust_remote_code = kwargs.pop("trust_remote_code", False)
- dduf_file: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_file", None)
+ dduf_file: dict[str, DDUFEntry] | None = kwargs.pop("dduf_file", None)
if dduf_file:
if custom_pipeline:
@@ -1617,7 +1623,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
allow_patterns = None
ignore_patterns = None
- model_info_call_error: Optional[Exception] = None
+ model_info_call_error: Exception | None = None
if not local_files_only:
try:
info = model_info(pretrained_model_name, token=token, revision=revision)
@@ -1838,19 +1844,48 @@ def _get_signature_keys(cls, obj):
@classmethod
def _get_signature_types(cls):
signature_types = {}
- for k, v in inspect.signature(cls.__init__).parameters.items():
- if inspect.isclass(v.annotation):
- signature_types[k] = (v.annotation,)
- elif get_origin(v.annotation) == Union:
- signature_types[k] = get_args(v.annotation)
- elif get_origin(v.annotation) in [List, Dict, list, dict]:
- signature_types[k] = (v.annotation,)
+ # Use get_type_hints to properly resolve string annotations (from __future__ import annotations)
+ try:
+ type_hints = get_type_hints(cls.__init__)
+ except Exception:
+ # Fallback to direct annotation access if get_type_hints fails
+ type_hints = {}
+ for k, v in inspect.signature(cls.__init__).parameters.items():
+ if v.annotation != inspect.Parameter.empty:
+ type_hints[k] = v.annotation
+
+ # Get all parameters from the signature to ensure we don't miss any
+ all_params = inspect.signature(cls.__init__).parameters
+
+ for param_name, param in all_params.items():
+ # Skip 'self' parameter
+ if param_name == "self":
+ continue
+
+ # If we have type hints, use them
+ if param_name in type_hints:
+ annotation = type_hints[param_name]
+ if inspect.isclass(annotation):
+ signature_types[param_name] = (annotation,)
+ elif get_origin(annotation) == Union:
+ signature_types[param_name] = get_args(annotation)
+ elif isinstance(annotation, types.UnionType):
+ # Handle PEP 604 union syntax (X | Y) introduced in Python 3.10+
+ signature_types[param_name] = get_args(annotation)
+ elif get_origin(annotation) in [List, Dict, list, dict]:
+ signature_types[param_name] = (annotation,)
+ else:
+ logger.warning(f"cannot get type annotation for Parameter {param_name} of {cls}.")
+ # Still add it with empty signature so it's in expected_types
+ signature_types[param_name] = (inspect.Signature.empty,)
else:
- logger.warning(f"cannot get type annotation for Parameter {k} of {cls}.")
+ # No type annotation found - add with empty signature
+ signature_types[param_name] = (inspect.Signature.empty,)
+
return signature_types
@property
- def parameters(self) -> Dict[str, Any]:
+ def parameters(self) -> dict[str, Any]:
r"""
The `self.parameters` property can be useful to run different pipelines with the same weights and
configurations without reallocating additional memory.
@@ -1880,7 +1915,7 @@ def parameters(self) -> Dict[str, Any]:
return pipeline_parameters
@property
- def components(self) -> Dict[str, Any]:
+ def components(self) -> dict[str, Any]:
r"""
The `self.components` property can be useful to run different pipelines with the same weights and
configurations without reallocating additional memory.
@@ -1947,7 +1982,7 @@ def progress_bar(self, iterable=None, total=None):
def set_progress_bar_config(self, **kwargs):
self._progress_bar_config = kwargs
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
r"""
Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). When this
option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed
@@ -1984,9 +2019,7 @@ def disable_xformers_memory_efficient_attention(self):
"""
self.set_use_memory_efficient_attention_xformers(False)
- def set_use_memory_efficient_attention_xformers(
- self, valid: bool, attention_op: Optional[Callable] = None
- ) -> None:
+ def set_use_memory_efficient_attention_xformers(self, valid: bool, attention_op: Callable | None = None) -> None:
# Recursively walk through all the children.
# Any children which exposes the set_use_memory_efficient_attention_xformers method
# gets the message
@@ -2004,7 +2037,7 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
for module in modules:
fn_recursive_set_mem_eff(module)
- def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):
+ def enable_attention_slicing(self, slice_size: str | int = "auto"):
r"""
Enable sliced attention computation. When this option is enabled, the attention module splits the input tensor
in slices to compute attention in several steps. For more than one attention head, the computation is performed
@@ -2049,7 +2082,7 @@ def disable_attention_slicing(self):
# set slice_size = `None` to disable `attention slicing`
self.enable_attention_slicing(None)
- def set_attention_slice(self, slice_size: Optional[int]):
+ def set_attention_slice(self, slice_size: int | None):
module_names, _ = self._get_signature_keys(self)
modules = [getattr(self, n, None) for n in module_names]
modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attention_slice")]
@@ -2183,7 +2216,7 @@ def from_pipe(cls, pipeline, **kwargs):
return new_pipeline
def _maybe_raise_error_if_group_offload_active(
- self, raise_error: bool = False, module: Optional[torch.nn.Module] = None
+ self, raise_error: bool = False, module: torch.nn.Module | None = None
) -> bool:
from ..hooks.group_offloading import _is_group_offload_enabled
@@ -2200,6 +2233,21 @@ def _maybe_raise_error_if_group_offload_active(
return True
return False
+ def _is_pipeline_device_mapped(self):
+ # We support passing `device_map="cuda"`, for example. This is helpful, in case
+ # users want to pass `device_map="cpu"` when initializing a pipeline. This explicit declaration is desirable
+ # in limited VRAM environments because quantized models often initialize directly on the accelerator.
+ device_map = self.hf_device_map
+ is_device_type_map = False
+ if isinstance(device_map, str):
+ try:
+ torch.device(device_map)
+ is_device_type_map = True
+ except RuntimeError:
+ pass
+
+ return not is_device_type_map and isinstance(device_map, dict) and len(device_map) > 1
+
class StableDiffusionMixin:
r"""
diff --git a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
index 2ecc13ef71bf..604e51d88583 100644
--- a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
+++ b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
@@ -16,7 +16,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import T5EncoderModel, T5Tokenizer
@@ -181,10 +181,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -199,15 +199,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -302,15 +302,15 @@ def __init__(
# Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 120,
**kwargs,
@@ -319,9 +319,9 @@ def encode_prompt(
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -687,50 +687,50 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 20,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 4.5,
- num_images_per_prompt: Optional[int] = 1,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ num_images_per_prompt: int | None = 1,
+ height: int | None = None,
+ width: int | None = None,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
clean_caption: bool = True,
use_resolution_binning: bool = True,
max_sequence_length: int = 120,
**kwargs,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -749,7 +749,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py
index f53d3c5630f0..286695aa8eb9 100644
--- a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py
+++ b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py
@@ -16,7 +16,7 @@
import inspect
import re
import urllib.parse as ul
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable
import torch
from transformers import T5EncoderModel, T5Tokenizer
@@ -125,10 +125,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -143,15 +143,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -246,15 +246,15 @@ def __init__(
# Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->300
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
**kwargs,
@@ -263,9 +263,9 @@ def encode_prompt(
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -632,50 +632,50 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 20,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 4.5,
- num_images_per_prompt: Optional[int] = 1,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ num_images_per_prompt: int | None = 1,
+ height: int | None = None,
+ width: int | None = None,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
clean_caption: bool = True,
use_resolution_binning: bool = True,
max_sequence_length: int = 300,
**kwargs,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -694,7 +694,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/prx/__init__.py b/src/diffusers/pipelines/prx/__init__.py
index 87aaefbd1368..ad2948e92e04 100644
--- a/src/diffusers/pipelines/prx/__init__.py
+++ b/src/diffusers/pipelines/prx/__init__.py
@@ -24,14 +24,25 @@
else:
_import_structure["pipeline_prx"] = ["PRXPipeline"]
-# Import T5GemmaEncoder for pipeline loading compatibility
+# Wrap T5GemmaEncoder to pass config.encoder (T5GemmaModuleConfig) instead of the
+# composite T5GemmaConfig, which lacks flat attributes expected by T5GemmaEncoder.__init__.
try:
if is_transformers_available():
import transformers
- from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder
+ from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder as _T5GemmaEncoder
+
+ class T5GemmaEncoder(_T5GemmaEncoder):
+ @classmethod
+ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+ if "config" not in kwargs:
+ from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig
+
+ config = T5GemmaConfig.from_pretrained(pretrained_model_name_or_path)
+ if hasattr(config, "encoder"):
+ kwargs["config"] = config.encoder
+ return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
_additional_imports["T5GemmaEncoder"] = T5GemmaEncoder
- # Patch transformers module directly for serialization
if not hasattr(transformers, "T5GemmaEncoder"):
transformers.T5GemmaEncoder = T5GemmaEncoder
except ImportError:
diff --git a/src/diffusers/pipelines/prx/pipeline_output.py b/src/diffusers/pipelines/prx/pipeline_output.py
index ea1bc9bf418a..81f92c294735 100644
--- a/src/diffusers/pipelines/prx/pipeline_output.py
+++ b/src/diffusers/pipelines/prx/pipeline_output.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -27,9 +26,9 @@ class PRXPipelineOutput(BaseOutput):
Output class for PRX pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/prx/pipeline_prx.py b/src/diffusers/pipelines/prx/pipeline_prx.py
index 873f25316e6d..e14815b91c41 100644
--- a/src/diffusers/pipelines/prx/pipeline_prx.py
+++ b/src/diffusers/pipelines/prx/pipeline_prx.py
@@ -16,9 +16,8 @@
import inspect
import re
import urllib.parse as ul
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
-import ftfy
import torch
from transformers import (
AutoTokenizer,
@@ -34,13 +33,13 @@
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.prx.pipeline_output import PRXPipelineOutput
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import (
- logging,
- replace_example_docstring,
-)
+from diffusers.utils import is_ftfy_available, logging, replace_example_docstring
from diffusers.utils.torch_utils import randn_tensor
+if is_ftfy_available():
+ import ftfy
+
DEFAULT_RESOLUTION = 512
ASPECT_RATIO_256_BIN = {
@@ -284,9 +283,9 @@ def __init__(
transformer: PRXTransformer2DModel,
scheduler: FlowMatchEulerDiscreteScheduler,
text_encoder: T5GemmaEncoder,
- tokenizer: Union[T5TokenizerFast, GemmaTokenizerFast, AutoTokenizer],
- vae: Optional[Union[AutoencoderKL, AutoencoderDC]] = None,
- default_sample_size: Optional[int] = DEFAULT_RESOLUTION,
+ tokenizer: T5TokenizerFast | GemmaTokenizerFast | AutoTokenizer,
+ vae: AutoencoderKL | AutoencoderDC | None = None,
+ default_sample_size: int | None = DEFAULT_RESOLUTION,
):
super().__init__()
@@ -352,8 +351,8 @@ def prepare_latents(
width: int,
dtype: torch.dtype,
device: torch.device,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
):
"""Prepare initial latents for the diffusion process."""
if latents is None:
@@ -370,15 +369,15 @@ def prepare_latents(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- prompt_attention_mask: Optional[torch.BoolTensor] = None,
- negative_prompt_attention_mask: Optional[torch.BoolTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ prompt_attention_mask: torch.BoolTensor | None = None,
+ negative_prompt_attention_mask: torch.BoolTensor | None = None,
):
"""Encode text prompt using standard text encoder and tokenizer, or use precomputed embeddings."""
if device is None:
@@ -420,7 +419,7 @@ def encode_prompt(
negative_prompt_attention_mask if do_classifier_free_guidance else None,
)
- def _tokenize_prompts(self, prompts: List[str], device: torch.device):
+ def _tokenize_prompts(self, prompts: list[str], device: torch.device):
"""Tokenize and clean prompts."""
cleaned = [self.text_preprocessor.clean_text(text) for text in prompts]
tokens = self.tokenizer(
@@ -435,7 +434,7 @@ def _tokenize_prompts(self, prompts: List[str], device: torch.device):
def _encode_prompt_standard(
self,
- prompt: List[str],
+ prompt: list[str],
device: torch.device,
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
@@ -473,13 +472,13 @@ def _encode_prompt_standard(
def check_inputs(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
height: int,
width: int,
guidance_scale: float,
- callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ callback_on_step_end_tensor_inputs: list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
):
"""Check that all inputs are in correct format."""
if prompt is not None and prompt_embeds is not None:
@@ -527,31 +526,31 @@ def check_inputs(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 4.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- prompt_attention_mask: Optional[torch.BoolTensor] = None,
- negative_prompt_attention_mask: Optional[torch.BoolTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ prompt_attention_mask: torch.BoolTensor | None = None,
+ negative_prompt_attention_mask: torch.BoolTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
use_resolution_binning: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`
instead.
negative_prompt (`str`, *optional*, defaults to `""`):
@@ -564,7 +563,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 28):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -576,7 +575,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -610,7 +609,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self, step, timestep, callback_kwargs)`.
`callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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 tensors that are listed
in the `._callback_tensor_inputs` attribute.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_output.py b/src/diffusers/pipelines/qwenimage/pipeline_output.py
index eef4b60e3770..e4ed06856e6a 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_output.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class QwenImagePipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py
index 9938ae95fd55..1715aa4d4250 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -72,10 +72,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -90,15 +90,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -187,9 +187,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -225,17 +225,17 @@ def _get_qwen_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -447,36 +447,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
true_cfg_scale: float = 4.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
@@ -494,7 +494,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -510,7 +510,7 @@ def __call__(
enable classifier-free guidance computations).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -538,7 +538,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py
index e0cb9924cda3..85a936f9ec24 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -117,7 +117,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -132,10 +132,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -150,15 +150,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -218,7 +218,7 @@ def __init__(
text_encoder: Qwen2_5_VLForConditionalGeneration,
tokenizer: Qwen2Tokenizer,
transformer: QwenImageTransformer2DModel,
- controlnet: Union[QwenImageControlNetModel, QwenImageMultiControlNetModel],
+ controlnet: QwenImageControlNetModel | QwenImageMultiControlNetModel,
):
super().__init__()
@@ -251,9 +251,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
# Coped from diffusers.pipelines.qwenimage.pipeline_qwenimage.get_qwen_prompt_embeds
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -290,17 +290,17 @@ def _get_qwen_prompt_embeds(
# Coped from diffusers.pipelines.qwenimage.pipeline_qwenimage.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -548,40 +548,40 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
true_cfg_scale: float = 4.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = None,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = None,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
control_image: PipelineImageInput = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
@@ -599,7 +599,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -615,7 +615,7 @@ def __call__(
enable classifier-free guidance computations).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -643,7 +643,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py
index f8521318e630..b1da59cb4f6c 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -90,7 +90,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -105,10 +105,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -123,15 +123,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -233,9 +233,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
# Coped from diffusers.pipelines.qwenimage.pipeline_qwenimage.get_qwen_prompt_embeds
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -272,16 +272,16 @@ def _get_qwen_prompt_embeds(
# Coped from diffusers.pipelines.qwenimage.pipeline_qwenimage.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -595,41 +595,41 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
true_cfg_scale: float = 4.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 1.0,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
+ control_guidance_start: float | list[float] = 0.0,
+ control_guidance_end: float | list[float] = 1.0,
control_image: PipelineImageInput = None,
control_mask: PipelineImageInput = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
@@ -642,7 +642,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -654,7 +654,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -682,7 +682,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py
index 353aadcbf08a..15e72a010ce5 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -81,10 +81,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -99,15 +99,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -140,7 +140,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -225,10 +225,10 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- image: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -272,18 +272,18 @@ def _get_qwen_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- image: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ image: torch.Tensor | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
image (`torch.Tensor`, *optional*):
image to be encoded
@@ -550,43 +550,43 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[PipelineImageInput] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ image: PipelineImageInput | None = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
true_cfg_scale: float = 4.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
@@ -604,7 +604,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -620,7 +620,7 @@ def __call__(
enable classifier-free guidance computations).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -648,7 +648,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py
index bc397f357bf5..20a2748bc7f9 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -82,10 +82,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -100,15 +100,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -141,7 +141,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -236,10 +236,10 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._get_qwen_prompt_embeds
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- image: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -284,18 +284,18 @@ def _get_qwen_prompt_embeds(
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- image: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ image: torch.Tensor | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
image (`torch.Tensor`, *optional*):
image to be encoded
@@ -675,47 +675,47 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[PipelineImageInput] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ image: PipelineImageInput | None = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
mask_image: PipelineImageInput = None,
masked_image_latents: PipelineImageInput = None,
true_cfg_scale: float = 4.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
@@ -726,14 +726,14 @@ def __call__(
enabled by setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale
encourages to generate images that are closely linked to the text `prompt`, usually at the expense of
lower image quality.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
1)`, or `(H, W)`.
- mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`):
+ mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`):
`Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask
latents tensor will ge generated by `mask_image`.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -756,7 +756,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -772,7 +772,7 @@ def __call__(
enable classifier-free guidance computations).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -800,7 +800,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py
index bc688aeee319..588783458571 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -84,10 +84,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -102,15 +102,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -143,7 +143,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -228,10 +228,10 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- image: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -286,18 +286,18 @@ def _get_qwen_prompt_embeds(
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- image: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ image: torch.Tensor | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
image (`torch.Tensor`, *optional*):
image to be encoded
@@ -519,43 +519,43 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[PipelineImageInput] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ image: PipelineImageInput | None = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
true_cfg_scale: float = 4.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
@@ -573,7 +573,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -589,7 +589,7 @@ def __call__(
enable classifier-free guidance computations).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -617,7 +617,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py
index 522e1d203051..42e63f8919a2 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py
@@ -1,5 +1,5 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -45,7 +45,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -74,10 +74,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -92,15 +92,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -194,9 +194,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._get_qwen_prompt_embeds
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -268,17 +268,17 @@ def get_timesteps(self, num_inference_steps, strength, device):
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -521,42 +521,42 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
true_cfg_scale: float = 4.0,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -582,7 +582,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -598,7 +598,7 @@ def __call__(
enable classifier-free guidance computations).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -626,7 +626,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py
index 960196bec166..5baf5bf5f77d 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py
@@ -1,5 +1,5 @@
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -48,7 +48,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -77,10 +77,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -95,15 +95,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -204,9 +204,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._get_qwen_prompt_embeds
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -279,17 +279,17 @@ def get_timesteps(self, num_inference_steps, strength, device):
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -632,45 +632,45 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
true_cfg_scale: float = 4.0,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -683,14 +683,14 @@ def __call__(
setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale encourages to
generate images that are closely linked to the text `prompt`, usually at the expense of lower image
quality.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
1)`, or `(H, W)`.
- mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`):
+ mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`):
`Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask
latents tensor will be generated by `mask_image`.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -713,7 +713,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -729,7 +729,7 @@ def __call__(
enable classifier-free guidance computations).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -757,7 +757,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py
index ea455721e4e5..c7a44d880f9b 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -89,10 +89,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -107,15 +107,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -148,7 +148,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -250,9 +250,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor
def _get_qwen_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ prompt: str | list[str] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -291,17 +291,17 @@ def _get_qwen_prompt_embeds(
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -532,26 +532,26 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[PipelineImageInput] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ image: PipelineImageInput | None = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
true_cfg_scale: float = 4.0,
- layers: Optional[int] = 4,
+ layers: int | None = 4,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
- guidance_scale: Optional[float] = None,
+ sigmas: list[float] | None = None,
+ guidance_scale: float | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
resolution: int = 640,
cfg_normalize: bool = False,
@@ -561,16 +561,16 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
not greater than `1`).
@@ -584,7 +584,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -600,7 +600,7 @@ def __call__(
enable classifier-free guidance computations).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/sana/pipeline_output.py b/src/diffusers/pipelines/sana/pipeline_output.py
index f8ac12951644..b9d095906e06 100644
--- a/src/diffusers/pipelines/sana/pipeline_output.py
+++ b/src/diffusers/pipelines/sana/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class SanaPipelineOutput(BaseOutput):
Output class for Sana pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/sana/pipeline_sana.py b/src/diffusers/pipelines/sana/pipeline_sana.py
index 33f9de7d20f0..17e0be9ba7aa 100644
--- a/src/diffusers/pipelines/sana/pipeline_sana.py
+++ b/src/diffusers/pipelines/sana/pipeline_sana.py
@@ -17,7 +17,7 @@
import re
import urllib.parse as ul
import warnings
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
@@ -130,10 +130,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -148,15 +148,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -201,7 +201,7 @@ class SanaPipeline(DiffusionPipeline, SanaLoraLoaderMixin):
def __init__(
self,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
text_encoder: Gemma2PreTrainedModel,
vae: AutoencoderDC,
transformer: SanaTransformer2DModel,
@@ -275,18 +275,18 @@ def disable_vae_tiling(self):
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
dtype: torch.dtype,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
+ complex_human_instruction: list[str] | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
@@ -333,27 +333,27 @@ def _get_gemma_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
- lora_scale: Optional[float] = None,
+ complex_human_instruction: list[str] | None = None,
+ lora_scale: float | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -728,31 +728,31 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 20,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 4.5,
- num_images_per_prompt: Optional[int] = 1,
+ num_images_per_prompt: int | None = 1,
height: int = 1024,
width: int = 1024,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
clean_caption: bool = False,
use_resolution_binning: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 300,
- complex_human_instruction: List[str] = [
+ complex_human_instruction: list[str] = [
"Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:",
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
@@ -762,26 +762,26 @@ def __call__(
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
"User Prompt: ",
],
- ) -> Union[SanaPipelineOutput, Tuple]:
+ ) -> SanaPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 20):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -800,7 +800,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -838,13 +838,13 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to `300`):
Maximum sequence length to use with the `prompt`.
- complex_human_instruction (`List[str]`, *optional*):
+ complex_human_instruction (`list[str]`, *optional*):
Instructions for complex human attention:
https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.
diff --git a/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py b/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py
index 9d5e17c2ed48..d976c7035d9d 100644
--- a/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py
+++ b/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py
@@ -17,7 +17,7 @@
import re
import urllib.parse as ul
import warnings
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
@@ -137,10 +137,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -155,15 +155,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -208,7 +208,7 @@ class SanaControlNetPipeline(DiffusionPipeline, SanaLoraLoaderMixin):
def __init__(
self,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
text_encoder: Gemma2PreTrainedModel,
vae: AutoencoderDC,
transformer: SanaTransformer2DModel,
@@ -289,18 +289,18 @@ def disable_vae_tiling(self):
# Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
dtype: torch.dtype,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
+ complex_human_instruction: list[str] | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
@@ -348,27 +348,27 @@ def _get_gemma_prompt_embeds(
# Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
- lora_scale: Optional[float] = None,
+ complex_human_instruction: list[str] | None = None,
+ lora_scale: float | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -777,33 +777,33 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 20,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 4.5,
control_image: PipelineImageInput = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
- num_images_per_prompt: Optional[int] = 1,
+ controlnet_conditioning_scale: float | list[float] = 1.0,
+ num_images_per_prompt: int | None = 1,
height: int = 1024,
width: int = 1024,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
clean_caption: bool = False,
use_resolution_binning: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 300,
- complex_human_instruction: List[str] = [
+ complex_human_instruction: list[str] = [
"Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:",
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
@@ -813,26 +813,26 @@ def __call__(
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
"User Prompt: ",
],
- ) -> Union[SanaPipelineOutput, Tuple]:
+ ) -> SanaPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 20):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -842,15 +842,15 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,:
+ `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`):
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,
images must be passed as a list such that each element of the list can be correctly batched for input
to a single ControlNet.
- controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
the corresponding scale as a list.
@@ -863,7 +863,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -901,13 +901,13 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to `300`):
Maximum sequence length to use with the `prompt`.
- complex_human_instruction (`List[str]`, *optional*):
+ complex_human_instruction (`list[str]`, *optional*):
Instructions for complex human attention:
https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.
diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py
index 4c6d2247495d..c85f05275f51 100644
--- a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py
+++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py
@@ -17,7 +17,7 @@
import re
import urllib.parse as ul
import warnings
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
@@ -81,10 +81,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -99,15 +99,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -152,7 +152,7 @@ class SanaSprintPipeline(DiffusionPipeline, SanaLoraLoaderMixin):
def __init__(
self,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
text_encoder: Gemma2PreTrainedModel,
vae: AutoencoderDC,
transformer: SanaTransformer2DModel,
@@ -227,18 +227,18 @@ def disable_vae_tiling(self):
# Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
dtype: torch.dtype,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
+ complex_human_instruction: list[str] | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
@@ -285,21 +285,21 @@ def _get_gemma_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
- lora_scale: Optional[float] = None,
+ complex_human_instruction: list[str] | None = None,
+ lora_scale: float | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -616,29 +616,29 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_inference_steps: int = 2,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
max_timesteps: float = 1.57080,
intermediate_timesteps: float = 1.3,
guidance_scale: float = 4.5,
- num_images_per_prompt: Optional[int] = 1,
+ num_images_per_prompt: int | None = 1,
height: int = 1024,
width: int = 1024,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
clean_caption: bool = False,
use_resolution_binning: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 300,
- complex_human_instruction: List[str] = [
+ complex_human_instruction: list[str] = [
"Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:",
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
@@ -648,12 +648,12 @@ def __call__(
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
"User Prompt: ",
],
- ) -> Union[SanaPipelineOutput, Tuple]:
+ ) -> SanaPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
num_inference_steps (`int`, *optional*, defaults to 20):
@@ -663,7 +663,7 @@ def __call__(
The maximum timestep value used in the SCM scheduler.
intermediate_timesteps (`float`, *optional*, defaults to 1.3):
The intermediate timestep value used in SCM scheduler (only used when num_inference_steps=2).
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -682,7 +682,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -715,13 +715,13 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to `300`):
Maximum sequence length to use with the `prompt`.
- complex_human_instruction (`List[str]`, *optional*):
+ complex_human_instruction (`list[str]`, *optional*):
Instructions for complex human attention:
https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.
diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py
index 8899ed84c4e5..a17c494e88eb 100644
--- a/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py
+++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py
@@ -17,7 +17,7 @@
import re
import urllib.parse as ul
import warnings
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
import torch.nn.functional as F
@@ -87,10 +87,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -105,15 +105,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -159,7 +159,7 @@ class SanaSprintImg2ImgPipeline(DiffusionPipeline, SanaLoraLoaderMixin):
def __init__(
self,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
text_encoder: Gemma2PreTrainedModel,
vae: AutoencoderDC,
transformer: SanaTransformer2DModel,
@@ -237,18 +237,18 @@ def disable_vae_tiling(self):
# Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
dtype: torch.dtype,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
+ complex_human_instruction: list[str] | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
@@ -296,21 +296,21 @@ def _get_gemma_prompt_embeds(
# Copied from diffusers.pipelines.sana.pipeline_sana_sprint.SanaSprintPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
- lora_scale: Optional[float] = None,
+ complex_human_instruction: list[str] | None = None,
+ lora_scale: float | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -687,31 +687,31 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_inference_steps: int = 2,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
max_timesteps: float = 1.57080,
intermediate_timesteps: float = 1.3,
guidance_scale: float = 4.5,
image: PipelineImageInput = None,
strength: float = 0.6,
- num_images_per_prompt: Optional[int] = 1,
+ num_images_per_prompt: int | None = 1,
height: int = 1024,
width: int = 1024,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
clean_caption: bool = False,
use_resolution_binning: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 300,
- complex_human_instruction: List[str] = [
+ complex_human_instruction: list[str] = [
"Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:",
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
@@ -721,12 +721,12 @@ def __call__(
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
"User Prompt: ",
],
- ) -> Union[SanaPipelineOutput, Tuple]:
+ ) -> SanaPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
num_inference_steps (`int`, *optional*, defaults to 20):
@@ -736,7 +736,7 @@ def __call__(
The maximum timestep value used in the SCM scheduler.
intermediate_timesteps (`float`, *optional*, defaults to 1.3):
The intermediate timestep value used in SCM scheduler (only used when num_inference_steps=2).
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
@@ -755,7 +755,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
[`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -788,13 +788,13 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to `300`):
Maximum sequence length to use with the `prompt`.
- complex_human_instruction (`List[str]`, *optional*):
+ complex_human_instruction (`list[str]`, *optional*):
Instructions for complex human attention:
https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.
diff --git a/src/diffusers/pipelines/sana_video/pipeline_output.py b/src/diffusers/pipelines/sana_video/pipeline_output.py
index 4d37923889eb..a625ad3e1d34 100644
--- a/src/diffusers/pipelines/sana_video/pipeline_output.py
+++ b/src/diffusers/pipelines/sana_video/pipeline_output.py
@@ -11,7 +11,7 @@ class SanaVideoPipelineOutput(BaseOutput):
Output class for Sana-Video pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
diff --git a/src/diffusers/pipelines/sana_video/pipeline_sana_video.py b/src/diffusers/pipelines/sana_video/pipeline_sana_video.py
index a786275e45a9..8b44dfc1143c 100644
--- a/src/diffusers/pipelines/sana_video/pipeline_sana_video.py
+++ b/src/diffusers/pipelines/sana_video/pipeline_sana_video.py
@@ -17,7 +17,7 @@
import re
import urllib.parse as ul
import warnings
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
@@ -126,10 +126,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -144,15 +144,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -211,9 +211,9 @@ class SanaVideoPipeline(DiffusionPipeline, SanaLoraLoaderMixin):
def __init__(
self,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
text_encoder: Gemma2PreTrainedModel,
- vae: Union[AutoencoderDC, AutoencoderKLWan],
+ vae: AutoencoderDC | AutoencoderKLWan,
transformer: SanaVideoTransformer3DModel,
scheduler: DPMSolverMultistepScheduler,
):
@@ -233,18 +233,18 @@ def __init__(
# Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
dtype: torch.dtype,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
+ complex_human_instruction: list[str] | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
@@ -291,27 +291,27 @@ def _get_gemma_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
- lora_scale: Optional[float] = None,
+ complex_human_instruction: list[str] | None = None,
+ lora_scale: float | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -650,10 +650,10 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -702,32 +702,32 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 6.0,
- num_videos_per_prompt: Optional[int] = 1,
+ num_videos_per_prompt: int | None = 1,
height: int = 480,
width: int = 832,
frames: int = 81,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
clean_caption: bool = False,
use_resolution_binning: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 300,
- complex_human_instruction: List[str] = [
+ complex_human_instruction: list[str] = [
"Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for video generation. Evaluate the level of detail in the user prompt:",
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and temporal relationships to create vivid and dynamic scenes.",
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
@@ -737,26 +737,26 @@ def __call__(
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
"User Prompt: ",
],
- ) -> Union[SanaVideoPipelineOutput, Tuple]:
+ ) -> SanaVideoPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the video generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality video at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -777,7 +777,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -820,7 +820,7 @@ def __call__(
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to `300`):
Maximum sequence length to use with the `prompt`.
- complex_human_instruction (`List[str]`, *optional*):
+ complex_human_instruction (`list[str]`, *optional*):
Instructions for complex human attention:
https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.
diff --git a/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py b/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py
index e87880b64cee..b90d7c6f5a60 100644
--- a/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py
+++ b/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py
@@ -17,7 +17,7 @@
import re
import urllib.parse as ul
import warnings
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL
import torch
@@ -102,10 +102,10 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -120,15 +120,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -161,7 +161,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -201,9 +201,9 @@ class SanaImageToVideoPipeline(DiffusionPipeline, SanaLoraLoaderMixin):
def __init__(
self,
- tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+ tokenizer: GemmaTokenizer | GemmaTokenizerFast,
text_encoder: Gemma2PreTrainedModel,
- vae: Union[AutoencoderDC, AutoencoderKLWan],
+ vae: AutoencoderDC | AutoencoderKLWan,
transformer: SanaVideoTransformer3DModel,
scheduler: FlowMatchEulerDiscreteScheduler,
):
@@ -230,18 +230,18 @@ def __init__(
# Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds
def _get_gemma_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
device: torch.device,
dtype: torch.dtype,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
+ complex_human_instruction: list[str] | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`, *optional*):
torch device to place the resulting embeddings on
@@ -289,27 +289,27 @@ def _get_gemma_prompt_embeds(
# Copied from diffusers.pipelines.sana_video.pipeline_sana_video.SanaVideoPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_videos_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ device: torch.device | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
clean_caption: bool = False,
max_sequence_length: int = 300,
- complex_human_instruction: Optional[List[str]] = None,
- lora_scale: Optional[float] = None,
+ complex_human_instruction: list[str] | None = None,
+ lora_scale: float | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds`
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
PixArt-Alpha, this should be "".
@@ -653,10 +653,10 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
shape = (
@@ -726,32 +726,32 @@ def interrupt(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
negative_prompt: str = "",
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 6.0,
- num_videos_per_prompt: Optional[int] = 1,
+ num_videos_per_prompt: int | None = 1,
height: int = 480,
width: int = 832,
frames: int = 81,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_attention_mask: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_attention_mask: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_attention_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_attention_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
clean_caption: bool = False,
use_resolution_binning: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 300,
- complex_human_instruction: List[str] = [
+ complex_human_instruction: list[str] = [
"Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for video generation. Evaluate the level of detail in the user prompt:",
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and temporal relationships to create vivid and dynamic scenes.",
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
@@ -761,7 +761,7 @@ def __call__(
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
"User Prompt: ",
],
- ) -> Union[SanaVideoPipelineOutput, Tuple]:
+ ) -> SanaVideoPipelineOutput | tuple:
"""
Function invoked when calling the pipeline for generation.
@@ -769,21 +769,21 @@ def __call__(
image (`PipelineImageInput`):
The input image to condition the video generation on. The first frame of the generated video will be
conditioned on this image.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the video generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality video at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -804,7 +804,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -847,7 +847,7 @@ def __call__(
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to `300`):
Maximum sequence length to use with the `prompt`.
- complex_human_instruction (`List[str]`, *optional*):
+ complex_human_instruction (`list[str]`, *optional*):
Instructions for complex human attention:
https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.
diff --git a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py
index 349912993981..8e5429ce2a8d 100644
--- a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py
+++ b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -13,13 +12,13 @@ class SemanticStableDiffusionPipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
- nsfw_content_detected (`List[bool]`)
- List indicating whether the corresponding generated image contains “not-safe-for-work” (nsfw) content or
+ nsfw_content_detected (`list[bool]`)
+ list indicating whether the corresponding generated image contains “not-safe-for-work” (nsfw) content or
`None` if safety checking could not be performed.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_content_detected: Optional[List[bool]]
+ images: list[PIL.Image.Image] | np.ndarray
+ nsfw_content_detected: list[bool] | None
diff --git a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py
index 49b09e205cc5..05d28896e117 100644
--- a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py
+++ b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py
@@ -1,6 +1,6 @@
import inspect
from itertools import repeat
-from typing import Callable, List, Optional, Union
+from typing import Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
@@ -144,7 +144,6 @@ def prepare_extra_step_kwargs(self, generator, eta):
extra_step_kwargs["generator"] = generator
return extra_step_kwargs
- # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs
def check_inputs(
self,
prompt,
@@ -223,37 +222,37 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str],
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
num_images_per_prompt: int = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- editing_prompt: Optional[Union[str, List[str]]] = None,
- editing_prompt_embeddings: Optional[torch.Tensor] = None,
- reverse_editing_direction: Optional[Union[bool, List[bool]]] = False,
- edit_guidance_scale: Optional[Union[float, List[float]]] = 5,
- edit_warmup_steps: Optional[Union[int, List[int]]] = 10,
- edit_cooldown_steps: Optional[Union[int, List[int]]] = None,
- edit_threshold: Optional[Union[float, List[float]]] = 0.9,
- edit_momentum_scale: Optional[float] = 0.1,
- edit_mom_beta: Optional[float] = 0.4,
- edit_weights: Optional[List[float]] = None,
- sem_guidance: Optional[List[torch.Tensor]] = None,
+ editing_prompt: str | list[str] | None = None,
+ editing_prompt_embeddings: torch.Tensor | None = None,
+ reverse_editing_direction: bool | list[bool] | None = False,
+ edit_guidance_scale: float | list[float] | None = 5,
+ edit_warmup_steps: int | list[int] | None = 10,
+ edit_cooldown_steps: int | list[int] | None = None,
+ edit_threshold: float | list[float] | None = 0.9,
+ edit_momentum_scale: float | None = 0.1,
+ edit_mom_beta: float | None = 0.4,
+ edit_weights: list[float] | None = None,
+ sem_guidance: list[torch.Tensor] | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -265,7 +264,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -273,7 +272,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -291,24 +290,24 @@ def __call__(
callback_steps (`int`, *optional*, defaults to 1):
The frequency at which the `callback` function is called. If not specified, the callback is called at
every step.
- editing_prompt (`str` or `List[str]`, *optional*):
+ editing_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to use for semantic guidance. Semantic guidance is disabled by setting
`editing_prompt = None`. Guidance direction of prompt should be specified via
`reverse_editing_direction`.
editing_prompt_embeddings (`torch.Tensor`, *optional*):
Pre-computed embeddings to use for semantic guidance. Guidance direction of embedding should be
specified via `reverse_editing_direction`.
- reverse_editing_direction (`bool` or `List[bool]`, *optional*, defaults to `False`):
+ reverse_editing_direction (`bool` or `list[bool]`, *optional*, defaults to `False`):
Whether the corresponding prompt in `editing_prompt` should be increased or decreased.
- edit_guidance_scale (`float` or `List[float]`, *optional*, defaults to 5):
+ edit_guidance_scale (`float` or `list[float]`, *optional*, defaults to 5):
Guidance scale for semantic guidance. If provided as a list, values should correspond to
`editing_prompt`.
- edit_warmup_steps (`float` or `List[float]`, *optional*, defaults to 10):
+ edit_warmup_steps (`float` or `list[float]`, *optional*, defaults to 10):
Number of diffusion steps (for each prompt) for which semantic guidance is not applied. Momentum is
calculated for those steps and applied once all warmup periods are over.
- edit_cooldown_steps (`float` or `List[float]`, *optional*, defaults to `None`):
+ edit_cooldown_steps (`float` or `list[float]`, *optional*, defaults to `None`):
Number of diffusion steps (for each prompt) after which semantic guidance is longer applied.
- edit_threshold (`float` or `List[float]`, *optional*, defaults to 0.9):
+ edit_threshold (`float` or `list[float]`, *optional*, defaults to 0.9):
Threshold of semantic guidance.
edit_momentum_scale (`float`, *optional*, defaults to 0.1):
Scale of the momentum to be added to the semantic guidance at each diffusion step. If set to 0.0,
@@ -318,11 +317,11 @@ def __call__(
Defines how semantic guidance momentum builds up. `edit_mom_beta` indicates how much of the previous
momentum is kept. Momentum is already built up during warmup (for diffusion steps smaller than
`edit_warmup_steps`).
- edit_weights (`List[float]`, *optional*, defaults to `None`):
+ edit_weights (`list[float]`, *optional*, defaults to `None`):
Indicates how much each individual concept should influence the overall guidance. If no weights are
provided all concepts are applied equally.
- sem_guidance (`List[torch.Tensor]`, *optional*):
- List of pre-generated guidance vectors to be applied at generation. Length of the list has to
+ sem_guidance (`list[torch.Tensor]`, *optional*):
+ list of pre-generated guidance vectors to be applied at generation. Length of the list has to
correspond to `num_inference_steps`.
Examples:
@@ -458,7 +457,7 @@ def __call__(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
diff --git a/src/diffusers/pipelines/shap_e/camera.py b/src/diffusers/pipelines/shap_e/camera.py
index 31e1759d6154..81807b6ff35a 100644
--- a/src/diffusers/pipelines/shap_e/camera.py
+++ b/src/diffusers/pipelines/shap_e/camera.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Tuple
import numpy as np
import torch
@@ -33,7 +32,7 @@ class DifferentiableProjectiveCamera:
height: int
x_fov: float
y_fov: float
- shape: Tuple[int]
+ shape: tuple[int]
def __post_init__(self):
assert self.x.shape[0] == self.y.shape[0] == self.z.shape[0] == self.origin.shape[0]
diff --git a/src/diffusers/pipelines/shap_e/pipeline_shap_e.py b/src/diffusers/pipelines/shap_e/pipeline_shap_e.py
index 49ddfd1196bf..44967dfb3349 100644
--- a/src/diffusers/pipelines/shap_e/pipeline_shap_e.py
+++ b/src/diffusers/pipelines/shap_e/pipeline_shap_e.py
@@ -14,7 +14,6 @@
import math
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -82,7 +81,7 @@ class ShapEPipelineOutput(BaseOutput):
A list of images for 3D rendering.
"""
- images: Union[List[List[PIL.Image.Image]], List[List[np.ndarray]]]
+ images: list[list[PIL.Image.Image]] | list[list[np.ndarray]]
class ShapEPipeline(DiffusionPipeline):
@@ -195,25 +194,25 @@ def __call__(
prompt: str,
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
guidance_scale: float = 4.0,
frame_size: int = 64,
- output_type: Optional[str] = "pil", # pil, np, latent, mesh
+ output_type: str | None = "pil", # pil, np, latent, mesh
return_dict: bool = True,
):
"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
num_inference_steps (`int`, *optional*, defaults to 25):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py b/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py
index 55d8b85822c4..964db30e7de2 100644
--- a/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py
+++ b/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -83,7 +82,7 @@ class ShapEPipelineOutput(BaseOutput):
A list of images for 3D rendering.
"""
- images: Union[PIL.Image.Image, np.ndarray]
+ images: PIL.Image.Image | np.ndarray
class ShapEImg2ImgPipeline(DiffusionPipeline):
@@ -147,7 +146,7 @@ def _encode_image(
num_images_per_prompt,
do_classifier_free_guidance,
):
- if isinstance(image, List) and isinstance(image[0], torch.Tensor):
+ if isinstance(image, list) and isinstance(image[0], torch.Tensor):
image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
if not isinstance(image, torch.Tensor):
@@ -174,21 +173,21 @@ def _encode_image(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Union[PIL.Image.Image, List[PIL.Image.Image]],
+ image: PIL.Image.Image | list[PIL.Image.Image],
num_images_per_prompt: int = 1,
num_inference_steps: int = 25,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
guidance_scale: float = 4.0,
frame_size: int = 64,
- output_type: Optional[str] = "pil", # pil, np, latent, mesh
+ output_type: str | None = "pil", # pil, np, latent, mesh
return_dict: bool = True,
):
"""
The call function to the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image` or tensor representing an image batch to be used as the starting point. Can also accept image
latents as image, but if passing latents directly it is not encoded again.
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -196,7 +195,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 25):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -231,7 +230,7 @@ def __call__(
batch_size = len(image)
else:
raise ValueError(
- f"`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `List[PIL.Image.Image]` or `List[torch.Tensor]` but is {type(image)}"
+ f"`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `list[PIL.Image.Image]` or `list[torch.Tensor]` but is {type(image)}"
)
device = self._execution_device
diff --git a/src/diffusers/pipelines/shap_e/renderer.py b/src/diffusers/pipelines/shap_e/renderer.py
index d1d05c894595..0c2058c887fc 100644
--- a/src/diffusers/pipelines/shap_e/renderer.py
+++ b/src/diffusers/pipelines/shap_e/renderer.py
@@ -14,7 +14,6 @@
import math
from dataclasses import dataclass
-from typing import Dict, Optional, Tuple
import numpy as np
import torch
@@ -131,7 +130,7 @@ def _convert_srgb_to_linear(u: torch.Tensor):
def _create_flat_edge_indices(
flat_cube_indices: torch.Tensor,
- grid_size: Tuple[int, int, int],
+ grid_size: tuple[int, int, int],
):
num_xs = (grid_size[0] - 1) * grid_size[1] * grid_size[2]
y_offset = num_xs
@@ -301,7 +300,7 @@ def intersect(
self,
origin: torch.Tensor,
direction: torch.Tensor,
- t0_lower: Optional[torch.Tensor] = None,
+ t0_lower: torch.Tensor | None = None,
epsilon=1e-6,
):
"""
@@ -479,7 +478,7 @@ class MeshDecoderOutput(BaseOutput):
verts: torch.Tensor
faces: torch.Tensor
- vertex_channels: Dict[str, torch.Tensor]
+ vertex_channels: dict[str, torch.Tensor]
class MeshDecoder(nn.Module):
@@ -742,13 +741,13 @@ class ShapEParamsProjModel(ModelMixin, ConfigMixin):
def __init__(
self,
*,
- param_names: Tuple[str, ...] = (
+ param_names: tuple[str] = (
"nerstf.mlp.0.weight",
"nerstf.mlp.1.weight",
"nerstf.mlp.2.weight",
"nerstf.mlp.3.weight",
),
- param_shapes: Tuple[Tuple[int]] = (
+ param_shapes: tuple[tuple[int]] = (
(256, 93),
(256, 256),
(256, 256),
@@ -786,13 +785,13 @@ class ShapERenderer(ModelMixin, ConfigMixin):
def __init__(
self,
*,
- param_names: Tuple[str, ...] = (
+ param_names: tuple[str] = (
"nerstf.mlp.0.weight",
"nerstf.mlp.1.weight",
"nerstf.mlp.2.weight",
"nerstf.mlp.3.weight",
),
- param_shapes: Tuple[Tuple[int, int], ...] = (
+ param_shapes: tuple[tuple[int]] = (
(256, 93),
(256, 256),
(256, 256),
@@ -804,7 +803,7 @@ def __init__(
n_hidden_layers: int = 6,
act_fn: str = "swish",
insert_direction_at: int = 4,
- background: Tuple[float, ...] = (
+ background: tuple[float] = (
255.0,
255.0,
255.0,
@@ -953,7 +952,7 @@ def decode_to_mesh(
device,
grid_size: int = 128,
query_batch_size: int = 4096,
- texture_channels: Tuple = ("R", "G", "B"),
+ texture_channels: tuple = ("R", "G", "B"),
):
# 1. project the parameters from the generated latents
projected_params = self.params_proj(latents)
diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_output.py b/src/diffusers/pipelines/skyreels_v2/pipeline_output.py
index 7a170d24c39a..dac2316362ec 100644
--- a/src/diffusers/pipelines/skyreels_v2/pipeline_output.py
+++ b/src/diffusers/pipelines/skyreels_v2/pipeline_output.py
@@ -11,8 +11,8 @@ class SkyReelsV2PipelineOutput(BaseOutput):
Output class for SkyReelsV2 pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py
index 1b1c8ee097c5..c92608fad3b6 100644
--- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py
+++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py
@@ -13,11 +13,11 @@
# limitations under the License.
import html
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import regex as re
import torch
-from transformers import AutoTokenizer, UMT5EncoderModel
+from transformers import AutoTokenizer, T5EncoderModel, UMT5EncoderModel
from ...callbacks import MultiPipelineCallbacks, PipelineCallback
from ...loaders import SkyReelsV2LoraLoaderMixin
@@ -132,7 +132,7 @@ class SkyReelsV2Pipeline(DiffusionPipeline, SkyReelsV2LoraLoaderMixin):
def __init__(
self,
tokenizer: AutoTokenizer,
- text_encoder: UMT5EncoderModel,
+ text_encoder: T5EncoderModel | UMT5EncoderModel,
transformer: SkyReelsV2Transformer3DModel,
vae: AutoencoderKLWan,
scheduler: UniPCMultistepScheduler,
@@ -154,11 +154,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -196,23 +196,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -324,10 +324,10 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -377,32 +377,30 @@ def attention_kwargs(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 544,
width: int = 960,
num_frames: int = 97,
num_inference_steps: int = 50,
guidance_scale: float = 6.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, defaults to `544`):
@@ -422,7 +420,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -445,7 +443,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py
index 4bc0d0aaea83..8751240a1af9 100644
--- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py
+++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py
@@ -16,11 +16,10 @@
import math
import re
from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
-import ftfy
import torch
-from transformers import AutoTokenizer, UMT5EncoderModel
+from transformers import AutoTokenizer, T5EncoderModel, UMT5EncoderModel
from ...callbacks import MultiPipelineCallbacks, PipelineCallback
from ...loaders import SkyReelsV2LoraLoaderMixin
@@ -114,7 +113,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -154,7 +153,7 @@ class SkyReelsV2DiffusionForcingPipeline(DiffusionPipeline, SkyReelsV2LoraLoader
def __init__(
self,
tokenizer: AutoTokenizer,
- text_encoder: UMT5EncoderModel,
+ text_encoder: T5EncoderModel | UMT5EncoderModel,
transformer: SkyReelsV2Transformer3DModel,
vae: AutoencoderKLWan,
scheduler: UniPCMultistepScheduler,
@@ -176,11 +175,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -218,23 +217,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -354,15 +353,15 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 97,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- base_latent_num_frames: Optional[int] = None,
- video_latents: Optional[torch.Tensor] = None,
- causal_block_size: Optional[int] = None,
- overlap_history_latent_frames: Optional[int] = None,
- long_video_iter: Optional[int] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ base_latent_num_frames: int | None = None,
+ video_latents: torch.Tensor | None = None,
+ causal_block_size: int | None = None,
+ overlap_history_latent_frames: int | None = None,
+ long_video_iter: int | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -462,7 +461,7 @@ def generate_timestep_matrix(
num_latent_frames]
- step_update_mask (torch.Tensor): Boolean mask indicating which frames to update Shape:
[num_iterations, num_latent_frames]
- - valid_interval (list[tuple]): List of (start, end) intervals for each iteration
+ - valid_interval (list[tuple]): list of (start, end) intervals for each iteration
Raises:
ValueError: If ar_step is too small for the given configuration
@@ -599,41 +598,39 @@ def attention_kwargs(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] = None,
height: int = 544,
width: int = 960,
num_frames: int = 97,
num_inference_steps: int = 50,
guidance_scale: float = 6.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
- overlap_history: Optional[int] = None,
+ overlap_history: int | None = None,
addnoise_condition: float = 0,
base_num_frames: int = 97,
ar_step: int = 0,
- causal_block_size: Optional[int] = None,
+ causal_block_size: int | None = None,
fps: int = 24,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -654,7 +651,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality. (**6.0 for T2V**, **5.0 for I2V**)
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -680,7 +677,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py
index 3e2004533258..a8f1b3a84a4a 100644
--- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py
+++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py
@@ -16,12 +16,11 @@
import math
import re
from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
-import ftfy
import PIL
import torch
-from transformers import AutoTokenizer, UMT5EncoderModel
+from transformers import AutoTokenizer, T5EncoderModel, UMT5EncoderModel
from diffusers.image_processor import PipelineImageInput
from diffusers.utils.torch_utils import randn_tensor
@@ -119,7 +118,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -159,7 +158,7 @@ class SkyReelsV2DiffusionForcingImageToVideoPipeline(DiffusionPipeline, SkyReels
def __init__(
self,
tokenizer: AutoTokenizer,
- text_encoder: UMT5EncoderModel,
+ text_encoder: T5EncoderModel | UMT5EncoderModel,
transformer: SkyReelsV2Transformer3DModel,
vae: AutoencoderKLWan,
scheduler: UniPCMultistepScheduler,
@@ -181,11 +180,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -223,23 +222,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -367,23 +366,23 @@ def check_inputs(
def prepare_latents(
self,
- image: Optional[PipelineImageInput],
+ image: PipelineImageInput | None,
batch_size: int,
num_channels_latents: int = 16,
height: int = 480,
width: int = 832,
num_frames: int = 97,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- last_image: Optional[torch.Tensor] = None,
- video_latents: Optional[torch.Tensor] = None,
- base_latent_num_frames: Optional[int] = None,
- causal_block_size: Optional[int] = None,
- overlap_history_latent_frames: Optional[int] = None,
- long_video_iter: Optional[int] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ last_image: torch.Tensor | None = None,
+ video_latents: torch.Tensor | None = None,
+ base_latent_num_frames: int | None = None,
+ causal_block_size: int | None = None,
+ overlap_history_latent_frames: int | None = None,
+ long_video_iter: int | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
latent_width = width // self.vae_scale_factor_spatial
@@ -508,7 +507,7 @@ def generate_timestep_matrix(
num_latent_frames]
- step_update_mask (torch.Tensor): Boolean mask indicating which frames to update Shape:
[num_iterations, num_latent_frames]
- - valid_interval (list[tuple]): List of (start, end) intervals for each iteration
+ - valid_interval (list[tuple]): list of (start, end) intervals for each iteration
Raises:
ValueError: If ar_step is too small for the given configuration
@@ -646,33 +645,31 @@ def attention_kwargs(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 544,
width: int = 960,
num_frames: int = 97,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- image_embeds: Optional[torch.Tensor] = None,
- last_image: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ image_embeds: torch.Tensor | None = None,
+ last_image: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
- overlap_history: Optional[int] = None,
+ overlap_history: int | None = None,
addnoise_condition: float = 0,
base_num_frames: int = 97,
ar_step: int = 0,
- causal_block_size: Optional[int] = None,
+ causal_block_size: int | None = None,
fps: int = 24,
):
r"""
@@ -681,10 +678,10 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -705,7 +702,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality. (**6.0 for T2V**, **5.0 for I2V**)
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -737,7 +734,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py
index 234ec531b862..924acb850d09 100644
--- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py
+++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py
@@ -17,12 +17,11 @@
import math
import re
from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
-import ftfy
import torch
from PIL import Image
-from transformers import AutoTokenizer, UMT5EncoderModel
+from transformers import AutoTokenizer, T5EncoderModel, UMT5EncoderModel
from ...callbacks import MultiPipelineCallbacks, PipelineCallback
from ...loaders import SkyReelsV2LoraLoaderMixin
@@ -116,10 +115,10 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -134,15 +133,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -175,7 +174,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -215,7 +214,7 @@ class SkyReelsV2DiffusionForcingVideoToVideoPipeline(DiffusionPipeline, SkyReels
def __init__(
self,
tokenizer: AutoTokenizer,
- text_encoder: UMT5EncoderModel,
+ text_encoder: T5EncoderModel | UMT5EncoderModel,
transformer: SkyReelsV2Transformer3DModel,
vae: AutoencoderKLWan,
scheduler: UniPCMultistepScheduler,
@@ -237,11 +236,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -279,23 +278,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -421,16 +420,16 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 97,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- video_latents: Optional[torch.Tensor] = None,
- base_latent_num_frames: Optional[int] = None,
- overlap_history: Optional[int] = None,
- causal_block_size: Optional[int] = None,
- overlap_history_latent_frames: Optional[int] = None,
- long_video_iter: Optional[int] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ video_latents: torch.Tensor | None = None,
+ base_latent_num_frames: int | None = None,
+ overlap_history: int | None = None,
+ causal_block_size: int | None = None,
+ overlap_history_latent_frames: int | None = None,
+ long_video_iter: int | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -546,7 +545,7 @@ def generate_timestep_matrix(
num_latent_frames]
- step_update_mask (torch.Tensor): Boolean mask indicating which frames to update Shape:
[num_iterations, num_latent_frames]
- - valid_interval (list[tuple]): List of (start, end) intervals for each iteration
+ - valid_interval (list[tuple]): list of (start, end) intervals for each iteration
Raises:
ValueError: If ar_step is too small for the given configuration
@@ -683,44 +682,42 @@ def attention_kwargs(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- video: List[Image.Image],
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ video: list[Image.Image],
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 544,
width: int = 960,
num_frames: int = 120,
num_inference_steps: int = 50,
guidance_scale: float = 6.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
- overlap_history: Optional[int] = None,
+ overlap_history: int | None = None,
addnoise_condition: float = 0,
base_num_frames: int = 97,
ar_step: int = 0,
- causal_block_size: Optional[int] = None,
+ causal_block_size: int | None = None,
fps: int = 24,
):
r"""
The call function to the pipeline for generation.
Args:
- video (`List[Image.Image]`):
+ video (`list[Image.Image]`):
The video to guide the video generation.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the video generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -741,7 +738,7 @@ def __call__(
usually at the expense of lower image quality. (**6.0 for T2V**, **5.0 for I2V**)
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -767,7 +764,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py
index d1df7f5f34cb..7c24b898e0bb 100644
--- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py
+++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py
@@ -13,12 +13,12 @@
# limitations under the License.
import html
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL
import regex as re
import torch
-from transformers import AutoTokenizer, CLIPProcessor, CLIPVisionModelWithProjection, UMT5EncoderModel
+from transformers import AutoTokenizer, CLIPProcessor, CLIPVisionModelWithProjection, T5EncoderModel, UMT5EncoderModel
from ...callbacks import MultiPipelineCallbacks, PipelineCallback
from ...image_processor import PipelineImageInput
@@ -112,7 +112,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -157,7 +157,7 @@ class SkyReelsV2ImageToVideoPipeline(DiffusionPipeline, SkyReelsV2LoraLoaderMixi
def __init__(
self,
tokenizer: AutoTokenizer,
- text_encoder: UMT5EncoderModel,
+ text_encoder: T5EncoderModel | UMT5EncoderModel,
image_encoder: CLIPVisionModelWithProjection,
image_processor: CLIPProcessor,
transformer: SkyReelsV2Transformer3DModel,
@@ -184,11 +184,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan_i2v.WanImageToVideoPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -227,7 +227,7 @@ def _get_t5_prompt_embeds(
def encode_image(
self,
image: PipelineImageInput,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
image = self.image_processor(images=image, return_tensors="pt").to(device)
@@ -237,23 +237,23 @@ def encode_image(
# Copied from diffusers.pipelines.wan.pipeline_wan_i2v.WanImageToVideoPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -378,12 +378,12 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- last_image: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ last_image: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
latent_width = width // self.vae_scale_factor_spatial
@@ -478,27 +478,25 @@ def attention_kwargs(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 544,
width: int = 960,
num_frames: int = 97,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- image_embeds: Optional[torch.Tensor] = None,
- last_image: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ image_embeds: torch.Tensor | None = None,
+ last_image: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
@@ -507,10 +505,10 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -531,7 +529,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -560,7 +558,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py
index 07b382dfc49f..d40269411bc0 100644
--- a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py
+++ b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py
@@ -14,7 +14,6 @@
from dataclasses import dataclass
from math import pi
-from typing import Optional
import torch
import torch.nn as nn
@@ -57,9 +56,9 @@ class StableAudioProjectionModelOutput(BaseOutput):
Sequence of hidden-states obtained by linearly projecting the audio end hidden states.
"""
- text_hidden_states: Optional[torch.Tensor] = None
- seconds_start_hidden_states: Optional[torch.Tensor] = None
- seconds_end_hidden_states: Optional[torch.Tensor] = None
+ text_hidden_states: torch.Tensor | None = None
+ seconds_start_hidden_states: torch.Tensor | None = None
+ seconds_end_hidden_states: torch.Tensor | None = None
class StableAudioNumberConditioner(nn.Module):
@@ -82,7 +81,7 @@ def __init__(
number_embedding_dim,
min_value,
max_value,
- internal_dim: Optional[int] = 256,
+ internal_dim: int | None = 256,
):
super().__init__()
self.time_positional_embedding = nn.Sequential(
@@ -138,9 +137,9 @@ def __init__(self, text_encoder_dim, conditioning_dim, min_value, max_value):
def forward(
self,
- text_hidden_states: Optional[torch.Tensor] = None,
- start_seconds: Optional[torch.Tensor] = None,
- end_seconds: Optional[torch.Tensor] = None,
+ text_hidden_states: torch.Tensor | None = None,
+ start_seconds: torch.Tensor | None = None,
+ end_seconds: torch.Tensor | None = None,
):
text_hidden_states = (
text_hidden_states if text_hidden_states is None else self.text_projection(text_hidden_states)
diff --git a/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py b/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py
index b7faf097ab0d..351c8b65de0e 100644
--- a/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py
+++ b/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import torch
from transformers import (
@@ -108,7 +108,7 @@ def __init__(
vae: AutoencoderOobleck,
text_encoder: T5EncoderModel,
projection_model: StableAudioProjectionModel,
- tokenizer: Union[T5Tokenizer, T5TokenizerFast],
+ tokenizer: T5Tokenizer | T5TokenizerFast,
transformer: StableAudioDiTModel,
scheduler: EDMDPMSolverMultistepScheduler,
):
@@ -158,10 +158,10 @@ def encode_prompt(
device,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.LongTensor] = None,
- negative_attention_mask: Optional[torch.LongTensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ negative_attention_mask: torch.LongTensor | None = None,
):
if prompt is not None and isinstance(prompt, str):
batch_size = 1
@@ -206,7 +206,7 @@ def encode_prompt(
prompt_embeds = prompt_embeds[0]
if do_classifier_free_guidance and negative_prompt is not None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -491,32 +491,32 @@ def prepare_latents(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- audio_end_in_s: Optional[float] = None,
- audio_start_in_s: Optional[float] = 0.0,
+ prompt: str | list[str] = None,
+ audio_end_in_s: float | None = None,
+ audio_start_in_s: float | None = 0.0,
num_inference_steps: int = 100,
guidance_scale: float = 7.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_waveforms_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_waveforms_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- initial_audio_waveforms: Optional[torch.Tensor] = None,
- initial_audio_sampling_rate: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.LongTensor] = None,
- negative_attention_mask: Optional[torch.LongTensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ initial_audio_waveforms: torch.Tensor | None = None,
+ initial_audio_sampling_rate: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ negative_attention_mask: torch.LongTensor | None = None,
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
- output_type: Optional[str] = "pt",
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
+ output_type: str | None = "pt",
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide audio generation. If not defined, you need to pass `prompt_embeds`.
audio_end_in_s (`float`, *optional*, defaults to 47.55):
Audio end index in seconds.
@@ -528,7 +528,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.0):
A higher guidance scale value encourages the model to generate audio that is closely linked to the text
`prompt` at the expense of lower sound quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in audio generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_waveforms_per_prompt (`int`, *optional*, defaults to 1):
@@ -536,7 +536,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py
index a6a60ad94be6..ef40078bfbb9 100644
--- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py
+++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
@@ -139,10 +139,10 @@ def encode_prompt(
do_classifier_free_guidance,
prompt=None,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_pooled: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_pooled: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_pooled: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_pooled: torch.Tensor | None = None,
):
if prompt_embeds is None:
# get prompt text embeddings
@@ -184,7 +184,7 @@ def encode_prompt(
prompt_embeds_pooled = prompt_embeds_pooled.repeat_interleave(num_images_per_prompt, dim=0)
if negative_prompt_embeds is None and do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -305,30 +305,30 @@ def get_timestep_ratio_conditioning(self, t, alphas_cumprod):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image_embeddings: Union[torch.Tensor, List[torch.Tensor]],
- prompt: Union[str, List[str]] = None,
+ image_embeddings: torch.Tensor | list[torch.Tensor],
+ prompt: str | list[str] = None,
num_inference_steps: int = 10,
guidance_scale: float = 0.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_pooled: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_pooled: Optional[torch.Tensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_pooled: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_pooled: torch.Tensor | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
"""
Function invoked when calling the pipeline for generation.
Args:
- image_embedding (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embedding (`torch.Tensor` or `list[torch.Tensor]`):
Image Embeddings either extracted from an image or generated by a Prior Model.
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
num_inference_steps (`int`, *optional*, defaults to 12):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
@@ -339,7 +339,7 @@ def __call__(
equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by
setting `decoder_guidance_scale > 1`. Higher guidance scale encourages to generate images that are
closely linked to the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `decoder_guidance_scale` is less than `1`).
prompt_embeds (`torch.Tensor`, *optional*):
@@ -358,7 +358,7 @@ def __call__(
input argument.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -375,7 +375,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py
index 838b93faaa0c..50e6c02b6017 100644
--- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py
+++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py
@@ -11,7 +11,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.
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import PIL
import torch
@@ -90,8 +90,8 @@ def __init__(
prior_text_encoder: CLIPTextModelWithProjection,
prior_tokenizer: CLIPTokenizer,
prior_scheduler: DDPMWuerstchenScheduler,
- prior_feature_extractor: Optional[CLIPImageProcessor] = None,
- prior_image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ prior_feature_extractor: CLIPImageProcessor | None = None,
+ prior_image_encoder: CLIPVisionModelWithProjection | None = None,
):
super().__init__()
@@ -124,10 +124,10 @@ def __init__(
vqgan=vqgan,
)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
- def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
@@ -137,7 +137,7 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t
self.prior_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device)
self.decoder_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device)
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models (`unet`, `text_encoder`, `vae`, and `safety checker` state dicts) to CPU using 🤗
Accelerate, significantly reducing memory usage. Models are moved to a `torch.device('meta')` and loaded on a
@@ -159,38 +159,38 @@ def set_progress_bar_config(self, **kwargs):
@replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- images: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]] = None,
+ prompt: str | list[str] | None = None,
+ images: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image] = None,
height: int = 512,
width: int = 512,
prior_num_inference_steps: int = 60,
prior_guidance_scale: float = 4.0,
num_inference_steps: int = 12,
decoder_guidance_scale: float = 0.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_pooled: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_pooled: Optional[torch.Tensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_pooled: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_pooled: torch.Tensor | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ prior_callback_on_step_end: Callable[[int, int], None] | None = None,
+ prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation for the prior and decoder.
- images (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, *optional*):
+ images (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, *optional*):
The images to guide the image generation for the prior.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
prompt_embeds (`torch.Tensor`, *optional*):
@@ -219,7 +219,7 @@ def __call__(
equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by
setting `prior_guidance_scale > 1`. Higher guidance scale encourages to generate images that are
closely linked to the text `prompt`, usually at the expense of lower image quality.
- prior_num_inference_steps (`Union[int, Dict[float, int]]`, *optional*, defaults to 60):
+ prior_num_inference_steps (`int | dict[float, int]`, *optional*, defaults to 60):
The number of prior denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference. For more specific timestep spacing, you can pass customized
`prior_timesteps`
@@ -233,7 +233,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -249,7 +249,7 @@ def __call__(
A function that calls at the end of each denoising steps during the inference. The function is called
with the following arguments: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep:
int, callback_kwargs: Dict)`.
- prior_callback_on_step_end_tensor_inputs (`List`, *optional*):
+ prior_callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `prior_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.
@@ -258,7 +258,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py
index 29ad8b5429d7..0c5ea9ed61b4 100644
--- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py
+++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py
@@ -14,7 +14,7 @@
from dataclasses import dataclass
from math import ceil
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL
@@ -70,11 +70,11 @@ class StableCascadePriorPipelineOutput(BaseOutput):
Text embeddings for the negative prompt.
"""
- image_embeddings: Union[torch.Tensor, np.ndarray]
- prompt_embeds: Union[torch.Tensor, np.ndarray]
- prompt_embeds_pooled: Union[torch.Tensor, np.ndarray]
- negative_prompt_embeds: Union[torch.Tensor, np.ndarray]
- negative_prompt_embeds_pooled: Union[torch.Tensor, np.ndarray]
+ image_embeddings: torch.Tensor | np.ndarray
+ prompt_embeds: torch.Tensor | np.ndarray
+ prompt_embeds_pooled: torch.Tensor | np.ndarray
+ negative_prompt_embeds: torch.Tensor | np.ndarray
+ negative_prompt_embeds_pooled: torch.Tensor | np.ndarray
class StableCascadePriorPipeline(DeprecatedPipelineMixin, DiffusionPipeline):
@@ -118,8 +118,8 @@ def __init__(
prior: StableCascadeUNet,
scheduler: DDPMWuerstchenScheduler,
resolution_multiple: float = 42.67,
- feature_extractor: Optional[CLIPImageProcessor] = None,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ feature_extractor: CLIPImageProcessor | None = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
) -> None:
super().__init__()
self.register_modules(
@@ -160,10 +160,10 @@ def encode_prompt(
do_classifier_free_guidance,
prompt=None,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_pooled: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_pooled: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_pooled: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_pooled: torch.Tensor | None = None,
):
if prompt_embeds is None:
# get prompt text embeddings
@@ -205,7 +205,7 @@ def encode_prompt(
prompt_embeds_pooled = prompt_embeds_pooled.repeat_interleave(num_images_per_prompt, dim=0)
if negative_prompt_embeds is None and do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -376,32 +376,32 @@ def get_timestep_ratio_conditioning(self, t, alphas_cumprod):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- images: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]] = None,
+ prompt: str | list[str] | None = None,
+ images: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image] = None,
height: int = 1024,
width: int = 1024,
num_inference_steps: int = 20,
- timesteps: List[float] = None,
+ timesteps: list[float] = None,
guidance_scale: float = 4.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- prompt_embeds_pooled: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds_pooled: Optional[torch.Tensor] = None,
- image_embeds: Optional[torch.Tensor] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pt",
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_pooled: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_pooled: torch.Tensor | None = None,
+ image_embeds: torch.Tensor | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pt",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
height (`int`, *optional*, defaults to 1024):
The height in pixels of the generated image.
@@ -416,7 +416,7 @@ def __call__(
equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by
setting `decoder_guidance_scale > 1`. Higher guidance scale encourages to generate images that are
closely linked to the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `decoder_guidance_scale` is less than `1`).
prompt_embeds (`torch.Tensor`, *optional*):
@@ -438,7 +438,7 @@ def __call__(
not provided, image embeddings will be generated from `image` input argument if existing.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -455,7 +455,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion/__init__.py b/src/diffusers/pipelines/stable_diffusion/__init__.py
index b05a3ce2a7c9..c2eebf586ef8 100644
--- a/src/diffusers/pipelines/stable_diffusion/__init__.py
+++ b/src/diffusers/pipelines/stable_diffusion/__init__.py
@@ -6,8 +6,6 @@
_LazyModule,
get_objects_from_module,
is_flax_available,
- is_k_diffusion_available,
- is_k_diffusion_version,
is_onnx_available,
is_torch_available,
is_transformers_available,
diff --git a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
index 6c0221d2092a..8fa4bd5941d2 100644
--- a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
+++ b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
@@ -17,7 +17,6 @@
import re
from contextlib import nullcontext
from io import BytesIO
-from typing import Dict, Optional, Union
import requests
import torch
@@ -1049,7 +1048,7 @@ def stable_unclip_image_encoder(original_config, local_files_only=False):
def stable_unclip_image_noising_components(
- original_config, clip_stats_path: Optional[str] = None, device: Optional[str] = None
+ original_config, clip_stats_path: str | None = None, device: str | None = None
):
"""
Returns the noising components for the img2img and txt2img unclip pipelines.
@@ -1144,25 +1143,25 @@ def convert_controlnet_checkpoint(
def download_from_original_stable_diffusion_ckpt(
- checkpoint_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+ checkpoint_path_or_dict: str | dict[str, torch.Tensor],
original_config_file: str = None,
- image_size: Optional[int] = None,
+ image_size: int | None = None,
prediction_type: str = None,
model_type: str = None,
extract_ema: bool = False,
scheduler_type: str = "pndm",
- num_in_channels: Optional[int] = None,
- upcast_attention: Optional[bool] = None,
+ num_in_channels: int | None = None,
+ upcast_attention: bool | None = None,
device: str = None,
from_safetensors: bool = False,
- stable_unclip: Optional[str] = None,
- stable_unclip_prior: Optional[str] = None,
- clip_stats_path: Optional[str] = None,
- controlnet: Optional[bool] = None,
- adapter: Optional[bool] = None,
+ stable_unclip: str | None = None,
+ stable_unclip_prior: str | None = None,
+ clip_stats_path: str | None = None,
+ controlnet: bool | None = None,
+ adapter: bool | None = None,
load_safety_checker: bool = True,
- safety_checker: Optional[StableDiffusionSafetyChecker] = None,
- feature_extractor: Optional[AutoFeatureExtractor] = None,
+ safety_checker: StableDiffusionSafetyChecker | None = None,
+ feature_extractor: AutoFeatureExtractor | None = None,
pipeline_class: DiffusionPipeline = None,
local_files_only=False,
vae_path=None,
@@ -1237,7 +1236,7 @@ def download_from_original_stable_diffusion_ckpt(
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer)
to use. If this parameter is `None`, the function will load a new instance of [CLIPTokenizer] by itself, if
needed.
- config_files (`Dict[str, str]`, *optional*, defaults to `None`):
+ config_files (`dict[str, str]`, *optional*, defaults to `None`):
A dictionary mapping from config file names to their contents. If this parameter is `None`, the function
will load the config files by itself, if needed. Valid keys are:
- `v1`: Config file for Stable Diffusion v1
@@ -1827,12 +1826,12 @@ def download_controlnet_from_original_ckpt(
original_config_file: str,
image_size: int = 512,
extract_ema: bool = False,
- num_in_channels: Optional[int] = None,
- upcast_attention: Optional[bool] = None,
+ num_in_channels: int | None = None,
+ upcast_attention: bool | None = None,
device: str = None,
from_safetensors: bool = False,
- use_linear_projection: Optional[bool] = None,
- cross_attention_dim: Optional[bool] = None,
+ use_linear_projection: bool | None = None,
+ cross_attention_dim: bool | None = None,
) -> DiffusionPipeline:
if from_safetensors:
from safetensors import safe_open
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
index 6befe77aa4b1..c15163917a7c 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
@@ -14,7 +14,6 @@
import warnings
from functools import partial
-from typing import Dict, List, Optional, Union
import jax
import jax.numpy as jnp
@@ -112,9 +111,7 @@ def __init__(
text_encoder: FlaxCLIPTextModel,
tokenizer: CLIPTokenizer,
unet: FlaxUNet2DConditionModel,
- scheduler: Union[
- FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler
- ],
+ scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler,
safety_checker: FlaxStableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
dtype: jnp.dtype = jnp.float32,
@@ -168,7 +165,7 @@ def __init__(
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
- def prepare_inputs(self, prompt: Union[str, List[str]]):
+ def prepare_inputs(self, prompt: str | list[str]):
if not isinstance(prompt, (str, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
@@ -218,14 +215,14 @@ def _run_safety_checker(self, images, safety_model_params, jit=False):
def _generate(
self,
prompt_ids: jnp.array,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
num_inference_steps: int,
height: int,
width: int,
guidance_scale: float,
- latents: Optional[jnp.ndarray] = None,
- neg_prompt_ids: Optional[jnp.ndarray] = None,
+ latents: jnp.ndarray | None = None,
+ neg_prompt_ids: jnp.ndarray | None = None,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
@@ -315,12 +312,12 @@ def loop_body(step, args):
def __call__(
self,
prompt_ids: jnp.array,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
num_inference_steps: int = 50,
- height: Optional[int] = None,
- width: Optional[int] = None,
- guidance_scale: Union[float, jnp.ndarray] = 7.5,
+ height: int | None = None,
+ width: int | None = None,
+ guidance_scale: float | jnp.ndarray = 7.5,
latents: jnp.ndarray = None,
neg_prompt_ids: jnp.ndarray = None,
return_dict: bool = True,
@@ -330,7 +327,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py
index 81656beba7e1..e1c004ceb44c 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py
@@ -14,7 +14,6 @@
import warnings
from functools import partial
-from typing import Dict, List, Optional, Union
import jax
import jax.numpy as jnp
@@ -136,9 +135,7 @@ def __init__(
text_encoder: FlaxCLIPTextModel,
tokenizer: CLIPTokenizer,
unet: FlaxUNet2DConditionModel,
- scheduler: Union[
- FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler
- ],
+ scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler,
safety_checker: FlaxStableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
dtype: jnp.dtype = jnp.float32,
@@ -167,7 +164,7 @@ def __init__(
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
- def prepare_inputs(self, prompt: Union[str, List[str]], image: Union[Image.Image, List[Image.Image]]):
+ def prepare_inputs(self, prompt: str | list[str], image: Image.Image | list[Image.Image]):
if not isinstance(prompt, (str, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
@@ -234,15 +231,15 @@ def _generate(
self,
prompt_ids: jnp.ndarray,
image: jnp.ndarray,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
start_timestep: int,
num_inference_steps: int,
height: int,
width: int,
guidance_scale: float,
- noise: Optional[jnp.ndarray] = None,
- neg_prompt_ids: Optional[jnp.ndarray] = None,
+ noise: jnp.ndarray | None = None,
+ neg_prompt_ids: jnp.ndarray | None = None,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
@@ -339,13 +336,13 @@ def __call__(
self,
prompt_ids: jnp.ndarray,
image: jnp.ndarray,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
strength: float = 0.8,
num_inference_steps: int = 50,
- height: Optional[int] = None,
- width: Optional[int] = None,
- guidance_scale: Union[float, jnp.ndarray] = 7.5,
+ height: int | None = None,
+ width: int | None = None,
+ guidance_scale: float | jnp.ndarray = 7.5,
noise: jnp.ndarray = None,
neg_prompt_ids: jnp.ndarray = None,
return_dict: bool = True,
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py
index 5938fe232a71..dbda93deebff 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py
@@ -14,7 +14,6 @@
import warnings
from functools import partial
-from typing import Dict, List, Optional, Union
import jax
import jax.numpy as jnp
@@ -135,9 +134,7 @@ def __init__(
text_encoder: FlaxCLIPTextModel,
tokenizer: CLIPTokenizer,
unet: FlaxUNet2DConditionModel,
- scheduler: Union[
- FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler
- ],
+ scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler,
safety_checker: FlaxStableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
dtype: jnp.dtype = jnp.float32,
@@ -193,9 +190,9 @@ def __init__(
def prepare_inputs(
self,
- prompt: Union[str, List[str]],
- image: Union[Image.Image, List[Image.Image]],
- mask: Union[Image.Image, List[Image.Image]],
+ prompt: str | list[str],
+ image: Image.Image | list[Image.Image],
+ mask: Image.Image | list[Image.Image],
):
if not isinstance(prompt, (str, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
@@ -269,14 +266,14 @@ def _generate(
prompt_ids: jnp.ndarray,
mask: jnp.ndarray,
masked_image: jnp.ndarray,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
num_inference_steps: int,
height: int,
width: int,
guidance_scale: float,
- latents: Optional[jnp.ndarray] = None,
- neg_prompt_ids: Optional[jnp.ndarray] = None,
+ latents: jnp.ndarray | None = None,
+ neg_prompt_ids: jnp.ndarray | None = None,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
@@ -397,12 +394,12 @@ def __call__(
prompt_ids: jnp.ndarray,
mask: jnp.ndarray,
masked_image: jnp.ndarray,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
num_inference_steps: int = 50,
- height: Optional[int] = None,
- width: Optional[int] = None,
- guidance_scale: Union[float, jnp.ndarray] = 7.5,
+ height: int | None = None,
+ width: int | None = None,
+ guidance_scale: float | jnp.ndarray = 7.5,
latents: jnp.ndarray = None,
neg_prompt_ids: jnp.ndarray = None,
return_dict: bool = True,
@@ -412,7 +409,7 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py
index 6ebe0986a1ab..0f66ca909e7d 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -36,7 +36,7 @@ class OnnxStableDiffusionPipeline(DiffusionPipeline):
text_encoder: OnnxRuntimeModel
tokenizer: CLIPTokenizer
unet: OnnxRuntimeModel
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler
safety_checker: OnnxRuntimeModel
feature_extractor: CLIPImageProcessor
@@ -50,7 +50,7 @@ def __init__(
text_encoder: OnnxRuntimeModel,
tokenizer: CLIPTokenizer,
unet: OnnxRuntimeModel,
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler,
safety_checker: OnnxRuntimeModel,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = True,
@@ -114,24 +114,24 @@ def __init__(
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- num_images_per_prompt: Optional[int],
+ prompt: str | list[str],
+ num_images_per_prompt: int | None,
do_classifier_free_guidance: bool,
- negative_prompt: Optional[str],
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
+ negative_prompt: str | None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
num_images_per_prompt (`int`):
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
prompt_embeds (`np.ndarray`, *optional*):
@@ -176,7 +176,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -217,13 +217,13 @@ def _encode_prompt(
def check_inputs(
self,
- prompt: Union[str, List[str]],
- height: Optional[int],
- width: Optional[int],
+ prompt: str | list[str],
+ height: int | None,
+ width: int | None,
callback_steps: int,
- negative_prompt: Optional[str] = None,
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
+ negative_prompt: str | None = None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
@@ -264,31 +264,31 @@ def check_inputs(
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = 512,
- width: Optional[int] = 512,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[np.random.RandomState] = None,
- latents: Optional[np.ndarray] = None,
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
- output_type: Optional[str] = "pil",
+ prompt: str | list[str] = None,
+ height: int | None = 512,
+ width: int | None = 512,
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 7.5,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: np.random.RandomState | None = None,
+ latents: np.ndarray | None = None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
+ callback: Callable[[int, int, np.ndarray], None] | None = None,
callback_steps: int = 1,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- image (`PIL.Image.Image` or List[`PIL.Image.Image`] or `torch.Tensor`):
+ image (`PIL.Image.Image` or list[`PIL.Image.Image`] or `torch.Tensor`):
`Image`, or tensor representing an image batch which will be upscaled. *
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
@@ -299,7 +299,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale`
is less than `1`).
@@ -470,7 +470,7 @@ def __init__(
text_encoder: OnnxRuntimeModel,
tokenizer: CLIPTokenizer,
unet: OnnxRuntimeModel,
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler,
safety_checker: OnnxRuntimeModel,
feature_extractor: CLIPImageProcessor,
):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py
index d63bf3bf4564..2bb3eb6990e1 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -89,7 +89,7 @@ class OnnxStableDiffusionImg2ImgPipeline(DiffusionPipeline):
text_encoder: OnnxRuntimeModel
tokenizer: CLIPTokenizer
unet: OnnxRuntimeModel
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler
safety_checker: OnnxRuntimeModel
feature_extractor: CLIPImageProcessor
@@ -103,7 +103,7 @@ def __init__(
text_encoder: OnnxRuntimeModel,
tokenizer: CLIPTokenizer,
unet: OnnxRuntimeModel,
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler,
safety_checker: OnnxRuntimeModel,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = True,
@@ -168,24 +168,24 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion.OnnxStableDiffusionPipeline._encode_prompt
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- num_images_per_prompt: Optional[int],
+ prompt: str | list[str],
+ num_images_per_prompt: int | None,
do_classifier_free_guidance: bool,
- negative_prompt: Optional[str],
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
+ negative_prompt: str | None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
num_images_per_prompt (`int`):
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
prompt_embeds (`np.ndarray`, *optional*):
@@ -230,7 +230,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -271,11 +271,11 @@ def _encode_prompt(
def check_inputs(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
callback_steps: int,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
):
if (callback_steps is None) or (
callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
@@ -313,27 +313,27 @@ def check_inputs(
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[np.ndarray, PIL.Image.Image] = None,
+ prompt: str | list[str],
+ image: np.ndarray | PIL.Image.Image = None,
strength: float = 0.8,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[np.random.RandomState] = None,
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
- output_type: Optional[str] = "pil",
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 7.5,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: np.random.RandomState | None = None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
+ callback: Callable[[int, int, np.ndarray], None] | None = None,
callback_steps: int = 1,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
image (`np.ndarray` or `PIL.Image.Image`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
@@ -353,7 +353,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py
index 158bcabbebfd..7e3416a33832 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -87,7 +87,7 @@ class OnnxStableDiffusionInpaintPipeline(DiffusionPipeline):
text_encoder: OnnxRuntimeModel
tokenizer: CLIPTokenizer
unet: OnnxRuntimeModel
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler
safety_checker: OnnxRuntimeModel
feature_extractor: CLIPImageProcessor
@@ -101,7 +101,7 @@ def __init__(
text_encoder: OnnxRuntimeModel,
tokenizer: CLIPTokenizer,
unet: OnnxRuntimeModel,
- scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
+ scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler,
safety_checker: OnnxRuntimeModel,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = True,
@@ -167,24 +167,24 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion.OnnxStableDiffusionPipeline._encode_prompt
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- num_images_per_prompt: Optional[int],
+ prompt: str | list[str],
+ num_images_per_prompt: int | None,
do_classifier_free_guidance: bool,
- negative_prompt: Optional[str],
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
+ negative_prompt: str | None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
num_images_per_prompt (`int`):
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
prompt_embeds (`np.ndarray`, *optional*):
@@ -229,7 +229,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -271,13 +271,13 @@ def _encode_prompt(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion.OnnxStableDiffusionPipeline.check_inputs
def check_inputs(
self,
- prompt: Union[str, List[str]],
- height: Optional[int],
- width: Optional[int],
+ prompt: str | list[str],
+ height: int | None,
+ width: int | None,
callback_steps: int,
- negative_prompt: Optional[str] = None,
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
+ negative_prompt: str | None = None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
@@ -319,30 +319,30 @@ def check_inputs(
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
image: PIL.Image.Image,
mask_image: PIL.Image.Image,
- height: Optional[int] = 512,
- width: Optional[int] = 512,
+ height: int | None = 512,
+ width: int | None = 512,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[np.random.RandomState] = None,
- latents: Optional[np.ndarray] = None,
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
- output_type: Optional[str] = "pil",
+ generator: np.random.RandomState | None = None,
+ latents: np.ndarray | None = None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
+ callback: Callable[[int, int, np.ndarray], None] | None = None,
callback_steps: int = 1,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
image (`PIL.Image.Image`):
`Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will
@@ -365,7 +365,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py
index a765163175a2..c721a4f2ffa0 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -74,8 +74,8 @@ def __init__(
unet: OnnxRuntimeModel,
low_res_scheduler: DDPMScheduler,
scheduler: KarrasDiffusionSchedulers,
- safety_checker: Optional[OnnxRuntimeModel] = None,
- feature_extractor: Optional[CLIPImageProcessor] = None,
+ safety_checker: OnnxRuntimeModel | None = None,
+ feature_extractor: CLIPImageProcessor | None = None,
max_noise_level: int = 350,
num_latent_channels=4,
num_unet_input_channels=7,
@@ -144,7 +144,7 @@ def __init__(
def check_inputs(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
image,
noise_level,
callback_steps,
@@ -245,24 +245,24 @@ def decode_latents(self, latents):
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- num_images_per_prompt: Optional[int],
+ prompt: str | list[str],
+ num_images_per_prompt: int | None,
do_classifier_free_guidance: bool,
- negative_prompt: Optional[str],
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
+ negative_prompt: str | None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
num_images_per_prompt (`int`):
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
prompt_embeds (`np.ndarray`, *optional*):
@@ -307,7 +307,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -348,28 +348,28 @@ def _encode_prompt(
def __call__(
self,
- prompt: Union[str, List[str]],
- image: Union[np.ndarray, PIL.Image.Image, List[PIL.Image.Image]],
+ prompt: str | list[str],
+ image: np.ndarray | PIL.Image.Image | list[PIL.Image.Image],
num_inference_steps: int = 75,
guidance_scale: float = 9.0,
noise_level: int = 20,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[np.random.RandomState, List[np.random.RandomState]]] = None,
- latents: Optional[np.ndarray] = None,
- prompt_embeds: Optional[np.ndarray] = None,
- negative_prompt_embeds: Optional[np.ndarray] = None,
- output_type: Optional[str] = "pil",
+ generator: np.random.RandomState | list[np.random.RandomState] | None = None,
+ latents: np.ndarray | None = None,
+ prompt_embeds: np.ndarray | None = None,
+ negative_prompt_embeds: np.ndarray | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
- callback_steps: Optional[int] = 1,
+ callback: Callable[[int, int, np.ndarray], None] | None = None,
+ callback_steps: int | None = 1,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
image (`np.ndarray` or `PIL.Image.Image`):
`Image`, or tensor representing an image batch, that will be used as the starting point for the
@@ -385,7 +385,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
noise_level (`float`, defaults to 0.2):
Deteremines the amount of noise to add to the initial image before performing upscaling.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion/pipeline_output.py
index 5fb9b1a1412d..804e5655d369 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_output.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -13,16 +12,16 @@ class StableDiffusionPipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
- nsfw_content_detected (`List[bool]`)
- List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
+ nsfw_content_detected (`list[bool]`)
+ list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
`None` if safety checking could not be performed.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_content_detected: Optional[List[bool]]
+ images: list[PIL.Image.Image] | np.ndarray
+ nsfw_content_detected: list[bool] | None
if is_flax_available():
@@ -36,10 +35,10 @@ class FlaxStableDiffusionPipelineOutput(BaseOutput):
Args:
images (`np.ndarray`):
Denoised images of array shape of `(batch_size, height, width, num_channels)`.
- nsfw_content_detected (`List[bool]`):
- List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content
+ nsfw_content_detected (`list[bool]`):
+ list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content
or `None` if safety checking could not be performed.
"""
images: np.ndarray
- nsfw_content_detected: List[bool]
+ nsfw_content_detected: list[bool]
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
index d079d2a225cf..42d62f53a20a 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from packaging import version
@@ -94,10 +94,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -112,15 +112,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -304,9 +304,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -336,16 +336,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -353,7 +353,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -452,7 +452,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -779,38 +779,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -819,18 +817,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -838,7 +836,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -852,7 +850,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -877,7 +875,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py
index e957c6661f87..3ec64c30763f 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py
@@ -14,7 +14,7 @@
import contextlib
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -53,7 +53,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -175,9 +175,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -208,16 +208,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -225,7 +225,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -324,7 +324,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -634,33 +634,33 @@ def num_timesteps(self):
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
- depth_map: Optional[torch.Tensor] = None,
+ depth_map: torch.Tensor | None = None,
strength: float = 0.8,
- num_inference_steps: Optional[int] = 50,
- guidance_scale: Optional[float] = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ num_inference_steps: int | None = 50,
+ guidance_scale: float | None = 7.5,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image` or tensor representing an image batch to be used as the starting point. Can accept image
latents as `image` only if `depth_map` is not `None`.
depth_map (`torch.Tensor`, *optional*):
@@ -678,7 +678,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -686,7 +686,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -711,7 +711,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py
index d47e2f0593dd..5974aea6b448 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Callable, List, Optional, Union
+from typing import Callable
import PIL.Image
import torch
@@ -218,7 +218,7 @@ def check_inputs(self, image, height, width, callback_steps):
and not isinstance(image, list)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+ "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
f" {type(image)}"
)
@@ -259,25 +259,25 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@torch.no_grad()
def __call__(
self,
- image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- num_images_per_prompt: Optional[int] = 1,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
):
r"""
The call function to the pipeline for generation.
Args:
- image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`):
+ image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.Tensor`):
Image or images to guide image generation. If you provide a tensor, it needs to be compatible with
[`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
@@ -295,7 +295,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
index d0be0ee51317..abcd06a2bb3b 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -84,7 +84,7 @@
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -122,10 +122,10 @@ def preprocess(image):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -140,15 +140,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -329,9 +329,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -362,16 +362,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -379,7 +379,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -478,7 +478,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -859,38 +859,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
strength: float = 0.8,
- num_inference_steps: Optional[int] = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- guidance_scale: Optional[float] = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: Optional[float] = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ num_inference_steps: int | None = 50,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ guidance_scale: float | None = 7.5,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ eta: float | None = 0.0,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
clip_skip: int = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -905,18 +903,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference. This parameter is modulated by `strength`.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -924,7 +922,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -934,7 +932,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -955,7 +953,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
index 82902cc7dcd0..a3e09d1ed1ad 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -53,7 +53,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -68,10 +68,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -86,15 +86,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -173,7 +173,7 @@ class StableDiffusionInpaintPipeline(
def __init__(
self,
- vae: Union[AutoencoderKL, AsymmetricAutoencoderKL],
+ vae: AutoencoderKL | AsymmetricAutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
@@ -283,9 +283,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -316,16 +316,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -333,7 +333,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -432,7 +432,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -880,51 +880,49 @@ def interrupt(self):
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: torch.Tensor = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 1.0,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
clip_skip: int = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be inpainted (which parts of the image to
be masked out with `mask_image` and repainted according to `prompt`). For both numpy array and pytorch
tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the
expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the
expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but
if passing latents directly it is not encoded again.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
@@ -951,18 +949,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference. This parameter is modulated by `strength`.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -970,7 +968,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -984,7 +982,7 @@ def __call__(
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1005,7 +1003,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
index 843d25d67c10..c89d593d57be 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -68,7 +68,7 @@ def preprocess(image):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -134,7 +134,7 @@ def __init__(
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -172,36 +172,34 @@ def __init__(
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
num_inference_steps: int = 100,
guidance_scale: float = 7.5,
image_guidance_scale: float = 1.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ cross_attention_kwargs: dict[str, Any] | None = None,
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image` or tensor representing an image batch to be repainted according to `prompt`. Can also accept
image latents as `image`, but if passing latents directly it is not encoded again.
num_inference_steps (`int`, *optional*, defaults to 100):
@@ -215,7 +213,7 @@ def __call__(
`image_guidance_scale > 1`. Higher image guidance scale encourages generated images that are closely
linked to the source `image`, usually at the expense of lower image quality. This pipeline requires a
value of at least `1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -248,7 +246,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
@@ -497,14 +495,14 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -512,7 +510,7 @@ def _encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_ prompt (`str` or `List[str]`, *optional*):
+ negative_ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -579,7 +577,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py
index a1d0407caf5e..02dc483c277a 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py
@@ -13,7 +13,7 @@
# limitations under the License.
import warnings
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -42,7 +42,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -132,10 +132,10 @@ def _encode_prompt(
device,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -169,10 +169,10 @@ def encode_prompt(
device,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
@@ -184,7 +184,7 @@ def encode_prompt(
torch device
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
prompt_embeds (`torch.FloatTensor`, *optional*):
@@ -243,7 +243,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance:
if negative_prompt_embeds is None or negative_pooled_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -396,29 +396,29 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
num_inference_steps: int = 75,
guidance_scale: float = 9.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image upscaling.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image` or tensor representing an image batch to be upscaled. If it's a tensor, it can be either a
latent output from a Stable Diffusion model or an image tensor in the range `[-1, 1]`. It is considered
a `latent` if `image.shape[1]` is `4`; otherwise, it is considered to be an image representation and
@@ -429,13 +429,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
index f13cdc67073f..4befa44550b7 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
@@ -14,7 +14,7 @@
import inspect
import warnings
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -123,9 +123,9 @@ def __init__(
unet: UNet2DConditionModel,
low_res_scheduler: DDPMScheduler,
scheduler: KarrasDiffusionSchedulers,
- safety_checker: Optional[Any] = None,
- feature_extractor: Optional[CLIPImageProcessor] = None,
- watermarker: Optional[Any] = None,
+ safety_checker: Any | None = None,
+ feature_extractor: CLIPImageProcessor | None = None,
+ watermarker: Any | None = None,
max_noise_level: int = 350,
):
super().__init__()
@@ -188,9 +188,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -221,16 +221,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -238,7 +238,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -337,7 +337,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -535,32 +535,32 @@ def upcast_vae(self):
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
num_inference_steps: int = 75,
guidance_scale: float = 9.0,
noise_level: int = 20,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
clip_skip: int = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image` or tensor representing an image batch to be upscaled.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
@@ -568,7 +568,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -576,7 +576,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py
index a134244e3ee4..39857ed230e6 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
@@ -173,8 +173,8 @@ def _encode_prior_prompt(
device,
num_images_per_prompt,
do_classifier_free_guidance,
- text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None,
- text_attention_mask: Optional[torch.Tensor] = None,
+ text_model_output: CLIPTextModelOutput | tuple | None = None,
+ text_attention_mask: torch.Tensor | None = None,
):
if text_model_output is None:
batch_size = len(prompt) if isinstance(prompt, list) else 1
@@ -268,9 +268,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -301,16 +301,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -318,7 +318,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -417,7 +417,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -600,8 +600,8 @@ def noise_image_embeddings(
self,
image_embeds: torch.Tensor,
noise_level: int,
- noise: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ noise: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
):
"""
Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher
@@ -647,35 +647,35 @@ def noise_image_embeddings(
def __call__(
self,
# regular denoising process args
- prompt: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 20,
guidance_scale: float = 10.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
noise_level: int = 0,
# prior args
prior_num_inference_steps: int = 25,
prior_guidance_scale: float = 4.0,
- prior_latents: Optional[torch.Tensor] = None,
- clip_skip: Optional[int] = None,
+ prior_latents: torch.Tensor | None = None,
+ clip_skip: int | None = None,
):
"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -687,7 +687,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 10.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -695,7 +695,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py
index abb4cc3a05d5..bb96e5db0295 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -175,9 +175,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -263,16 +263,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -280,7 +280,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -379,7 +379,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -547,7 +547,7 @@ def check_inputs(
and not isinstance(image, list)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+ "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
f" {type(image)}"
)
@@ -579,8 +579,8 @@ def noise_image_embeddings(
self,
image_embeds: torch.Tensor,
noise_level: int,
- noise: Optional[torch.Tensor] = None,
- generator: Optional[torch.Generator] = None,
+ noise: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
):
"""
Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher
@@ -625,33 +625,33 @@ def noise_image_embeddings(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Union[torch.Tensor, PIL.Image.Image] = None,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ image: torch.Tensor | PIL.Image.Image = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 20,
guidance_scale: float = 10,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
noise_level: int = 0,
- image_embeds: Optional[torch.Tensor] = None,
- clip_skip: Optional[int] = None,
+ image_embeds: torch.Tensor | None = None,
+ clip_skip: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, either `prompt_embeds` will be
used or prompt is initialized to `""`.
image (`torch.Tensor` or `PIL.Image.Image`):
@@ -668,7 +668,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 10.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -676,7 +676,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py b/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py
index 55d5023b6869..77a31b91d52d 100644
--- a/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py
+++ b/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple
-
import jax
import jax.numpy as jnp
from flax import linen as nn
@@ -76,7 +74,7 @@ class FlaxStableDiffusionSafetyChecker(FlaxPreTrainedModel):
def __init__(
self,
config: CLIPConfig,
- input_shape: Optional[Tuple] = None,
+ input_shape: tuple | None = None,
seed: int = 0,
dtype: jnp.dtype = jnp.float32,
_do_init: bool = True,
@@ -87,7 +85,7 @@ def __init__(
module = self.module_class(config=config, dtype=dtype, **kwargs)
super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
- def init_weights(self, rng: jax.Array, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
+ def init_weights(self, rng: jax.Array, input_shape: tuple, params: FrozenDict = None) -> FrozenDict:
# init input tensor
clip_input = jax.random.normal(rng, input_shape)
diff --git a/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py b/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py
index ffd66792fe46..ba91a0f23923 100644
--- a/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py
+++ b/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Union
-
import torch
from torch import nn
@@ -41,8 +39,8 @@ def __init__(
def to(
self,
- torch_device: Optional[Union[str, torch.device]] = None,
- torch_dtype: Optional[torch.dtype] = None,
+ torch_device: str | torch.device | None = None,
+ torch_dtype: torch.dtype | None = None,
):
self.mean = nn.Parameter(self.mean.to(torch_device).to(torch_dtype))
self.std = nn.Parameter(self.std.to(torch_device).to(torch_dtype))
diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py
index 4655f446102a..43f801d81e2e 100644
--- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py
+++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,9 +12,9 @@ class StableDiffusion3PipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
index fcd108aef4c2..7764a79d7faf 100644
--- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
+++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import (
@@ -88,10 +88,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -106,15 +106,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -232,11 +232,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -288,10 +288,10 @@ def _get_t5_prompt_embeds(
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- clip_skip: Optional[int] = None,
+ device: torch.device | None = None,
+ clip_skip: int | None = None,
clip_model_index: int = 0,
):
device = device or self._execution_device
@@ -343,32 +343,32 @@ def _get_clip_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]],
- prompt_3: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ prompt_3: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ clip_skip: int | None = None,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -377,14 +377,14 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
prompt_embeds (`torch.FloatTensor`, *optional*):
@@ -715,9 +715,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch
# Adapted from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.prepare_ip_adapter_image_embeds
def prepare_ip_adapter_image_embeds(
self,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
) -> torch.Tensor:
@@ -773,50 +773,50 @@ def enable_sequential_cpu_offload(self, *args, **kwargs):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 28,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
- skip_guidance_layers: List[int] = None,
+ skip_guidance_layers: list[int] = None,
skip_layer_guidance_scale: float = 2.8,
skip_layer_guidance_stop: float = 0.2,
skip_layer_guidance_start: float = 0.01,
- mu: Optional[float] = None,
+ mu: float | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
will be used instead
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -826,7 +826,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -836,19 +836,19 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used instead
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used instead
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -890,12 +890,12 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`.
- skip_guidance_layers (`List[int]`, *optional*):
+ skip_guidance_layers (`list[int]`, *optional*):
A list of integers that specify layers to skip during guidance. If not provided, all layers will be
used for guidance. If provided, the guidance will only be applied to the layers specified in the list.
Recommended value by StabiltyAI for Stable Diffusion 3.5 Medium is [7, 8, 9].
diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py
index e6ddbb5544c7..7951b970cd0c 100644
--- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py
+++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -93,7 +93,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -108,10 +108,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -126,15 +126,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -218,8 +218,8 @@ def __init__(
tokenizer_2: CLIPTokenizer,
text_encoder_3: T5EncoderModel,
tokenizer_3: T5TokenizerFast,
- image_encoder: Optional[SiglipVisionModel] = None,
- feature_extractor: Optional[SiglipImageProcessor] = None,
+ image_encoder: SiglipVisionModel | None = None,
+ feature_extractor: SiglipImageProcessor | None = None,
):
super().__init__()
@@ -256,11 +256,11 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -313,10 +313,10 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- clip_skip: Optional[int] = None,
+ device: torch.device | None = None,
+ clip_skip: int | None = None,
clip_model_index: int = 0,
):
device = device or self._execution_device
@@ -369,32 +369,32 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]],
- prompt_3: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ prompt_3: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ clip_skip: int | None = None,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -403,14 +403,14 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
prompt_embeds (`torch.FloatTensor`, *optional*):
@@ -771,9 +771,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_ip_adapter_image_embeds
def prepare_ip_adapter_image_embeds(
self,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
) -> torch.Tensor:
@@ -830,48 +830,48 @@ def enable_sequential_cpu_offload(self, *args, **kwargs):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
image: PipelineImageInput = None,
strength: float = 0.6,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- 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,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
- mu: Optional[float] = None,
+ mu: float | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
will be used instead
height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor):
@@ -881,7 +881,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -891,19 +891,19 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used instead
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used instead
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -945,7 +945,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py
index b1b30efc7da3..d3594b868f89 100644
--- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py
+++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import (
@@ -92,7 +92,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -107,10 +107,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -125,15 +125,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -217,8 +217,8 @@ def __init__(
tokenizer_2: CLIPTokenizer,
text_encoder_3: T5EncoderModel,
tokenizer_3: T5TokenizerFast,
- image_encoder: Optional[SiglipVisionModel] = None,
- feature_extractor: Optional[SiglipImageProcessor] = None,
+ image_encoder: SiglipVisionModel | None = None,
+ feature_extractor: SiglipImageProcessor | None = None,
):
super().__init__()
@@ -262,11 +262,11 @@ def __init__(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 256,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -319,10 +319,10 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
- clip_skip: Optional[int] = None,
+ device: torch.device | None = None,
+ clip_skip: int | None = None,
clip_model_index: int = 0,
):
device = device or self._execution_device
@@ -375,32 +375,32 @@ def _get_clip_prompt_embeds(
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- prompt_2: Union[str, List[str]],
- prompt_3: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ prompt_3: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+ clip_skip: int | None = None,
max_sequence_length: int = 256,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in all text-encoders
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
used in all text-encoders
device: (`torch.device`):
@@ -409,14 +409,14 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
prompt_embeds (`torch.FloatTensor`, *optional*):
@@ -863,9 +863,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch
# Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_ip_adapter_image_embeds
def prepare_ip_adapter_image_embeds(
self,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- device: Optional[torch.device] = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
) -> torch.Tensor:
@@ -922,67 +922,67 @@ def enable_sequential_cpu_offload(self, *args, **kwargs):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- prompt_3: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: PipelineImageInput = None,
height: int = None,
width: int = None,
- padding_mask_crop: Optional[int] = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.6,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 7.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- negative_prompt_3: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 256,
- mu: Optional[float] = None,
+ mu: float | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
will be used instead
- prompt_3 (`str` or `List[str]`, *optional*):
+ prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
will be used instead
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
1)`, or `(H, W)`.
- mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`):
+ mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`):
`Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask
latents tensor will be generated by `mask_image`.
height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor):
@@ -1005,7 +1005,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1015,19 +1015,19 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used instead
- negative_prompt_3 (`str` or `List[str]`, *optional*):
+ negative_prompt_3 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
`text_encoder_3`. If not defined, `negative_prompt` is used instead
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -1069,7 +1069,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py b/src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py
index a1ff99b6aa34..80b0c09bc9a5 100644
--- a/src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py
+++ b/src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py
@@ -14,7 +14,7 @@
import inspect
import math
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -111,7 +111,7 @@ def get_average_attention(self):
average_attention = self.attention_store
return average_attention
- def aggregate_attention(self, from_where: List[str]) -> torch.Tensor:
+ def aggregate_attention(self, from_where: list[str]) -> torch.Tensor:
"""Aggregates the attention across the different layers and heads at the specified resolution."""
out = []
attention_maps = self.get_average_attention()
@@ -267,9 +267,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -300,16 +300,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -317,7 +317,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -416,7 +416,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -619,8 +619,8 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@staticmethod
def _compute_max_attention_per_index(
attention_maps: torch.Tensor,
- indices: List[int],
- ) -> List[torch.Tensor]:
+ indices: list[int],
+ ) -> list[torch.Tensor]:
"""Computes the maximum attention value for each of the tokens we wish to alter."""
attention_for_text = attention_maps[:, :, 1:-1]
attention_for_text *= 100
@@ -641,7 +641,7 @@ def _compute_max_attention_per_index(
def _aggregate_and_get_max_attention_per_token(
self,
- indices: List[int],
+ indices: list[int],
):
"""Aggregates the attention for each token and computes the max activation value for each token to alter."""
attention_maps = self.attention_store.aggregate_attention(
@@ -654,7 +654,7 @@ def _aggregate_and_get_max_attention_per_token(
return max_attention_per_index
@staticmethod
- def _compute_loss(max_attention_per_index: List[torch.Tensor]) -> torch.Tensor:
+ def _compute_loss(max_attention_per_index: list[torch.Tensor]) -> torch.Tensor:
"""Computes the attend-and-excite loss using the maximum attention value for each token."""
losses = [max(0, 1.0 - curr_max) for curr_max in max_attention_per_index]
loss = max(losses)
@@ -670,7 +670,7 @@ def _update_latent(latents: torch.Tensor, loss: torch.Tensor, step_size: float)
def _perform_iterative_refinement_step(
self,
latents: torch.Tensor,
- indices: List[int],
+ indices: list[int],
loss: torch.Tensor,
threshold: float,
text_embeddings: torch.Tensor,
@@ -740,7 +740,7 @@ def register_attention_control(self):
self.unet.set_attn_processor(attn_procs)
self.attention_store.num_att_layers = cross_att_count
- def get_indices(self, prompt: str) -> Dict[str, int]:
+ def get_indices(self, prompt: str) -> dict[str, int]:
"""Utility function to list the indices of the tokens you wish to alte"""
ids = self.tokenizer(prompt).input_ids
indices = {i: tok for tok, i in zip(self.tokenizer.convert_ids_to_tokens(ids), range(len(ids)))}
@@ -750,37 +750,37 @@ def get_indices(self, prompt: str) -> Dict[str, int]:
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]],
- token_indices: Union[List[int], List[List[int]]],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str],
+ token_indices: list[int] | list[list[int]],
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
num_images_per_prompt: int = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
max_iter_to_alter: int = 25,
thresholds: dict = {0: 0.05, 10: 0.5, 20: 0.8},
scale_factor: int = 20,
- attn_res: Optional[Tuple[int]] = (16, 16),
- clip_skip: Optional[int] = None,
+ attn_res: tuple[int] | None = (16, 16),
+ clip_skip: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- token_indices (`List[int]`):
+ token_indices (`list[int]`):
The token indices to alter with attend-and-excite.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -792,7 +792,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -800,7 +800,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py b/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py
index 65c25ffbe492..43bc2eb955c7 100644
--- a/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py
+++ b/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -63,14 +63,14 @@ class DiffEditInversionPipelineOutput(BaseOutput):
Args:
latents (`torch.Tensor`)
inverted latents tensor
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `num_timesteps * batch_size` or numpy array of shape `(num_timesteps,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `num_timesteps * batch_size` or numpy array of shape `(num_timesteps,
batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the
diffusion pipeline.
"""
latents: torch.Tensor
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
EXAMPLE_DOC_STRING = """
@@ -393,9 +393,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -426,16 +426,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -443,7 +443,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -542,7 +542,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -844,23 +844,23 @@ def get_epsilon(self, model_output: torch.Tensor, sample: torch.Tensor, timestep
@replace_example_docstring(EXAMPLE_DOC_STRING)
def generate_mask(
self,
- image: Union[torch.Tensor, PIL.Image.Image] = None,
- target_prompt: Optional[Union[str, List[str]]] = None,
- target_negative_prompt: Optional[Union[str, List[str]]] = None,
- target_prompt_embeds: Optional[torch.Tensor] = None,
- target_negative_prompt_embeds: Optional[torch.Tensor] = None,
- source_prompt: Optional[Union[str, List[str]]] = None,
- source_negative_prompt: Optional[Union[str, List[str]]] = None,
- source_prompt_embeds: Optional[torch.Tensor] = None,
- source_negative_prompt_embeds: Optional[torch.Tensor] = None,
- num_maps_per_mask: Optional[int] = 10,
- mask_encode_strength: Optional[float] = 0.5,
- mask_thresholding_ratio: Optional[float] = 3.0,
+ image: torch.Tensor | PIL.Image.Image = None,
+ target_prompt: str | list[str] | None = None,
+ target_negative_prompt: str | list[str] | None = None,
+ target_prompt_embeds: torch.Tensor | None = None,
+ target_negative_prompt_embeds: torch.Tensor | None = None,
+ source_prompt: str | list[str] | None = None,
+ source_negative_prompt: str | list[str] | None = None,
+ source_prompt_embeds: torch.Tensor | None = None,
+ source_negative_prompt_embeds: torch.Tensor | None = None,
+ num_maps_per_mask: int | None = 10,
+ mask_encode_strength: float | None = 0.5,
+ mask_thresholding_ratio: float | None = 3.0,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- output_type: Optional[str] = "np",
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ output_type: str | None = "np",
+ cross_attention_kwargs: dict[str, Any] | None = None,
):
r"""
Generate a latent mask given a mask prompt, a target prompt, and an image.
@@ -868,10 +868,10 @@ def generate_mask(
Args:
image (`PIL.Image.Image`):
`Image` or tensor representing an image batch to be used for computing the mask.
- target_prompt (`str` or `List[str]`, *optional*):
+ target_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide semantic mask generation. If not defined, you need to pass
`prompt_embeds`.
- target_negative_prompt (`str` or `List[str]`, *optional*):
+ target_negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
target_prompt_embeds (`torch.Tensor`, *optional*):
@@ -880,10 +880,10 @@ def generate_mask(
target_negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
- source_prompt (`str` or `List[str]`, *optional*):
+ source_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide semantic mask generation using DiffEdit. If not defined, you need to
pass `source_prompt_embeds` or `source_image` instead.
- source_negative_prompt (`str` or `List[str]`, *optional*):
+ source_negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide semantic mask generation away from using DiffEdit. If not defined, you
need to pass `source_negative_prompt_embeds` or `source_image` instead.
source_prompt_embeds (`torch.Tensor`, *optional*):
@@ -908,7 +908,7 @@ def generate_mask(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
output_type (`str`, *optional*, defaults to `"pil"`):
@@ -921,8 +921,8 @@ def generate_mask(
Examples:
Returns:
- `List[PIL.Image.Image]` or `np.array`:
- When returning a `List[PIL.Image.Image]`, the list consists of a batch of single-channel binary images
+ `list[PIL.Image.Image]` or `np.array`:
+ When returning a `list[PIL.Image.Image]`, the list consists of a batch of single-channel binary images
with dimensions `(height // self.vae_scale_factor, width // self.vae_scale_factor)`. If it's
`np.array`, the shape is `(batch_size, height // self.vae_scale_factor, width //
self.vae_scale_factor)`.
@@ -1063,21 +1063,21 @@ def generate_mask(
@replace_example_docstring(EXAMPLE_INVERT_DOC_STRING)
def invert(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- image: Union[torch.Tensor, PIL.Image.Image] = None,
+ prompt: str | list[str] | None = None,
+ image: torch.Tensor | PIL.Image.Image = None,
num_inference_steps: int = 50,
inpaint_strength: float = 0.8,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
decode_latents: bool = False,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
+ cross_attention_kwargs: dict[str, Any] | None = None,
lambda_auto_corr: float = 20.0,
lambda_kl: float = 20.0,
num_reg_steps: int = 0,
@@ -1087,7 +1087,7 @@ def invert(
Generate inverted latents given a prompt and image.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
image (`PIL.Image.Image`):
`Image` or tensor representing an image batch to produce the inverted latents guided by `prompt`.
@@ -1102,7 +1102,7 @@ def invert(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
generator (`torch.Generator`, *optional*):
@@ -1301,31 +1301,31 @@ def invert(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- mask_image: Union[torch.Tensor, PIL.Image.Image] = None,
- image_latents: Union[torch.Tensor, PIL.Image.Image] = None,
- inpaint_strength: Optional[float] = 0.8,
+ prompt: str | list[str] | None = None,
+ mask_image: torch.Tensor | PIL.Image.Image = None,
+ image_latents: torch.Tensor | PIL.Image.Image = None,
+ inpaint_strength: float | None = 0.8,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
clip_skip: int = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
mask_image (`PIL.Image.Image`):
`Image` or tensor representing an image batch to mask the generated image. White pixels in the mask are
@@ -1345,7 +1345,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
diff --git a/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py b/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py
index 78b026684cfa..fa5bc9376e53 100644
--- a/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py
+++ b/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py
@@ -14,7 +14,7 @@
import inspect
import warnings
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -191,9 +191,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -224,16 +224,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -241,7 +241,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -340,7 +340,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -540,34 +540,34 @@ def target_size_center_crop(self, im, new_hw):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
gligen_scheduled_sampling_beta: float = 0.3,
- gligen_phrases: List[str] = None,
- gligen_boxes: List[List[float]] = None,
- gligen_inpaint_image: Optional[PIL.Image.Image] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ gligen_phrases: list[str] = None,
+ gligen_boxes: list[list[float]] = None,
+ gligen_inpaint_image: PIL.Image.Image | None = None,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -579,13 +579,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- gligen_phrases (`List[str]`):
+ gligen_phrases (`list[str]`):
The phrases to guide what to include in each of the regions defined by the corresponding
`gligen_boxes`. There should only be one phrase per bounding box.
- gligen_boxes (`List[List[float]]`):
+ gligen_boxes (`list[list[float]]`):
The bounding boxes that identify rectangular regions of the image that are going to be filled with the
content described by the corresponding `gligen_phrases`. Each rectangular box is defined as a
- `List[float]` of 4 elements `[xmin, ymin, xmax, ymax]` where each value is between [0,1].
+ `list[float]` of 4 elements `[xmin, ymin, xmax, ymax]` where each value is between [0,1].
gligen_inpaint_image (`PIL.Image.Image`, *optional*):
The input image, if provided, is inpainted with objects described by the `gligen_boxes` and
`gligen_phrases`. Otherwise, it is treated as a generation task on a blank input image.
@@ -593,7 +593,7 @@ def __call__(
Scheduled Sampling factor from [GLIGEN: Open-Set Grounded Text-to-Image
Generation](https://huggingface.co/papers/2301.07093). Scheduled Sampling factor is only varied for
scheduled sampling during inference for improved quality and controllability.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -601,7 +601,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py b/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py
index 05cbad139d92..62e8a9fa95ae 100644
--- a/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py
+++ b/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py
@@ -14,7 +14,7 @@
import inspect
import warnings
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -255,16 +255,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -272,7 +272,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -371,7 +371,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -715,30 +715,30 @@ def get_cross_attention_kwargs_without_grounded(self, hidden_size, repeat_batch,
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
gligen_scheduled_sampling_beta: float = 0.3,
- gligen_phrases: List[str] = None,
- gligen_images: List[PIL.Image.Image] = None,
- input_phrases_mask: Union[int, List[int]] = None,
- input_images_mask: Union[int, List[int]] = None,
- gligen_boxes: List[List[float]] = None,
- gligen_inpaint_image: Optional[PIL.Image.Image] = None,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ gligen_phrases: list[str] = None,
+ gligen_images: list[PIL.Image.Image] = None,
+ input_phrases_mask: int | list[int] = None,
+ input_images_mask: int | list[int] = None,
+ gligen_boxes: list[list[float]] = None,
+ gligen_inpaint_image: PIL.Image.Image | None = None,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
gligen_normalize_constant: float = 28.7,
clip_skip: int = None,
):
@@ -746,7 +746,7 @@ def __call__(
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -758,20 +758,20 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- gligen_phrases (`List[str]`):
+ gligen_phrases (`list[str]`):
The phrases to guide what to include in each of the regions defined by the corresponding
`gligen_boxes`. There should only be one phrase per bounding box.
- gligen_images (`List[PIL.Image.Image]`):
+ gligen_images (`list[PIL.Image.Image]`):
The images to guide what to include in each of the regions defined by the corresponding `gligen_boxes`.
There should only be one image per bounding box
- input_phrases_mask (`int` or `List[int]`):
+ input_phrases_mask (`int` or `list[int]`):
pre phrases mask input defined by the correspongding `input_phrases_mask`
- input_images_mask (`int` or `List[int]`):
+ input_images_mask (`int` or `list[int]`):
pre images mask input defined by the correspongding `input_images_mask`
- gligen_boxes (`List[List[float]]`):
+ gligen_boxes (`list[list[float]]`):
The bounding boxes that identify rectangular regions of the image that are going to be filled with the
content described by the corresponding `gligen_phrases`. Each rectangular box is defined as a
- `List[float]` of 4 elements `[xmin, ymin, xmax, ymax]` where each value is between [0,1].
+ `list[float]` of 4 elements `[xmin, ymin, xmax, ymax]` where each value is between [0,1].
gligen_inpaint_image (`PIL.Image.Image`, *optional*):
The input image, if provided, is inpainted with objects described by the `gligen_boxes` and
`gligen_phrases`. Otherwise, it is treated as a generation task on a blank input image.
@@ -779,7 +779,7 @@ def __call__(
Scheduled Sampling factor from [GLIGEN: Open-Set Grounded Text-to-Image
Generation](https://huggingface.co/papers/2301.07093). Scheduled Sampling factor is only varied for
scheduled sampling during inference for improved quality and controllability.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -787,7 +787,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py b/src/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py
deleted file mode 100644
index 7eb5bf8c229b..000000000000
--- a/src/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py
+++ /dev/null
@@ -1,62 +0,0 @@
-from typing import TYPE_CHECKING
-
-from ...utils import (
- DIFFUSERS_SLOW_IMPORT,
- OptionalDependencyNotAvailable,
- _LazyModule,
- get_objects_from_module,
- is_k_diffusion_available,
- is_k_diffusion_version,
- is_torch_available,
- is_transformers_available,
-)
-
-
-_dummy_objects = {}
-_import_structure = {}
-
-
-try:
- if not (
- is_transformers_available()
- and is_torch_available()
- and is_k_diffusion_available()
- and is_k_diffusion_version(">=", "0.0.12")
- ):
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- from ...utils import dummy_torch_and_transformers_and_k_diffusion_objects # noqa F403
-
- _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_k_diffusion_objects))
-else:
- _import_structure["pipeline_stable_diffusion_k_diffusion"] = ["StableDiffusionKDiffusionPipeline"]
- _import_structure["pipeline_stable_diffusion_xl_k_diffusion"] = ["StableDiffusionXLKDiffusionPipeline"]
-
-if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
- try:
- if not (
- is_transformers_available()
- and is_torch_available()
- and is_k_diffusion_available()
- and is_k_diffusion_version(">=", "0.0.12")
- ):
- raise OptionalDependencyNotAvailable()
-
- except OptionalDependencyNotAvailable:
- from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import *
- else:
- from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline
- from .pipeline_stable_diffusion_xl_k_diffusion import StableDiffusionXLKDiffusionPipeline
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(
- __name__,
- globals()["__file__"],
- _import_structure,
- module_spec=__spec__,
- )
-
- for name, value in _dummy_objects.items():
- setattr(sys.modules[__name__], name, value)
diff --git a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py b/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py
deleted file mode 100755
index feebd6adf8f8..000000000000
--- a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py
+++ /dev/null
@@ -1,689 +0,0 @@
-# Copyright 2025 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import importlib
-import inspect
-from typing import Callable, List, Optional, Union
-
-import torch
-from k_diffusion.external import CompVisDenoiser, CompVisVDenoiser
-from k_diffusion.sampling import BrownianTreeNoiseSampler, get_sigmas_karras
-from transformers import (
- CLIPImageProcessor,
- CLIPTextModel,
- CLIPTokenizer,
- CLIPTokenizerFast,
-)
-
-from ...image_processor import VaeImageProcessor
-from ...loaders import (
- StableDiffusionLoraLoaderMixin,
- TextualInversionLoaderMixin,
-)
-from ...models import AutoencoderKL, UNet2DConditionModel
-from ...models.lora import adjust_lora_scale_text_encoder
-from ...schedulers import KarrasDiffusionSchedulers, LMSDiscreteScheduler
-from ...utils import (
- USE_PEFT_BACKEND,
- deprecate,
- logging,
- scale_lora_layers,
- unscale_lora_layers,
-)
-from ...utils.torch_utils import randn_tensor
-from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin
-from ..stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-
-class ModelWrapper:
- def __init__(self, model, alphas_cumprod):
- self.model = model
- self.alphas_cumprod = alphas_cumprod
-
- def apply_model(self, *args, **kwargs):
- if len(args) == 3:
- encoder_hidden_states = args[-1]
- args = args[:2]
- if kwargs.get("cond", None) is not None:
- encoder_hidden_states = kwargs.pop("cond")
- return self.model(*args, encoder_hidden_states=encoder_hidden_states, **kwargs).sample
-
-
-class StableDiffusionKDiffusionPipeline(
- DeprecatedPipelineMixin,
- DiffusionPipeline,
- StableDiffusionMixin,
- TextualInversionLoaderMixin,
- StableDiffusionLoraLoaderMixin,
-):
- r"""
- Pipeline for text-to-image generation using Stable Diffusion.
-
- This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
- library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
-
- The pipeline also inherits the following loading methods:
- - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
- - [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights
- - [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for saving LoRA weights
-
- > [!WARNING] > This is an experimental pipeline and is likely to change in the future.
-
- Args:
- vae ([`AutoencoderKL`]):
- Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
- text_encoder ([`CLIPTextModel`]):
- Frozen text-encoder. Stable Diffusion uses the text portion of
- [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
- the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
- tokenizer (`CLIPTokenizer`):
- Tokenizer of class
- [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
- unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
- scheduler ([`SchedulerMixin`]):
- A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
- [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
- safety_checker ([`StableDiffusionSafetyChecker`]):
- Classification module that estimates whether generated images could be considered offensive or harmful.
- Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for
- details.
- feature_extractor ([`CLIPImageProcessor`]):
- Model that extracts features from generated images to be used as inputs for the `safety_checker`.
- """
-
- _last_supported_version = "0.33.1"
-
- model_cpu_offload_seq = "text_encoder->unet->vae"
- _optional_components = ["safety_checker", "feature_extractor"]
- _exclude_from_cpu_offload = ["safety_checker"]
-
- def __init__(
- self,
- vae: AutoencoderKL,
- text_encoder: CLIPTextModel,
- tokenizer: Union[CLIPTokenizer, CLIPTokenizerFast],
- unet: UNet2DConditionModel,
- scheduler: KarrasDiffusionSchedulers,
- safety_checker: StableDiffusionSafetyChecker,
- feature_extractor: CLIPImageProcessor,
- requires_safety_checker: bool = True,
- ):
- super().__init__()
-
- logger.info(
- f"{self.__class__} is an experimental pipeline and is likely to change in the future. We recommend to use"
- " this pipeline for fast experimentation / iteration if needed, but advice to rely on existing pipelines"
- " as defined in https://huggingface.co/docs/diffusers/api/schedulers#implemented-schedulers for"
- " production settings."
- )
-
- # get correct sigmas from LMS
- scheduler = LMSDiscreteScheduler.from_config(scheduler.config)
- self.register_modules(
- vae=vae,
- text_encoder=text_encoder,
- tokenizer=tokenizer,
- unet=unet,
- scheduler=scheduler,
- safety_checker=safety_checker,
- feature_extractor=feature_extractor,
- )
- self.register_to_config(requires_safety_checker=requires_safety_checker)
- self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
- self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
-
- model = ModelWrapper(unet, scheduler.alphas_cumprod)
- if scheduler.config.prediction_type == "v_prediction":
- self.k_diffusion_model = CompVisVDenoiser(model)
- else:
- self.k_diffusion_model = CompVisDenoiser(model)
-
- def set_scheduler(self, scheduler_type: str):
- library = importlib.import_module("k_diffusion")
- sampling = getattr(library, "sampling")
- try:
- self.sampler = getattr(sampling, scheduler_type)
- except Exception:
- valid_samplers = []
- for s in dir(sampling):
- if "sample_" in s:
- valid_samplers.append(s)
-
- raise ValueError(f"Invalid scheduler type {scheduler_type}. Please choose one of {valid_samplers}.")
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
- def _encode_prompt(
- self,
- prompt,
- device,
- num_images_per_prompt,
- do_classifier_free_guidance,
- negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- **kwargs,
- ):
- deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
- deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
-
- prompt_embeds_tuple = self.encode_prompt(
- prompt=prompt,
- device=device,
- num_images_per_prompt=num_images_per_prompt,
- do_classifier_free_guidance=do_classifier_free_guidance,
- negative_prompt=negative_prompt,
- prompt_embeds=prompt_embeds,
- negative_prompt_embeds=negative_prompt_embeds,
- lora_scale=lora_scale,
- **kwargs,
- )
-
- # concatenate for backwards comp
- prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])
-
- return prompt_embeds
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt
- def encode_prompt(
- self,
- prompt,
- device,
- num_images_per_prompt,
- do_classifier_free_guidance,
- negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
- ):
- r"""
- Encodes the prompt into text encoder hidden states.
-
- Args:
- prompt (`str` or `List[str]`, *optional*):
- prompt to be encoded
- device: (`torch.device`):
- torch device
- num_images_per_prompt (`int`):
- number of images that should be generated per prompt
- do_classifier_free_guidance (`bool`):
- whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation. If not defined, one has to pass
- `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
- less than `1`).
- prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
- provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
- argument.
- lora_scale (`float`, *optional*):
- A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
- 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.
- """
- # set lora scale so that monkey patched LoRA
- # function of text encoder can correctly access it
- if lora_scale is not None and isinstance(self, StableDiffusionLoraLoaderMixin):
- self._lora_scale = lora_scale
-
- # dynamically adjust the LoRA scale
- if not USE_PEFT_BACKEND:
- adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
- else:
- scale_lora_layers(self.text_encoder, lora_scale)
-
- 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]
-
- if prompt_embeds is None:
- # textual inversion: process multi-vector tokens if necessary
- if isinstance(self, TextualInversionLoaderMixin):
- prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
-
- text_inputs = self.tokenizer(
- prompt,
- padding="max_length",
- max_length=self.tokenizer.model_max_length,
- truncation=True,
- return_tensors="pt",
- )
- text_input_ids = text_inputs.input_ids
- untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
- if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
- text_input_ids, untruncated_ids
- ):
- removed_text = self.tokenizer.batch_decode(
- untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
- )
- logger.warning(
- "The following part of your input was truncated because CLIP can only handle sequences up to"
- f" {self.tokenizer.model_max_length} tokens: {removed_text}"
- )
-
- if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
- attention_mask = text_inputs.attention_mask.to(device)
- else:
- attention_mask = None
-
- if clip_skip is None:
- prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
- prompt_embeds = prompt_embeds[0]
- else:
- prompt_embeds = self.text_encoder(
- text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
- )
- # Access the `hidden_states` first, that contains a tuple of
- # all the hidden states from the encoder layers. Then index into
- # the tuple to access the hidden states from the desired layer.
- prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]
- # We also need to apply the final LayerNorm here to not mess with the
- # representations. The `last_hidden_states` that we typically use for
- # obtaining the final prompt representations passes through the LayerNorm
- # layer.
- prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)
-
- if self.text_encoder is not None:
- prompt_embeds_dtype = self.text_encoder.dtype
- elif self.unet is not None:
- prompt_embeds_dtype = self.unet.dtype
- else:
- prompt_embeds_dtype = prompt_embeds.dtype
-
- prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
- bs_embed, seq_len, _ = prompt_embeds.shape
- # duplicate text embeddings for each generation per prompt, using mps friendly method
- prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
- prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
- # get unconditional embeddings for classifier free guidance
- if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
- if negative_prompt is None:
- uncond_tokens = [""] * batch_size
- elif prompt is not None and type(prompt) is not type(negative_prompt):
- raise TypeError(
- f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
- f" {type(prompt)}."
- )
- elif isinstance(negative_prompt, str):
- uncond_tokens = [negative_prompt]
- elif batch_size != len(negative_prompt):
- raise ValueError(
- f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
- f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
- " the batch size of `prompt`."
- )
- else:
- uncond_tokens = negative_prompt
-
- # textual inversion: process multi-vector tokens if necessary
- if isinstance(self, TextualInversionLoaderMixin):
- uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
-
- max_length = prompt_embeds.shape[1]
- uncond_input = self.tokenizer(
- uncond_tokens,
- padding="max_length",
- max_length=max_length,
- truncation=True,
- return_tensors="pt",
- )
-
- if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
- attention_mask = uncond_input.attention_mask.to(device)
- else:
- attention_mask = None
-
- negative_prompt_embeds = self.text_encoder(
- uncond_input.input_ids.to(device),
- attention_mask=attention_mask,
- )
- negative_prompt_embeds = negative_prompt_embeds[0]
-
- if do_classifier_free_guidance:
- # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
- seq_len = negative_prompt_embeds.shape[1]
-
- negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
- negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
- negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
- if self.text_encoder is not None:
- if isinstance(self, StableDiffusionLoraLoaderMixin) and USE_PEFT_BACKEND:
- # Retrieve the original scale by scaling back the LoRA layers
- unscale_lora_layers(self.text_encoder, lora_scale)
-
- return prompt_embeds, negative_prompt_embeds
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
- def run_safety_checker(self, image, device, dtype):
- if self.safety_checker is None:
- has_nsfw_concept = None
- else:
- if torch.is_tensor(image):
- feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
- else:
- feature_extractor_input = self.image_processor.numpy_to_pil(image)
- safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
- image, has_nsfw_concept = self.safety_checker(
- images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
- )
- return image, has_nsfw_concept
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
- def decode_latents(self, latents):
- deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
- deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
-
- latents = 1 / self.vae.config.scaling_factor * latents
- image = self.vae.decode(latents, return_dict=False)[0]
- image = (image / 2 + 0.5).clamp(0, 1)
- # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
- image = image.cpu().permute(0, 2, 3, 1).float().numpy()
- return image
-
- def check_inputs(
- self,
- prompt,
- height,
- width,
- callback_steps,
- negative_prompt=None,
- prompt_embeds=None,
- negative_prompt_embeds=None,
- callback_on_step_end_tensor_inputs=None,
- ):
- if height % 8 != 0 or width % 8 != 0:
- raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
- if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
- raise ValueError(
- f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
- f" {type(callback_steps)}."
- )
- if callback_on_step_end_tensor_inputs is not None and not all(
- k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
- ):
- raise ValueError(
- f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
- )
-
- if prompt is not None and prompt_embeds is not None:
- raise ValueError(
- f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
- " only forward one of the two."
- )
- elif prompt is None and prompt_embeds is None:
- raise ValueError(
- "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
- )
- elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
- raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
- if negative_prompt is not None and negative_prompt_embeds is not None:
- raise ValueError(
- f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
- f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
- )
-
- if prompt_embeds is not None and negative_prompt_embeds is not None:
- if prompt_embeds.shape != negative_prompt_embeds.shape:
- raise ValueError(
- "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
- f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
- f" {negative_prompt_embeds.shape}."
- )
-
- def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
- shape = (
- batch_size,
- num_channels_latents,
- int(height) // self.vae_scale_factor,
- int(width) // self.vae_scale_factor,
- )
- if latents is None:
- latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
- else:
- if latents.shape != shape:
- raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
- latents = latents.to(device)
-
- # scale the initial noise by the standard deviation required by the scheduler
- return latents
-
- @torch.no_grad()
- def __call__(
- self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- num_inference_steps: int = 50,
- guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: int = 1,
- use_karras_sigmas: Optional[bool] = False,
- noise_sampler_seed: Optional[int] = None,
- clip_skip: int = None,
- ):
- r"""
- Function invoked when calling the pipeline for generation.
-
- Args:
- prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
- instead.
- height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
- The height in pixels of the generated image.
- width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
- The width in pixels of the generated image.
- num_inference_steps (`int`, *optional*, defaults to 50):
- The number of denoising steps. More denoising steps usually lead to a higher quality image at the
- expense of slower inference.
- guidance_scale (`float`, *optional*, defaults to 7.5):
- Guidance scale as defined in [Classifier-Free Diffusion
- Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
- of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
- `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
- the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation. If not defined, one has to pass
- `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale`
- is less than `1`).
- num_images_per_prompt (`int`, *optional*, defaults to 1):
- The number of images to generate per prompt.
- eta (`float`, *optional*, defaults to 0.0):
- Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
- applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator`, *optional*):
- One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
- to make generation deterministic.
- latents (`torch.Tensor`, *optional*):
- Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
- generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
- tensor will be generated by sampling using the supplied random `generator`.
- prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
- provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
- argument.
- output_type (`str`, *optional*, defaults to `"pil"`):
- The output format of the generate image. Choose between
- [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
- return_dict (`bool`, *optional*, defaults to `True`):
- Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
- plain tuple.
- callback (`Callable`, *optional*):
- A function that will be called every `callback_steps` steps during inference. The function will be
- called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.
- callback_steps (`int`, *optional*, defaults to 1):
- The frequency at which the `callback` function will be called. If not specified, the callback will be
- called at every step.
- use_karras_sigmas (`bool`, *optional*, defaults to `False`):
- Use karras sigmas. For example, specifying `sample_dpmpp_2m` to `set_scheduler` will be equivalent to
- `DPM++2M` in stable-diffusion-webui. On top of that, setting this option to True will make it `DPM++2M
- Karras`.
- noise_sampler_seed (`int`, *optional*, defaults to `None`):
- The random seed to use for the noise sampler. If `None`, a random seed will be generated.
- 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.
- Returns:
- [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
- [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
- When returning a tuple, the first element is a list with the generated images, and the second element is a
- list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
- (nsfw) content, according to the `safety_checker`.
- """
- # 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
-
- # 1. Check inputs. Raise error if not correct
- self.check_inputs(
- prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
- )
-
- # 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]
-
- device = self._execution_device
- # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
- # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`
- # corresponds to doing no classifier free guidance.
- do_classifier_free_guidance = True
- if guidance_scale <= 1.0:
- raise ValueError("has to use guidance_scale")
-
- # 3. Encode input prompt
- prompt_embeds, negative_prompt_embeds = self.encode_prompt(
- prompt,
- device,
- num_images_per_prompt,
- do_classifier_free_guidance,
- negative_prompt,
- prompt_embeds=prompt_embeds,
- negative_prompt_embeds=negative_prompt_embeds,
- clip_skip=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 do_classifier_free_guidance:
- prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
-
- # 4. Prepare timesteps
- self.scheduler.set_timesteps(num_inference_steps, device=prompt_embeds.device)
-
- # 5. Prepare sigmas
- if use_karras_sigmas:
- sigma_min: float = self.k_diffusion_model.sigmas[0].item()
- sigma_max: float = self.k_diffusion_model.sigmas[-1].item()
- sigmas = get_sigmas_karras(n=num_inference_steps, sigma_min=sigma_min, sigma_max=sigma_max)
- else:
- sigmas = self.scheduler.sigmas
- sigmas = sigmas.to(device)
- sigmas = sigmas.to(prompt_embeds.dtype)
-
- # 6. Prepare latent variables
- num_channels_latents = self.unet.config.in_channels
- latents = self.prepare_latents(
- batch_size * num_images_per_prompt,
- num_channels_latents,
- height,
- width,
- prompt_embeds.dtype,
- device,
- generator,
- latents,
- )
- latents = latents * sigmas[0]
- self.k_diffusion_model.sigmas = self.k_diffusion_model.sigmas.to(latents.device)
- self.k_diffusion_model.log_sigmas = self.k_diffusion_model.log_sigmas.to(latents.device)
-
- # 7. Define model function
- def model_fn(x, t):
- latent_model_input = torch.cat([x] * 2)
- t = torch.cat([t] * 2)
-
- noise_pred = self.k_diffusion_model(latent_model_input, t, cond=prompt_embeds)
-
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
- noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
- return noise_pred
-
- # 8. Run k-diffusion solver
- sampler_kwargs = {}
-
- if "noise_sampler" in inspect.signature(self.sampler).parameters:
- min_sigma, max_sigma = sigmas[sigmas > 0].min(), sigmas.max()
- noise_sampler = BrownianTreeNoiseSampler(latents, min_sigma, max_sigma, noise_sampler_seed)
- sampler_kwargs["noise_sampler"] = noise_sampler
-
- if "generator" in inspect.signature(self.sampler).parameters:
- sampler_kwargs["generator"] = generator
-
- latents = self.sampler(model_fn, latents, sigmas, **sampler_kwargs)
-
- if not output_type == "latent":
- image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[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)
-
- # 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)
diff --git a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py b/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py
deleted file mode 100644
index f9a8abfcc568..000000000000
--- a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py
+++ /dev/null
@@ -1,889 +0,0 @@
-# Copyright 2025 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import importlib
-import inspect
-from typing import List, Optional, Tuple, Union
-
-import torch
-from k_diffusion.external import CompVisDenoiser, CompVisVDenoiser
-from k_diffusion.sampling import BrownianTreeNoiseSampler, get_sigmas_karras
-from transformers import (
- CLIPTextModel,
- CLIPTextModelWithProjection,
- CLIPTokenizer,
-)
-
-from ...image_processor import VaeImageProcessor
-from ...loaders import (
- FromSingleFileMixin,
- IPAdapterMixin,
- StableDiffusionXLLoraLoaderMixin,
- TextualInversionLoaderMixin,
-)
-from ...models import AutoencoderKL, UNet2DConditionModel
-from ...models.lora import adjust_lora_scale_text_encoder
-from ...schedulers import KarrasDiffusionSchedulers, LMSDiscreteScheduler
-from ...utils import (
- USE_PEFT_BACKEND,
- deprecate,
- logging,
- replace_example_docstring,
- scale_lora_layers,
- unscale_lora_layers,
-)
-from ...utils.torch_utils import randn_tensor
-from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin
-from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-EXAMPLE_DOC_STRING = """
- Examples:
- ```py
- >>> import torch
- >>> from diffusers import StableDiffusionXLKDiffusionPipeline
-
- >>> pipe = StableDiffusionXLKDiffusionPipeline.from_pretrained(
- ... "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
- ... )
- >>> pipe = pipe.to("cuda")
- >>> pipe.set_scheduler("sample_dpmpp_2m_sde")
-
- >>> prompt = "a photo of an astronaut riding a horse on mars"
- >>> image = pipe(prompt).images[0]
- ```
-"""
-
-
-# Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.ModelWrapper
-class ModelWrapper:
- def __init__(self, model, alphas_cumprod):
- self.model = model
- self.alphas_cumprod = alphas_cumprod
-
- def apply_model(self, *args, **kwargs):
- if len(args) == 3:
- encoder_hidden_states = args[-1]
- args = args[:2]
- if kwargs.get("cond", None) is not None:
- encoder_hidden_states = kwargs.pop("cond")
- return self.model(*args, encoder_hidden_states=encoder_hidden_states, **kwargs).sample
-
-
-class StableDiffusionXLKDiffusionPipeline(
- DeprecatedPipelineMixin,
- DiffusionPipeline,
- StableDiffusionMixin,
- FromSingleFileMixin,
- StableDiffusionXLLoraLoaderMixin,
- TextualInversionLoaderMixin,
- IPAdapterMixin,
-):
- _last_supported_version = "0.33.1"
-
- r"""
- Pipeline for text-to-image generation using Stable Diffusion XL and k-diffusion.
-
- This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
- library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
-
- The pipeline also inherits the following loading methods:
- - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
- - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
- - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights
- - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights
- - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
-
- Args:
- vae ([`AutoencoderKL`]):
- Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
- text_encoder ([`CLIPTextModel`]):
- Frozen text-encoder. Stable Diffusion XL uses the text portion of
- [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
- the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
- text_encoder_2 ([` CLIPTextModelWithProjection`]):
- Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of
- [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
- specifically the
- [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
- variant.
- tokenizer (`CLIPTokenizer`):
- Tokenizer of class
- [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
- tokenizer_2 (`CLIPTokenizer`):
- Second Tokenizer of class
- [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
- unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
- scheduler ([`SchedulerMixin`]):
- A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
- [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
- force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`):
- Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of
- `stabilityai/stable-diffusion-xl-base-1-0`.
- """
-
- model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
- _optional_components = [
- "tokenizer",
- "tokenizer_2",
- "text_encoder",
- "text_encoder_2",
- "feature_extractor",
- ]
-
- def __init__(
- self,
- vae: AutoencoderKL,
- text_encoder: CLIPTextModel,
- text_encoder_2: CLIPTextModelWithProjection,
- tokenizer: CLIPTokenizer,
- tokenizer_2: CLIPTokenizer,
- unet: UNet2DConditionModel,
- scheduler: KarrasDiffusionSchedulers,
- force_zeros_for_empty_prompt: bool = True,
- ):
- super().__init__()
-
- # get correct sigmas from LMS
- scheduler = LMSDiscreteScheduler.from_config(scheduler.config)
- self.register_modules(
- vae=vae,
- text_encoder=text_encoder,
- text_encoder_2=text_encoder_2,
- tokenizer=tokenizer,
- tokenizer_2=tokenizer_2,
- unet=unet,
- scheduler=scheduler,
- )
- self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
- self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
- self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
-
- self.default_sample_size = (
- self.unet.config.sample_size
- if hasattr(self, "unet") and self.unet is not None and hasattr(self.unet.config, "sample_size")
- else 128
- )
-
- model = ModelWrapper(unet, scheduler.alphas_cumprod)
- if scheduler.config.prediction_type == "v_prediction":
- self.k_diffusion_model = CompVisVDenoiser(model)
- else:
- self.k_diffusion_model = CompVisDenoiser(model)
-
- # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.set_scheduler
- def set_scheduler(self, scheduler_type: str):
- library = importlib.import_module("k_diffusion")
- sampling = getattr(library, "sampling")
- try:
- self.sampler = getattr(sampling, scheduler_type)
- except Exception:
- valid_samplers = []
- for s in dir(sampling):
- if "sample_" in s:
- valid_samplers.append(s)
-
- raise ValueError(f"Invalid scheduler type {scheduler_type}. Please choose one of {valid_samplers}.")
-
- # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt
- def encode_prompt(
- self,
- prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
- num_images_per_prompt: int = 1,
- do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
- ):
- r"""
- Encodes the prompt into text encoder hidden states.
-
- Args:
- prompt (`str` or `List[str]`, *optional*):
- prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
- The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
- used in both text-encoders
- device: (`torch.device`):
- torch device
- num_images_per_prompt (`int`):
- number of images that should be generated per prompt
- do_classifier_free_guidance (`bool`):
- whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation. If not defined, one has to pass
- `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
- less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
- `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
- prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
- provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
- argument.
- pooled_prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
- If not provided, pooled text embeddings will be generated from `prompt` input argument.
- negative_pooled_prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
- input argument.
- lora_scale (`float`, *optional*):
- A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
- 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.
- """
- device = device or self._execution_device
-
- # set lora scale so that monkey patched LoRA
- # function of text encoder can correctly access it
- if lora_scale is not None and isinstance(self, StableDiffusionXLLoraLoaderMixin):
- self._lora_scale = lora_scale
-
- # dynamically adjust the LoRA scale
- if self.text_encoder is not None:
- if not USE_PEFT_BACKEND:
- adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
- else:
- scale_lora_layers(self.text_encoder, lora_scale)
-
- if self.text_encoder_2 is not None:
- if not USE_PEFT_BACKEND:
- adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale)
- else:
- scale_lora_layers(self.text_encoder_2, lora_scale)
-
- prompt = [prompt] if isinstance(prompt, str) else prompt
-
- if prompt is not None:
- batch_size = len(prompt)
- else:
- batch_size = prompt_embeds.shape[0]
-
- # Define tokenizers and text encoders
- tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
- text_encoders = (
- [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
- )
-
- if prompt_embeds is None:
- prompt_2 = prompt_2 or prompt
- prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
-
- # textual inversion: process multi-vector tokens if necessary
- prompt_embeds_list = []
- prompts = [prompt, prompt_2]
- for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
- if isinstance(self, TextualInversionLoaderMixin):
- prompt = self.maybe_convert_prompt(prompt, tokenizer)
-
- text_inputs = tokenizer(
- prompt,
- padding="max_length",
- max_length=tokenizer.model_max_length,
- truncation=True,
- return_tensors="pt",
- )
-
- text_input_ids = text_inputs.input_ids
- untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
- if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
- text_input_ids, untruncated_ids
- ):
- removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
- logger.warning(
- "The following part of your input was truncated because CLIP can only handle sequences up to"
- f" {tokenizer.model_max_length} tokens: {removed_text}"
- )
-
- prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
-
- # We are only ALWAYS interested in the pooled output of the final text encoder
- if pooled_prompt_embeds is None and prompt_embeds[0].ndim == 2:
- pooled_prompt_embeds = prompt_embeds[0]
-
- if clip_skip is None:
- prompt_embeds = prompt_embeds.hidden_states[-2]
- else:
- # "2" because SDXL always indexes from the penultimate layer.
- prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
-
- prompt_embeds_list.append(prompt_embeds)
-
- prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
-
- # get unconditional embeddings for classifier free guidance
- zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
- if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
- negative_prompt_embeds = torch.zeros_like(prompt_embeds)
- negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
- elif do_classifier_free_guidance and negative_prompt_embeds is None:
- negative_prompt = negative_prompt or ""
- negative_prompt_2 = negative_prompt_2 or negative_prompt
-
- # normalize str to list
- negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
- negative_prompt_2 = (
- batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
- )
-
- uncond_tokens: List[str]
- if prompt is not None and type(prompt) is not type(negative_prompt):
- raise TypeError(
- f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
- f" {type(prompt)}."
- )
- elif batch_size != len(negative_prompt):
- raise ValueError(
- f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
- f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
- " the batch size of `prompt`."
- )
- else:
- uncond_tokens = [negative_prompt, negative_prompt_2]
-
- negative_prompt_embeds_list = []
- for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
- if isinstance(self, TextualInversionLoaderMixin):
- negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
-
- max_length = prompt_embeds.shape[1]
- uncond_input = tokenizer(
- negative_prompt,
- padding="max_length",
- max_length=max_length,
- truncation=True,
- return_tensors="pt",
- )
-
- negative_prompt_embeds = text_encoder(
- uncond_input.input_ids.to(device),
- output_hidden_states=True,
- )
-
- # We are only ALWAYS interested in the pooled output of the final text encoder
- if negative_pooled_prompt_embeds is None and negative_prompt_embeds[0].ndim == 2:
- negative_pooled_prompt_embeds = negative_prompt_embeds[0]
- negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
-
- negative_prompt_embeds_list.append(negative_prompt_embeds)
-
- negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
-
- if self.text_encoder_2 is not None:
- prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
- else:
- prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device)
-
- bs_embed, seq_len, _ = prompt_embeds.shape
- # duplicate text embeddings for each generation per prompt, using mps friendly method
- prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
- prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
- if do_classifier_free_guidance:
- # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
- seq_len = negative_prompt_embeds.shape[1]
-
- if self.text_encoder_2 is not None:
- negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
- else:
- negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device)
-
- negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
- negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
- pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
- bs_embed * num_images_per_prompt, -1
- )
- if do_classifier_free_guidance:
- negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
- bs_embed * num_images_per_prompt, -1
- )
-
- if self.text_encoder is not None:
- if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
- # Retrieve the original scale by scaling back the LoRA layers
- unscale_lora_layers(self.text_encoder, lora_scale)
-
- if self.text_encoder_2 is not None:
- if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
- # Retrieve the original scale by scaling back the LoRA layers
- unscale_lora_layers(self.text_encoder_2, lora_scale)
-
- return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
-
- def check_inputs(
- self,
- prompt,
- prompt_2,
- height,
- width,
- negative_prompt=None,
- negative_prompt_2=None,
- prompt_embeds=None,
- negative_prompt_embeds=None,
- pooled_prompt_embeds=None,
- negative_pooled_prompt_embeds=None,
- ):
- if height % 8 != 0 or width % 8 != 0:
- raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
- if prompt is not None and prompt_embeds is not None:
- raise ValueError(
- f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
- " only forward one of the two."
- )
- elif prompt_2 is not None and prompt_embeds is not None:
- raise ValueError(
- f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
- " only forward one of the two."
- )
- elif prompt is None and prompt_embeds is None:
- raise ValueError(
- "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
- )
- elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
- raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
- elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
- raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
-
- if negative_prompt is not None and negative_prompt_embeds is not None:
- raise ValueError(
- f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
- f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
- )
- elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
- raise ValueError(
- f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
- f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
- )
-
- if prompt_embeds is not None and negative_prompt_embeds is not None:
- if prompt_embeds.shape != negative_prompt_embeds.shape:
- raise ValueError(
- "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
- f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
- f" {negative_prompt_embeds.shape}."
- )
-
- if prompt_embeds is not None and pooled_prompt_embeds is None:
- raise ValueError(
- "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
- )
-
- if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
- raise ValueError(
- "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
- )
-
- def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
- shape = (
- batch_size,
- num_channels_latents,
- int(height) // self.vae_scale_factor,
- int(width) // self.vae_scale_factor,
- )
- if isinstance(generator, list) and len(generator) != batch_size:
- raise ValueError(
- f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
- f" size of {batch_size}. Make sure the batch size matches the length of the generators."
- )
- if latents is None:
- latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
- else:
- latents = latents.to(device)
-
- return latents
-
- # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._get_add_time_ids
- def _get_add_time_ids(
- self, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None
- ):
- add_time_ids = list(original_size + crops_coords_top_left + target_size)
-
- passed_add_embed_dim = (
- self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
- )
- expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
-
- if expected_add_embed_dim != passed_add_embed_dim:
- raise ValueError(
- f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
- )
-
- add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
- return add_time_ids
-
- # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.upcast_vae
- def upcast_vae(self):
- deprecate(
- "upcast_vae",
- "1.0.0",
- "`upcast_vae` is deprecated. Please use `pipe.vae.to(torch.float32)`. For more details, please refer to: https://github.com/huggingface/diffusers/pull/12619#issue-3606633695.",
- )
- self.vae.to(dtype=torch.float32)
-
- @property
- def guidance_scale(self):
- return self._guidance_scale
-
- @property
- def clip_skip(self):
- return self._clip_skip
-
- # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
- # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`
- # corresponds to doing no classifier free guidance.
- @property
- def do_classifier_free_guidance(self):
- return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
-
- @torch.no_grad()
- @replace_example_docstring(EXAMPLE_DOC_STRING)
- def __call__(
- self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- num_inference_steps: int = 50,
- guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- return_dict: bool = True,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- use_karras_sigmas: Optional[bool] = False,
- noise_sampler_seed: Optional[int] = None,
- clip_skip: Optional[int] = None,
- ):
- r"""
- Function invoked when calling the pipeline for generation.
-
- Args:
- prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
- instead.
- prompt_2 (`str` or `List[str]`, *optional*):
- The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
- used in both text-encoders
- height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
- The height in pixels of the generated image. This is set to 1024 by default for the best results.
- Anything below 512 pixels won't work well for
- [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
- and checkpoints that are not specifically fine-tuned on low resolutions.
- width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
- The width in pixels of the generated image. This is set to 1024 by default for the best results.
- Anything below 512 pixels won't work well for
- [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
- and checkpoints that are not specifically fine-tuned on low resolutions.
- num_inference_steps (`int`, *optional*, defaults to 50):
- The number of denoising steps. More denoising steps usually lead to a higher quality image at the
- expense of slower inference.
- guidance_scale (`float`, *optional*, defaults to 5.0):
- Guidance scale as defined in [Classifier-Free Diffusion
- Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
- of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
- `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
- the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation. If not defined, one has to pass
- `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
- less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
- `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
- num_images_per_prompt (`int`, *optional*, defaults to 1):
- The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
- One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
- to make generation deterministic.
- latents (`torch.Tensor`, *optional*):
- Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
- generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
- tensor will be generated by sampling using the supplied random `generator`.
- prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
- provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
- argument.
- pooled_prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
- If not provided, pooled text embeddings will be generated from `prompt` input argument.
- negative_pooled_prompt_embeds (`torch.Tensor`, *optional*):
- Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
- input argument.
- output_type (`str`, *optional*, defaults to `"pil"`):
- The output format of the generate image. Choose between
- [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
- return_dict (`bool`, *optional*, defaults to `True`):
- Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
- of a plain tuple.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
- If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
- `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
- explained in section 2.2 of
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
- `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
- `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
- `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
- For most cases, `target_size` should be set to the desired height and width of the generated image. If
- not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
- section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
- To negatively condition the generation process based on a specific image resolution. Part of SDXL's
- micro-conditioning as explained in section 2.2 of
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
- information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
- To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
- micro-conditioning as explained in section 2.2 of
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
- information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
- To negatively condition the generation process based on a target image resolution. It should be as same
- as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
- information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
-
- Examples:
-
- Returns:
- [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
- [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
- `tuple`. When returning a tuple, the first element is a list with the generated images.
- """
-
- # 0. Default height and width to unet
- height = height or self.default_sample_size * self.vae_scale_factor
- width = width or self.default_sample_size * self.vae_scale_factor
-
- original_size = original_size or (height, width)
- target_size = target_size or (height, width)
-
- # 1. Check inputs. Raise error if not correct
- self.check_inputs(
- prompt,
- prompt_2,
- height,
- width,
- negative_prompt,
- negative_prompt_2,
- prompt_embeds,
- negative_prompt_embeds,
- pooled_prompt_embeds,
- negative_pooled_prompt_embeds,
- )
-
- if guidance_scale <= 1.0:
- raise ValueError("has to use guidance_scale")
-
- self._guidance_scale = guidance_scale
- self._clip_skip = clip_skip
-
- # 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]
-
- device = self._execution_device
-
- # 3. Encode input prompt
- lora_scale = None
-
- (
- prompt_embeds,
- negative_prompt_embeds,
- pooled_prompt_embeds,
- negative_pooled_prompt_embeds,
- ) = self.encode_prompt(
- prompt=prompt,
- prompt_2=prompt_2,
- device=device,
- num_images_per_prompt=num_images_per_prompt,
- do_classifier_free_guidance=self.do_classifier_free_guidance,
- negative_prompt=negative_prompt,
- negative_prompt_2=negative_prompt_2,
- prompt_embeds=prompt_embeds,
- negative_prompt_embeds=negative_prompt_embeds,
- pooled_prompt_embeds=pooled_prompt_embeds,
- negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
- lora_scale=lora_scale,
- clip_skip=self.clip_skip,
- )
-
- # 4. Prepare timesteps
- self.scheduler.set_timesteps(num_inference_steps, device=prompt_embeds.device)
-
- # 5. Prepare sigmas
- if use_karras_sigmas:
- sigma_min: float = self.k_diffusion_model.sigmas[0].item()
- sigma_max: float = self.k_diffusion_model.sigmas[-1].item()
- sigmas = get_sigmas_karras(n=num_inference_steps, sigma_min=sigma_min, sigma_max=sigma_max)
- else:
- sigmas = self.scheduler.sigmas
- sigmas = sigmas.to(dtype=prompt_embeds.dtype, device=device)
-
- # 6. Prepare latent variables
- num_channels_latents = self.unet.config.in_channels
- latents = self.prepare_latents(
- batch_size * num_images_per_prompt,
- num_channels_latents,
- height,
- width,
- prompt_embeds.dtype,
- device,
- generator,
- latents,
- )
- latents = latents * sigmas[0]
-
- self.k_diffusion_model.sigmas = self.k_diffusion_model.sigmas.to(latents.device)
- self.k_diffusion_model.log_sigmas = self.k_diffusion_model.log_sigmas.to(latents.device)
-
- # 7. Prepare added time ids & embeddings
- add_text_embeds = pooled_prompt_embeds
- if self.text_encoder_2 is None:
- text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
- else:
- text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
-
- add_time_ids = self._get_add_time_ids(
- original_size,
- crops_coords_top_left,
- target_size,
- dtype=prompt_embeds.dtype,
- text_encoder_projection_dim=text_encoder_projection_dim,
- )
- if negative_original_size is not None and negative_target_size is not None:
- negative_add_time_ids = self._get_add_time_ids(
- negative_original_size,
- negative_crops_coords_top_left,
- negative_target_size,
- dtype=prompt_embeds.dtype,
- text_encoder_projection_dim=text_encoder_projection_dim,
- )
- else:
- negative_add_time_ids = add_time_ids
-
- if self.do_classifier_free_guidance:
- prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
- add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
- add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
-
- prompt_embeds = prompt_embeds.to(device)
- add_text_embeds = add_text_embeds.to(device)
- add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
-
- added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
-
- # 8. 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)
-
- # 9. Define model function
- def model_fn(x, t):
- latent_model_input = torch.cat([x] * 2)
- t = torch.cat([t] * 2)
-
- noise_pred = self.k_diffusion_model(
- latent_model_input,
- t,
- cond=prompt_embeds,
- timestep_cond=timestep_cond,
- added_cond_kwargs=added_cond_kwargs,
- )
-
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
- noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
- return noise_pred
-
- # 10. Run k-diffusion solver
- sampler_kwargs = {}
-
- if "noise_sampler" in inspect.signature(self.sampler).parameters:
- min_sigma, max_sigma = sigmas[sigmas > 0].min(), sigmas.max()
- noise_sampler = BrownianTreeNoiseSampler(latents, min_sigma, max_sigma, noise_sampler_seed)
- sampler_kwargs["noise_sampler"] = noise_sampler
-
- if "generator" in inspect.signature(self.sampler).parameters:
- sampler_kwargs["generator"] = generator
-
- latents = self.sampler(model_fn, latents, sigmas, **sampler_kwargs)
-
- if not output_type == "latent":
- # make sure the VAE is in float32 mode, as it overflows in float16
- needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
-
- if needs_upcasting:
- self.upcast_vae()
- latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
-
- image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-
- # cast back to fp16 if needed
- if needs_upcasting:
- self.vae.to(dtype=torch.float16)
- else:
- image = latents
-
- if not output_type == "latent":
- image = self.image_processor.postprocess(image, output_type=output_type)
-
- # Offload all models
- self.maybe_free_model_hooks()
-
- if not return_dict:
- return (image,)
-
- return StableDiffusionXLPipelineOutput(images=image)
diff --git a/src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py b/src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py
index c32121c88c9b..6de144aa7e8b 100644
--- a/src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py
+++ b/src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -98,10 +98,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -116,15 +116,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -161,20 +161,20 @@ class LDM3DPipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- rgb (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ rgb (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
- depth (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ depth (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
- nsfw_content_detected (`List[bool]`)
- List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
+ nsfw_content_detected (`list[bool]`)
+ list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
`None` if safety checking could not be performed.
"""
- rgb: Union[List[PIL.Image.Image], np.ndarray]
- depth: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_content_detected: Optional[List[bool]]
+ rgb: list[PIL.Image.Image] | np.ndarray
+ depth: list[PIL.Image.Image] | np.ndarray
+ nsfw_content_detected: list[bool] | None
class StableDiffusionLDM3DPipeline(
@@ -235,7 +235,7 @@ def __init__(
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
- image_encoder: Optional[CLIPVisionModelWithProjection],
+ image_encoder: CLIPVisionModelWithProjection | None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -278,9 +278,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -311,16 +311,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -328,7 +328,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -427,7 +427,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -748,36 +748,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 49,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -786,18 +786,18 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 5.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -805,7 +805,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -820,7 +820,7 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -841,7 +841,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py b/src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py
index 6d93e5feab4d..259fbd933430 100644
--- a/src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py
+++ b/src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py
@@ -13,7 +13,7 @@
import copy
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
@@ -98,10 +98,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -116,15 +116,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -211,7 +211,7 @@ def __init__(
scheduler: DDIMScheduler,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -254,9 +254,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -287,16 +287,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -304,7 +304,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -403,7 +403,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -735,7 +735,7 @@ def get_views(
window_size: int = 64,
stride: int = 8,
circular_padding: bool = False,
- ) -> List[Tuple[int, int, int, int]]:
+ ) -> list[tuple[int, int, int, int]]:
"""
Generates a list of views based on the given parameters. Here, we define the mappings F_i (see Eq. 7 in the
MultiDiffusion paper https://huggingface.co/papers/2302.08113). If panorama's height/width < window_size,
@@ -749,7 +749,7 @@ def get_views(
circular_padding (bool, optional): Whether to apply circular padding. Defaults to False.
Returns:
- List[Tuple[int, int, int, int]]: A list of tuples representing the views. Each tuple contains four integers
+ list[tuple[int, int, int, int]]: A list of tuples representing the views. Each tuple contains four integers
representing the start and end coordinates of the window in the panorama.
"""
@@ -802,37 +802,37 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = 512,
- width: Optional[int] = 2048,
+ prompt: str | list[str] = None,
+ height: int | None = 512,
+ width: int | None = 2048,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
+ timesteps: list[int] = None,
guidance_scale: float = 7.5,
view_batch_size: int = 1,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
circular_padding: bool = False,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs: Any,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to 512):
The height in pixels of the generated image.
@@ -842,7 +842,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
The timesteps at which to generate the images. If not specified, then the default timestep spacing
strategy of the scheduler is used.
guidance_scale (`float`, *optional*, defaults to 7.5):
@@ -851,7 +851,7 @@ def __call__(
view_batch_size (`int`, *optional*, defaults to 1):
The batch size to denoise split views. For some GPUs with high performance, higher view batch size can
speedup the generation and increase the VRAM usage.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -859,7 +859,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -874,7 +874,7 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -902,7 +902,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ callback_on_step_end_tensor_inputs (`list[str]`, *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.
diff --git a/src/diffusers/pipelines/stable_diffusion_safe/__init__.py b/src/diffusers/pipelines/stable_diffusion_safe/__init__.py
index b432b9418c46..b35015a9f729 100644
--- a/src/diffusers/pipelines/stable_diffusion_safe/__init__.py
+++ b/src/diffusers/pipelines/stable_diffusion_safe/__init__.py
@@ -1,6 +1,6 @@
from dataclasses import dataclass
from enum import Enum
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING
import numpy as np
import PIL
diff --git a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py
index 69a064d6638d..6b784bb0e102 100644
--- a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py
+++ b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -15,20 +14,20 @@ class StableDiffusionSafePipelineOutput(BaseOutput):
Output class for Safe Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
- nsfw_content_detected (`List[bool]`)
- List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
+ nsfw_content_detected (`list[bool]`)
+ list of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
(nsfw) content, or `None` if safety checking could not be performed.
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images that were flagged by the safety checker any may contain "not-safe-for-work"
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images that were flagged by the safety checker any may contain "not-safe-for-work"
(nsfw) content, or `None` if no safety check was performed or no images were flagged.
applied_safety_concept (`str`)
The safety concept that was applied for safety guidance, or `None` if safety guidance was disabled
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_content_detected: Optional[List[bool]]
- unsafe_images: Optional[Union[List[PIL.Image.Image], np.ndarray]]
- applied_safety_concept: Optional[str]
+ images: list[PIL.Image.Image] | np.ndarray
+ nsfw_content_detected: list[bool] | None
+ unsafe_images: list[PIL.Image.Image] | np.ndarray | None
+ applied_safety_concept: str | None
diff --git a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py
index d334107b0703..26bb5128ba9b 100644
--- a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py
+++ b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py
@@ -1,6 +1,6 @@
import inspect
import warnings
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -73,11 +73,11 @@ def __init__(
scheduler: KarrasDiffusionSchedulers,
safety_checker: SafeStableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
requires_safety_checker: bool = True,
):
super().__init__()
- safety_concept: Optional[str] = (
+ safety_concept: str | None = (
"an image showing hate, harassment, violence, suffering, humiliation, harm, suicide, sexual, nudity,"
" bodily fluids, blood, obscene gestures, illegal activity, drug use, theft, vandalism, weapons, child"
" abuse, brutality, cruelty"
@@ -199,7 +199,7 @@ def _encode_prompt(
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -207,7 +207,7 @@ def _encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`):
+ negative_prompt (`str` or `list[str]`):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
"""
@@ -248,7 +248,7 @@ def _encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -374,7 +374,6 @@ def prepare_extra_step_kwargs(self, generator, eta):
extra_step_kwargs["generator"] = generator
return extra_step_kwargs
- # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs
def check_inputs(
self,
prompt,
@@ -520,32 +519,32 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str],
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- sld_guidance_scale: Optional[float] = 1000,
- sld_warmup_steps: Optional[int] = 10,
- sld_threshold: Optional[float] = 0.01,
- sld_momentum_scale: Optional[float] = 0.3,
- sld_mom_beta: Optional[float] = 0.4,
+ sld_guidance_scale: float | None = 1000,
+ sld_warmup_steps: int | None = 10,
+ sld_threshold: float | None = 0.01,
+ sld_momentum_scale: float | None = 0.3,
+ sld_mom_beta: float | None = 0.4,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -557,7 +556,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -565,7 +564,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py b/src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py
index 48add535a81d..3cf604911f0b 100644
--- a/src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py
+++ b/src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
import torch.nn.functional as F
@@ -155,7 +155,7 @@ def __init__(
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
- image_encoder: Optional[CLIPVisionModelWithProjection] = None,
+ image_encoder: CLIPVisionModelWithProjection | None = None,
requires_safety_checker: bool = True,
):
super().__init__()
@@ -182,9 +182,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -215,16 +215,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -232,7 +232,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -331,7 +331,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -494,7 +494,6 @@ def prepare_extra_step_kwargs(self, generator, eta):
extra_step_kwargs["generator"] = generator
return extra_step_kwargs
- # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs
def check_inputs(
self,
prompt,
@@ -574,33 +573,33 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
sag_scale: float = 0.75,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
@@ -614,7 +613,7 @@ def __call__(
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
sag_scale (`float`, *optional*, defaults to 0.75):
Chosen between [0, 1.0] for better quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -622,7 +621,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -637,7 +636,7 @@ def __call__(
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*):
Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. If not provided, embeddings are computed from the
`ip_adapter_image` input argument.
output_type (`str`, *optional*, defaults to `"pil"`):
diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py
index 3227fd9a08a4..f80abc7073ed 100644
--- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py
+++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py
@@ -12,8 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
from functools import partial
-from typing import Dict, List, Optional, Union
import jax
import jax.numpy as jnp
@@ -48,9 +49,7 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: FlaxUNet2DConditionModel,
- scheduler: Union[
- FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler
- ],
+ scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler,
dtype: jnp.dtype = jnp.float32,
):
super().__init__()
@@ -67,7 +66,7 @@ def __init__(
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
- def prepare_inputs(self, prompt: Union[str, List[str]]):
+ def prepare_inputs(self, prompt: str | list[str]):
if not isinstance(prompt, (str, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
@@ -88,12 +87,12 @@ def prepare_inputs(self, prompt: Union[str, List[str]]):
def __call__(
self,
prompt_ids: jax.Array,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
num_inference_steps: int = 50,
- guidance_scale: Union[float, jax.Array] = 7.5,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ guidance_scale: float | jax.Array = 7.5,
+ height: int | None = None,
+ width: int | None = None,
latents: jnp.array = None,
neg_prompt_ids: jnp.array = None,
return_dict: bool = True,
@@ -169,14 +168,14 @@ def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, b
def _generate(
self,
prompt_ids: jnp.array,
- params: Union[Dict, FrozenDict],
+ params: dict | FrozenDict,
prng_seed: jax.Array,
num_inference_steps: int,
height: int,
width: int,
guidance_scale: float,
- latents: Optional[jnp.array] = None,
- neg_prompt_ids: Optional[jnp.array] = None,
+ latents: jnp.array | None = None,
+ neg_prompt_ids: jnp.array | None = None,
return_latents=False,
):
if height % 8 != 0 or width % 8 != 0:
diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py
index 0783f44486ee..6617e7cd058c 100644
--- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py
+++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,12 +12,12 @@ class StableDiffusionXLPipelineOutput(BaseOutput):
Output class for Stable Diffusion pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
if is_flax_available():
diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
index 3a63bb4f253a..2f6b105702e8 100644
--- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
+++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import torch
from transformers import (
@@ -110,10 +110,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -128,15 +128,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -248,7 +248,7 @@ def __init__(
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
):
super().__init__()
@@ -283,26 +283,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -311,11 +311,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -433,7 +433,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -822,52 +822,50 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
@@ -883,11 +881,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -904,11 +902,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -916,7 +914,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -938,7 +936,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -959,31 +957,31 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -993,7 +991,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py
index d1916b635f92..19ccfab3de0a 100644
--- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py
+++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -113,7 +113,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -128,10 +128,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -146,15 +146,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -270,7 +270,7 @@ def __init__(
feature_extractor: CLIPImageProcessor = None,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
):
super().__init__()
@@ -301,26 +301,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -329,11 +329,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -451,7 +451,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -974,58 +974,56 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
strength: float = 0.3,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_start: Optional[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_start: float | None = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
- image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`):
+ image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[np.ndarray]`):
The image(s) to modify with the pipeline.
strength (`float`, *optional*, defaults to 0.3):
Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`
@@ -1037,11 +1035,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1066,11 +1064,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -1078,7 +1076,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -1100,7 +1098,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1121,31 +1119,31 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1166,7 +1164,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py
index fcfddc192b8b..8de7d4f0bb7d 100644
--- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py
+++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -140,7 +140,7 @@ def mask_pil_to_torch(mask, height, width):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -155,10 +155,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -173,15 +173,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -300,7 +300,7 @@ def __init__(
feature_extractor: CLIPImageProcessor = None,
requires_aesthetics_score: bool = False,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
):
super().__init__()
@@ -405,26 +405,26 @@ def prepare_ip_adapter_image_embeds(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -433,11 +433,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -555,7 +555,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -1079,60 +1079,58 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
masked_image_latents: torch.Tensor = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- padding_mask_crop: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
strength: float = 0.9999,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_start: Optional[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_start: float | None = None,
+ denoising_end: float | None = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_skip: int | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
image (`PIL.Image.Image`):
@@ -1171,11 +1169,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -1200,11 +1198,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -1222,7 +1220,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1249,31 +1247,31 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
@@ -1294,7 +1292,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py
index 5e13362eb3d1..b79119a94a0c 100644
--- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py
+++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL.Image
import torch
@@ -84,7 +84,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -179,8 +179,8 @@ def __init__(
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
- is_cosxl_edit: Optional[bool] = False,
+ add_watermarker: bool | None = None,
+ is_cosxl_edit: bool | None = False,
):
super().__init__()
@@ -213,25 +213,25 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -240,11 +240,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -350,7 +350,7 @@ def encode_prompt(
negative_prompt = negative_prompt or ""
negative_prompt_2 = negative_prompt_2 or negative_prompt
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -596,46 +596,46 @@ def upcast_vae(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 100,
- denoising_end: Optional[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
image_guidance_scale: float = 1.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Tuple[int, int] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Tuple[int, int] = None,
+ original_size: tuple[int, int] = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
- image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`):
+ image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[np.ndarray]`):
The image(s) to modify with the pipeline.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The height in pixels of the generated image.
@@ -662,11 +662,11 @@ def __call__(
scale is enabled by setting `image_guidance_scale > 1`. Higher image guidance scale encourages to
generate images that are closely linked to the source image `image`, usually at the expense of lower
image quality. This pipeline requires a value of at least `1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders.
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -674,7 +674,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -717,17 +717,17 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
diff --git a/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py b/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py
index 633094239dca..05877f69d403 100644
--- a/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py
+++ b/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -73,10 +73,10 @@ def _append_dims(x, target_dims):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -91,15 +91,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -136,12 +136,12 @@ class StableVideoDiffusionPipelineOutput(BaseOutput):
Output class for Stable Video Diffusion pipeline.
Args:
- frames (`[List[List[PIL.Image.Image]]`, `np.ndarray`, `torch.Tensor`]):
- List of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size,
+ frames (`[list[list[PIL.Image.Image]]`, `np.ndarray`, `torch.Tensor`]):
+ list of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size,
num_frames, height, width, num_channels)`.
"""
- frames: Union[List[List[PIL.Image.Image]], np.ndarray, torch.Tensor]
+ frames: list[list[PIL.Image.Image]] | np.ndarray | torch.Tensor
class StableVideoDiffusionPipeline(DiffusionPipeline):
@@ -191,7 +191,7 @@ def __init__(
def _encode_image(
self,
image: PipelineImageInput,
- device: Union[str, torch.device],
+ device: str | torch.device,
num_videos_per_prompt: int,
do_classifier_free_guidance: bool,
) -> torch.Tensor:
@@ -239,7 +239,7 @@ def _encode_image(
def _encode_vae_image(
self,
image: torch.Tensor,
- device: Union[str, torch.device],
+ device: str | torch.device,
num_videos_per_prompt: int,
do_classifier_free_guidance: bool,
):
@@ -323,7 +323,7 @@ def check_inputs(self, image, height, width):
and not isinstance(image, list)
):
raise ValueError(
- "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+ "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
f" {type(image)}"
)
@@ -338,9 +338,9 @@ def prepare_latents(
height: int,
width: int,
dtype: torch.dtype,
- device: Union[str, torch.device],
+ device: str | torch.device,
generator: torch.Generator,
- latents: Optional[torch.Tensor] = None,
+ latents: torch.Tensor | None = None,
):
shape = (
batch_size,
@@ -385,31 +385,31 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor],
+ image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor,
height: int = 576,
width: int = 1024,
- num_frames: Optional[int] = None,
+ num_frames: int | None = None,
num_inference_steps: int = 25,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
min_guidance_scale: float = 1.0,
max_guidance_scale: float = 3.0,
fps: int = 7,
motion_bucket_id: int = 127,
noise_aug_strength: float = 0.02,
- decode_chunk_size: Optional[int] = None,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ decode_chunk_size: int | None = None,
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
return_dict: bool = True,
):
r"""
The call function to the pipeline for generation.
Args:
- image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`):
+ image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.Tensor`):
Image(s) to guide image generation. If you provide a tensor, the expected value range is between `[0,
1]`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
@@ -422,7 +422,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 25):
The number of denoising steps. More denoising steps usually lead to a higher quality video at the
expense of slower inference. This parameter is modulated by `strength`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -445,7 +445,7 @@ def __call__(
For lower memory usage, reduce `decode_chunk_size`.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of videos to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -460,7 +460,7 @@ def __call__(
`callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`.
`callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
@@ -473,7 +473,7 @@ def __call__(
Returns:
[`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] or `tuple`:
If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] is
- returned, otherwise a `tuple` of (`List[List[PIL.Image.Image]]` or `np.ndarray` or `torch.Tensor`) is
+ returned, otherwise a `tuple` of (`list[list[PIL.Image.Image]]` or `np.ndarray` or `torch.Tensor`) is
returned.
"""
# 0. Default height and width to unet
@@ -697,7 +697,7 @@ def _filter2d(input, kernel):
height, width = tmp_kernel.shape[-2:]
- padding_shape: List[int] = _compute_padding([height, width])
+ padding_shape: list[int] = _compute_padding([height, width])
input = torch.nn.functional.pad(input, padding_shape, mode="reflect")
# kernel and input tensor reshape to align element-wise or batch-wise params
diff --git a/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py b/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py
index 7b6673cf16f7..ffb877cfd0f6 100644
--- a/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py
+++ b/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py
@@ -14,7 +14,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -54,16 +54,16 @@
class StableDiffusionAdapterPipelineOutput(BaseOutput):
"""
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
- nsfw_content_detected (`List[bool]`)
- List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
+ nsfw_content_detected (`list[bool]`)
+ list of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
(nsfw) content, or `None` if safety checking could not be performed.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_content_detected: Optional[List[bool]]
+ images: list[PIL.Image.Image] | np.ndarray
+ nsfw_content_detected: list[bool] | None
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
@@ -131,10 +131,10 @@ def _preprocess_adapter_image(image, height, width):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -149,15 +149,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -197,11 +197,11 @@ class StableDiffusionAdapterPipeline(DiffusionPipeline, StableDiffusionMixin, Fr
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
- adapter ([`T2IAdapter`] or [`MultiAdapter`] or `List[T2IAdapter]`):
+ adapter ([`T2IAdapter`] or [`MultiAdapter`] or `list[T2IAdapter]`):
Provides additional conditioning to the unet during the denoising process. If you set multiple Adapter as a
list, the outputs from each Adapter are added together to create one combined additional conditioning.
- adapter_weights (`List[float]`, *optional*, defaults to None):
- List of floats representing the weight which will be multiply to each adapter's output before adding them
+ adapter_weights (`list[float]`, *optional*, defaults to None):
+ list of floats representing the weight which will be multiply to each adapter's output before adding them
together.
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
@@ -233,7 +233,7 @@ def __init__(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
- adapter: Union[T2IAdapter, MultiAdapter, List[T2IAdapter]],
+ adapter: T2IAdapter | MultiAdapter | list[T2IAdapter],
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
@@ -282,9 +282,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -315,16 +315,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -332,7 +332,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -431,7 +431,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -690,37 +690,37 @@ def do_classifier_free_guidance(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- image: Union[torch.Tensor, PIL.Image.Image, List[PIL.Image.Image]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | PIL.Image.Image | list[PIL.Image.Image] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- adapter_conditioning_scale: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ adapter_conditioning_scale: float | list[float] = 1.0,
+ clip_skip: int | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[list[PIL.Image.Image]]`):
The Adapter input condition. Adapter uses this input condition to generate guidance to Unet. If the
type is specified as `torch.Tensor`, it is passed to Adapter as is. PIL.Image.Image` can also be
accepted as an image. The control image is automatically resized to fit the output image.
@@ -731,11 +731,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -745,7 +745,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
@@ -754,7 +754,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -784,7 +784,7 @@ def __call__(
A kwargs dictionary that if specified is passed along to the `AttnProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- adapter_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ adapter_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the adapter are multiplied by `adapter_conditioning_scale` before they are added to the
residual in the original unet. If multiple adapters are specified in init, you can set the
corresponding scale as a list.
diff --git a/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py b/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py
index bf089bf540ba..a6dd07847de2 100644
--- a/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py
+++ b/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL.Image
@@ -153,10 +153,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -171,15 +171,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -233,11 +233,11 @@ class StableDiffusionXLAdapterPipeline(
- [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
Args:
- adapter ([`T2IAdapter`] or [`MultiAdapter`] or `List[T2IAdapter]`):
+ adapter ([`T2IAdapter`] or [`MultiAdapter`] or `list[T2IAdapter]`):
Provides additional conditioning to the unet during the denoising process. If you set multiple Adapter as a
list, the outputs from each Adapter are added together to create one combined additional conditioning.
- adapter_weights (`List[float]`, *optional*, defaults to None):
- List of floats representing the weight which will be multiply to each adapter's output before adding them
+ adapter_weights (`list[float]`, *optional*, defaults to None):
+ list of floats representing the weight which will be multiply to each adapter's output before adding them
together.
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
@@ -278,7 +278,7 @@ def __init__(
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
- adapter: Union[T2IAdapter, MultiAdapter, List[T2IAdapter]],
+ adapter: T2IAdapter | MultiAdapter | list[T2IAdapter],
scheduler: KarrasDiffusionSchedulers,
force_zeros_for_empty_prompt: bool = True,
feature_extractor: CLIPImageProcessor = None,
@@ -311,26 +311,26 @@ def __init__(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -339,11 +339,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -461,7 +461,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -857,55 +857,55 @@ def do_classifier_free_guidance(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] = None,
+ prompt_2: str | list[str] | None = None,
image: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- timesteps: List[int] = None,
- sigmas: List[float] = None,
- denoising_end: Optional[float] = None,
+ timesteps: list[int] = None,
+ sigmas: list[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 5.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ ip_adapter_image: PipelineImageInput | None = None,
+ ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
- negative_original_size: Optional[Tuple[int, int]] = None,
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
- negative_target_size: Optional[Tuple[int, int]] = None,
- adapter_conditioning_scale: Union[float, List[float]] = 1.0,
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
+ negative_original_size: tuple[int, int] | None = None,
+ negative_crops_coords_top_left: tuple[int, int] = (0, 0),
+ negative_target_size: tuple[int, int] | None = None,
+ adapter_conditioning_scale: float | list[float] = 1.0,
adapter_conditioning_factor: float = 1.0,
- clip_skip: Optional[int] = None,
+ clip_skip: int | None = None,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
- image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[List[PIL.Image.Image]]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[list[PIL.Image.Image]]`):
The Adapter input condition. Adapter uses this input condition to generate guidance to Unet. If the
type is specified as `torch.Tensor`, it is passed to Adapter as is. PIL.Image.Image` can also be
accepted as an image. The control image is automatically resized to fit the output image.
@@ -920,11 +920,11 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -941,11 +941,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -953,7 +953,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -975,7 +975,7 @@ def __call__(
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
@@ -1002,37 +1002,37 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
- adapter_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+ adapter_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0):
The outputs of the adapter are multiplied by `adapter_conditioning_scale` before they are added to the
residual in the original unet. If multiple adapters are specified in init, you can set the
corresponding scale as a list.
diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py
index 040bf0efba84..c94c5d2d144a 100644
--- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py
+++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py
@@ -1,5 +1,4 @@
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL
@@ -16,11 +15,11 @@ class TextToVideoSDPipelineOutput(BaseOutput):
Output class for text-to-video pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised
PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`
"""
- frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]]
+ frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]
diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py
index 3ce7b4d1990f..0ca64d33acda 100644
--- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py
+++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import CLIPTextModel, CLIPTokenizer
@@ -130,9 +130,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -163,16 +163,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -180,7 +180,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -279,7 +279,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -368,7 +368,6 @@ def prepare_extra_step_kwargs(self, generator, eta):
extra_step_kwargs["generator"] = generator
return extra_step_kwargs
- # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs
def check_inputs(
self,
prompt,
@@ -450,30 +449,30 @@ def prepare_latents(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_frames: int = 16,
num_inference_steps: int = 50,
guidance_scale: float = 9.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated video.
@@ -488,7 +487,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -496,7 +495,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py
index 9d0b7e3dbc32..6908f51eb21b 100644
--- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py
+++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -90,7 +90,7 @@
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -165,9 +165,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -198,16 +198,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -215,7 +215,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -314,7 +314,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -515,31 +515,31 @@ def prepare_latents(self, video, timestep, batch_size, dtype, device, generator=
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- video: Union[List[np.ndarray], torch.Tensor] = None,
+ prompt: str | list[str] = None,
+ video: list[np.ndarray] | torch.Tensor = None,
strength: float = 0.6,
num_inference_steps: int = 50,
guidance_scale: float = 15.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- clip_skip: Optional[int] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
+ clip_skip: int | None = None,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- video (`List[np.ndarray]` or `torch.Tensor`):
+ video (`list[np.ndarray]` or `torch.Tensor`):
`video` frames or tensor representing a video batch to be used as the starting point for the process.
Can also accept video latents as `image`, if passing latents directly, it will not be encoded again.
strength (`float`, *optional*, defaults to 0.8):
@@ -554,13 +554,13 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in video generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py
index 96316f8e91e5..66defb2f3745 100644
--- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py
+++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py
@@ -1,7 +1,7 @@
import copy
import inspect
from dataclasses import dataclass
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -199,16 +199,16 @@ class TextToVideoPipelineOutput(BaseOutput):
Output class for zero-shot text-to-video pipeline.
Args:
- images (`[List[PIL.Image.Image]`, `np.ndarray`]):
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ images (`[list[PIL.Image.Image]`, `np.ndarray`]):
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
- nsfw_content_detected (`[List[bool]]`):
- List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
+ nsfw_content_detected (`[list[bool]]`):
+ list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
`None` if safety checking could not be performed.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
- nsfw_content_detected: Optional[List[bool]]
+ images: list[PIL.Image.Image] | np.ndarray
+ nsfw_content_detected: list[bool] | None
def coords_grid(batch, ht, wd, device):
@@ -374,7 +374,7 @@ def forward_loop(self, x_t0, t0, t1, generator):
Timestep at t0.
t1:
Timestamp at t1.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
@@ -466,7 +466,6 @@ def backward_loop(
return latents.clone().detach()
- # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs
def check_inputs(
self,
prompt,
@@ -545,32 +544,32 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
- video_length: Optional[int] = 8,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str],
+ video_length: int | None = 8,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
motion_field_strength_x: float = 12,
motion_field_strength_y: float = 12,
- output_type: Optional[str] = "tensor",
+ output_type: str | None = "tensor",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
- callback_steps: Optional[int] = 1,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
+ callback_steps: int | None = 1,
t0: int = 44,
t1: int = 47,
- frame_ids: Optional[List[int]] = None,
+ frame_ids: list[int] | None = None,
):
"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
video_length (`int`, *optional*, defaults to 8):
The number of generated video frames.
@@ -584,7 +583,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in video generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_videos_per_prompt (`int`, *optional*, defaults to 1):
@@ -592,7 +591,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -623,7 +622,7 @@ def __call__(
t1 (`int`, *optional*, defaults to 47):
Timestep t0. Should be in the range [t0 + 1, num_inference_steps - 1]. See the
[paper](https://huggingface.co/papers/2303.13439), Sect. 3.3.1.
- frame_ids (`List[int]`, *optional*):
+ frame_ids (`list[int]`, *optional*):
Indexes of the frames that are being generated. This is used when generating longer videos
chunk-by-chunk.
@@ -821,16 +820,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -838,7 +837,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -937,7 +936,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py
index c8dce75e2671..a3286cd940fd 100644
--- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py
+++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py
@@ -1,7 +1,7 @@
import copy
import inspect
from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import numpy as np
import PIL
@@ -218,12 +218,12 @@ class TextToVideoSDXLPipelineOutput(BaseOutput):
Output class for zero-shot text-to-video pipeline.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.coords_grid
@@ -402,7 +402,7 @@ def __init__(
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
force_zeros_for_empty_prompt: bool = True,
- add_watermarker: Optional[bool] = None,
+ add_watermarker: bool | None = None,
):
super().__init__()
self.register_modules(
@@ -585,26 +585,26 @@ def check_inputs(
def encode_prompt(
self,
prompt: str,
- prompt_2: Optional[str] = None,
- device: Optional[torch.device] = None,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[str] = None,
- negative_prompt_2: Optional[str] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
device: (`torch.device`):
@@ -613,11 +613,11 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
prompt_embeds (`torch.Tensor`, *optional*):
@@ -735,7 +735,7 @@ def encode_prompt(
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -832,7 +832,7 @@ def forward_loop(self, x_t0, t0, t1, generator):
Timestep at t0.
t1:
Timestamp at t1.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
@@ -937,36 +937,36 @@ def backward_loop(
@torch.no_grad()
def __call__(
self,
- prompt: Union[str, List[str]],
- prompt_2: Optional[Union[str, List[str]]] = None,
- video_length: Optional[int] = 8,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ video_length: int | None = 8,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- denoising_end: Optional[float] = None,
+ denoising_end: float | None = None,
guidance_scale: float = 7.5,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
- num_videos_per_prompt: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ num_videos_per_prompt: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- frame_ids: Optional[List[int]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- pooled_prompt_embeds: Optional[torch.Tensor] = None,
- negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
- latents: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ frame_ids: list[int] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ latents: torch.Tensor | None = None,
motion_field_strength_x: float = 12,
motion_field_strength_y: float = 12,
- output_type: Optional[str] = "tensor",
+ output_type: str | None = "tensor",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ cross_attention_kwargs: dict[str, Any] | None = None,
guidance_rescale: float = 0.0,
- original_size: Optional[Tuple[int, int]] = None,
- crops_coords_top_left: Tuple[int, int] = (0, 0),
- target_size: Optional[Tuple[int, int]] = None,
+ original_size: tuple[int, int] | None = None,
+ crops_coords_top_left: tuple[int, int] = (0, 0),
+ target_size: tuple[int, int] | None = None,
t0: int = 44,
t1: int = 47,
):
@@ -974,10 +974,10 @@ def __call__(
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- prompt_2 (`str` or `List[str]`, *optional*):
+ prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
video_length (`int`, *optional*, defaults to 8):
@@ -1002,11 +1002,11 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- negative_prompt_2 (`str` or `List[str]`, *optional*):
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_videos_per_prompt (`int`, *optional*, defaults to 1):
@@ -1014,10 +1014,10 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
applies to [`schedulers.DDIMScheduler`], will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
- frame_ids (`List[int]`, *optional*):
+ frame_ids (`list[int]`, *optional*):
Indexes of the frames that are being generated. This is used when generating longer videos
chunk-by-chunk.
prompt_embeds (`torch.Tensor`, *optional*):
@@ -1066,17 +1066,17 @@ def __call__(
[Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
diff --git a/src/diffusers/pipelines/transformers_loading_utils.py b/src/diffusers/pipelines/transformers_loading_utils.py
index b52d154d6ba2..02f8ee612601 100644
--- a/src/diffusers/pipelines/transformers_loading_utils.py
+++ b/src/diffusers/pipelines/transformers_loading_utils.py
@@ -15,7 +15,7 @@
import contextlib
import os
import tempfile
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING
from huggingface_hub import DDUFEntry
from tqdm import tqdm
@@ -34,7 +34,7 @@
def _load_tokenizer_from_dduf(
- cls: "PreTrainedTokenizer", name: str, dduf_entries: Dict[str, DDUFEntry], **kwargs
+ cls: "PreTrainedTokenizer", name: str, dduf_entries: dict[str, DDUFEntry], **kwargs
) -> "PreTrainedTokenizer":
"""
Load a tokenizer from a DDUF archive.
@@ -57,7 +57,7 @@ def _load_tokenizer_from_dduf(
def _load_transformers_model_from_dduf(
- cls: "PreTrainedModel", name: str, dduf_entries: Dict[str, DDUFEntry], **kwargs
+ cls: "PreTrainedModel", name: str, dduf_entries: dict[str, DDUFEntry], **kwargs
) -> "PreTrainedModel":
"""
Load a transformers model from a DDUF archive.
@@ -112,6 +112,8 @@ def _load_transformers_model_from_dduf(
tensors = safetensors.torch.load(mmap)
# Update the state dictionary with tensors
state_dict.update(tensors)
+ # `from_pretrained` sets the model to eval mode by default, which is the
+ # correct behavior for inference. Do not call `model.train()` here.
return cls.from_pretrained(
pretrained_model_name_or_path=None,
config=config,
diff --git a/src/diffusers/pipelines/unclip/pipeline_unclip.py b/src/diffusers/pipelines/unclip/pipeline_unclip.py
index bbb9b0eb3ab2..430f1a1e5265 100644
--- a/src/diffusers/pipelines/unclip/pipeline_unclip.py
+++ b/src/diffusers/pipelines/unclip/pipeline_unclip.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Tuple, Union
import torch
from torch.nn import functional as F
@@ -131,8 +130,8 @@ def _encode_prompt(
device,
num_images_per_prompt,
do_classifier_free_guidance,
- text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None,
- text_attention_mask: Optional[torch.Tensor] = None,
+ text_model_output: CLIPTextModelOutput | tuple | None = None,
+ text_attention_mask: torch.Tensor | None = None,
):
if text_model_output is None:
batch_size = len(prompt) if isinstance(prompt, list) else 1
@@ -219,27 +218,27 @@ def _encode_prompt(
@torch.no_grad()
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] | None = None,
num_images_per_prompt: int = 1,
prior_num_inference_steps: int = 25,
decoder_num_inference_steps: int = 25,
super_res_num_inference_steps: int = 7,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- prior_latents: Optional[torch.Tensor] = None,
- decoder_latents: Optional[torch.Tensor] = None,
- super_res_latents: Optional[torch.Tensor] = None,
- text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None,
- text_attention_mask: Optional[torch.Tensor] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ prior_latents: torch.Tensor | None = None,
+ decoder_latents: torch.Tensor | None = None,
+ super_res_latents: torch.Tensor | None = None,
+ text_model_output: CLIPTextModelOutput | tuple | None = None,
+ text_attention_mask: torch.Tensor | None = None,
prior_guidance_scale: float = 4.0,
decoder_guidance_scale: float = 8.0,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide image generation. This can only be left undefined if `text_model_output`
and `text_attention_mask` is passed.
num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -253,7 +252,7 @@ def __call__(
super_res_num_inference_steps (`int`, *optional*, defaults to 7):
The number of denoising steps for super resolution. More denoising steps usually lead to a higher
quality image at the expense of slower inference.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
prior_latents (`torch.Tensor` of shape (batch size, embeddings dimension), *optional*):
diff --git a/src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py b/src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py
index 31710a000e0a..d0d8bdc44787 100644
--- a/src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py
+++ b/src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py
@@ -13,7 +13,6 @@
# limitations under the License.
import inspect
-from typing import List, Optional, Union
import PIL.Image
import torch
@@ -190,7 +189,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr
return prompt_embeds, text_encoder_hidden_states, text_mask
- def _encode_image(self, image, device, num_images_per_prompt, image_embeddings: Optional[torch.Tensor] = None):
+ def _encode_image(self, image, device, num_images_per_prompt, image_embeddings: torch.Tensor | None = None):
dtype = next(self.image_encoder.parameters()).dtype
if image_embeddings is None:
@@ -207,23 +206,23 @@ def _encode_image(self, image, device, num_images_per_prompt, image_embeddings:
@torch.no_grad()
def __call__(
self,
- image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor]] = None,
+ image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor | None = None,
num_images_per_prompt: int = 1,
decoder_num_inference_steps: int = 25,
super_res_num_inference_steps: int = 7,
- generator: Optional[torch.Generator] = None,
- decoder_latents: Optional[torch.Tensor] = None,
- super_res_latents: Optional[torch.Tensor] = None,
- image_embeddings: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ decoder_latents: torch.Tensor | None = None,
+ super_res_latents: torch.Tensor | None = None,
+ image_embeddings: torch.Tensor | None = None,
decoder_guidance_scale: float = 8.0,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
):
"""
The call function to the pipeline for generation.
Args:
- image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`):
+ image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.Tensor`):
`Image` or tensor representing an image batch to be used as the starting point. If you provide a
tensor, it needs to be compatible with the [`CLIPImageProcessor`]
[configuration](https://huggingface.co/fusing/karlo-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
diff --git a/src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py b/src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py
index 0ddcbf735770..c68b5d9ab5a8 100644
--- a/src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py
+++ b/src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py
@@ -1,5 +1,3 @@
-from typing import Optional
-
import numpy as np
import torch
from torch import nn
@@ -68,13 +66,13 @@ def __init__(
self,
prefix_length: int,
prefix_inner_dim: int,
- prefix_hidden_dim: Optional[int] = None,
+ prefix_hidden_dim: int | None = None,
vocab_size: int = 50257, # Start of GPT2 config args
n_positions: int = 1024,
n_embd: int = 768,
n_layer: int = 12,
n_head: int = 12,
- n_inner: Optional[int] = None,
+ n_inner: int | None = None,
activation_function: str = "gelu_new",
resid_pdrop: float = 0.1,
embd_pdrop: float = 0.1,
@@ -132,8 +130,8 @@ def forward(
self,
input_ids: torch.Tensor,
prefix_embeds: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
+ attention_mask: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
):
"""
Args:
@@ -180,7 +178,7 @@ def generate_captions(self, features, eos_token_id, device):
Device to perform text generation on.
Returns:
- `List[str]`: A list of strings generated from the decoder model.
+ `list[str]`: A list of strings generated from the decoder model.
"""
features = torch.split(features, 1, dim=0)
@@ -207,7 +205,7 @@ def generate_beam(
beam_size: int = 5,
entry_length: int = 67,
temperature: float = 1.0,
- eos_token_id: Optional[int] = None,
+ eos_token_id: int | None = None,
):
"""
Generates text using the given tokenizer and text prompt or token embedding via beam search. This
@@ -233,7 +231,7 @@ def generate_beam(
The temperature to use when performing the softmax over logits from the decoding model.
Returns:
- `Tuple(torch.Tensor, torch.Tensor)`: A tuple of tensors where the first element is a tensor of generated
+ `tuple(torch.Tensor, torch.Tensor)`: A tuple of tensors where the first element is a tensor of generated
token sequences sorted by score in descending order, and the second element is the sequence lengths
corresponding to those sequences.
"""
diff --git a/src/diffusers/pipelines/unidiffuser/modeling_uvit.py b/src/diffusers/pipelines/unidiffuser/modeling_uvit.py
index 2a04ec2e4030..125188196c1e 100644
--- a/src/diffusers/pipelines/unidiffuser/modeling_uvit.py
+++ b/src/diffusers/pipelines/unidiffuser/modeling_uvit.py
@@ -1,5 +1,4 @@
import math
-from typing import Optional, Union
import torch
from torch import nn
@@ -178,9 +177,9 @@ def __init__(
num_attention_heads: int,
attention_head_dim: int,
dropout=0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
attention_bias: bool = False,
only_cross_attention: bool = False,
double_self_attention: bool = False,
@@ -373,9 +372,9 @@ def __init__(
num_attention_heads: int,
attention_head_dim: int,
dropout=0.0,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
attention_bias: bool = False,
only_cross_attention: bool = False,
double_self_attention: bool = False,
@@ -591,18 +590,18 @@ def __init__(
self,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
- out_channels: Optional[int] = None,
+ in_channels: int | None = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
norm_num_groups: int = 32,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
- sample_size: Optional[int] = None,
- num_vector_embeds: Optional[int] = None,
- patch_size: Optional[int] = 2,
+ sample_size: int | None = None,
+ num_vector_embeds: int | None = None,
+ patch_size: int | None = 2,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
use_linear_projection: bool = False,
only_cross_attention: bool = False,
upcast_attention: bool = False,
@@ -902,18 +901,18 @@ def __init__(
num_text_tokens: int = 77,
num_attention_heads: int = 16,
attention_head_dim: int = 88,
- in_channels: Optional[int] = None,
- out_channels: Optional[int] = None,
+ in_channels: int | None = None,
+ out_channels: int | None = None,
num_layers: int = 1,
dropout: float = 0.0,
norm_num_groups: int = 32,
- cross_attention_dim: Optional[int] = None,
+ cross_attention_dim: int | None = None,
attention_bias: bool = False,
- sample_size: Optional[int] = None,
- num_vector_embeds: Optional[int] = None,
- patch_size: Optional[int] = None,
+ sample_size: int | None = None,
+ num_vector_embeds: int | None = None,
+ patch_size: int | None = None,
activation_fn: str = "geglu",
- num_embeds_ada_norm: Optional[int] = None,
+ num_embeds_ada_norm: int | None = None,
use_linear_projection: bool = False,
only_cross_attention: bool = False,
upcast_attention: bool = False,
@@ -1040,9 +1039,9 @@ def forward(
latent_image_embeds: torch.Tensor,
image_embeds: torch.Tensor,
prompt_embeds: torch.Tensor,
- timestep_img: Union[torch.Tensor, float, int],
- timestep_text: Union[torch.Tensor, float, int],
- data_type: Optional[Union[torch.Tensor, float, int]] = 1,
+ timestep_img: torch.Tensor | float | int,
+ timestep_text: torch.Tensor | float | int,
+ data_type: torch.Tensor | float | int | None = 1,
encoder_hidden_states=None,
cross_attention_kwargs=None,
):
diff --git a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py
index f9298d5b86f8..81d2ce95dc53 100644
--- a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py
+++ b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py
@@ -1,6 +1,6 @@
import inspect
from dataclasses import dataclass
-from typing import Callable, List, Optional, Union
+from typing import Callable
import numpy as np
import PIL.Image
@@ -50,16 +50,16 @@ class ImageTextPipelineOutput(BaseOutput):
Output class for joint image-text pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
- List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
num_channels)`.
- text (`List[str]` or `List[List[str]]`)
- List of generated text strings of length `batch_size` or a list of list of strings whose outer list has
+ text (`list[str]` or `list[list[str]]`)
+ list of generated text strings of length `batch_size` or a list of list of strings whose outer list has
length `batch_size`.
"""
- images: Optional[Union[List[PIL.Image.Image], np.ndarray]]
- text: Optional[Union[List[str], List[List[str]]]]
+ images: list[PIL.Image.Image] | np.ndarray | None
+ text: list[str] | list[list[str]] | None
class UniDiffuserPipeline(DeprecatedPipelineMixin, DiffusionPipeline):
@@ -392,9 +392,9 @@ def _encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
**kwargs,
):
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
@@ -425,16 +425,16 @@ def encode_prompt(
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
@@ -442,7 +442,7 @@ def encode_prompt(
number of images that should be generated per prompt
do_classifier_free_guidance (`bool`):
whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -541,7 +541,7 @@ def encode_prompt(
# get unconditional embeddings for classifier free guidance
if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif prompt is not None and type(prompt) is not type(negative_prompt):
@@ -1119,34 +1119,34 @@ def check_inputs(
@torch.no_grad()
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
- image: Optional[Union[torch.Tensor, PIL.Image.Image]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- data_type: Optional[int] = 1,
+ prompt: str | list[str] | None = None,
+ image: torch.Tensor | PIL.Image.Image | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ data_type: int | None = 1,
num_inference_steps: int = 50,
guidance_scale: float = 8.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- num_prompts_per_image: Optional[int] = 1,
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ num_prompts_per_image: int | None = 1,
eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_latents: Optional[torch.Tensor] = None,
- vae_latents: Optional[torch.Tensor] = None,
- clip_latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_latents: torch.Tensor | None = None,
+ vae_latents: torch.Tensor | None = None,
+ clip_latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback: Callable[[int, int, torch.Tensor], None] | None = None,
callback_steps: int = 1,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
Required for text-conditioned image generation (`text2img`) mode.
image (`torch.Tensor` or `PIL.Image.Image`, *optional*):
@@ -1166,7 +1166,7 @@ def __call__(
guidance_scale (`float`, *optional*, defaults to 8.0):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). Used in
text-conditioned image generation (`text2img`) mode.
@@ -1181,7 +1181,7 @@ def __call__(
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only
applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py
index 91a54e1ae82f..f640fddc2bc5 100644
--- a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py
+++ b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from PIL import Image
@@ -250,24 +250,24 @@ def check_inputs(
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- task_prompt: Union[str, List[str]] = None,
- content_prompt: Union[str, List[str]] = None,
- image: Optional[torch.FloatTensor] = None,
- upsampling_height: Optional[int] = None,
- upsampling_width: Optional[int] = None,
+ task_prompt: str | list[str] = None,
+ content_prompt: str | list[str] = None,
+ image: torch.FloatTensor | None = None,
+ upsampling_height: int | None = None,
+ upsampling_width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 30.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
upsampling_strength: float = 1.0,
):
@@ -275,11 +275,11 @@ def __call__(
Function invoked when calling the VisualCloze pipeline for generation.
Args:
- task_prompt (`str` or `List[str]`, *optional*):
+ task_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to define the task intention.
- content_prompt (`str` or `List[str]`, *optional*):
+ content_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to define the content or caption of the target image to be generated.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -297,7 +297,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -309,7 +309,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -336,7 +336,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py
index e12995106bcf..dd5d0603d6d0 100644
--- a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py
+++ b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import numpy as np
import torch
@@ -192,11 +192,11 @@ def __init__(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_images_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -242,9 +242,9 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
def _get_clip_prompt_embeds(
self,
- prompt: Union[str, List[str]],
+ prompt: str | list[str],
num_images_per_prompt: int = 1,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
@@ -287,25 +287,25 @@ def _get_clip_prompt_embeds(
# Modified from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
def encode_prompt(
self,
- layout_prompt: Union[str, List[str]],
- task_prompt: Union[str, List[str]],
- content_prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ layout_prompt: str | list[str],
+ task_prompt: str | list[str],
+ content_prompt: str | list[str],
+ device: torch.device | None = None,
num_images_per_prompt: int = 1,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
- lora_scale: Optional[float] = None,
+ lora_scale: float | None = None,
):
r"""
Args:
- layout_prompt (`str` or `List[str]`, *optional*):
+ layout_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to define the number of in-context examples and the number of images involved in
the task.
- task_prompt (`str` or `List[str]`, *optional*):
+ task_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to define the task intention.
- content_prompt (`str` or `List[str]`, *optional*):
+ content_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to define the content or caption of the target image to be generated.
device: (`torch.device`):
torch device
@@ -709,33 +709,33 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- task_prompt: Union[str, List[str]] = None,
- content_prompt: Union[str, List[str]] = None,
- image: Optional[torch.FloatTensor] = None,
+ task_prompt: str | list[str] = None,
+ content_prompt: str | list[str] = None,
+ image: torch.FloatTensor | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 30.0,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
- output_type: Optional[str] = "pil",
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the VisualCloze pipeline for generation.
Args:
- task_prompt (`str` or `List[str]`, *optional*):
+ task_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to define the task intention.
- content_prompt (`str` or `List[str]`, *optional*):
+ content_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to define the content or caption of the target image to be generated.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
@@ -743,7 +743,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -755,7 +755,7 @@ def __call__(
the text `prompt`, usually at the expense of lower image quality.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
@@ -782,7 +782,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/visualcloze/visualcloze_utils.py b/src/diffusers/pipelines/visualcloze/visualcloze_utils.py
index efe5dff47623..1a96ff8d2d13 100644
--- a/src/diffusers/pipelines/visualcloze/visualcloze_utils.py
+++ b/src/diffusers/pipelines/visualcloze/visualcloze_utils.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Dict, List, Optional, Tuple, Union
import torch
from PIL import Image
@@ -40,8 +39,8 @@ def __init__(self, *args, resolution: int = 384, **kwargs):
self.resolution = resolution
def preprocess_image(
- self, input_images: List[List[Optional[Image.Image]]], vae_scale_factor: int
- ) -> Tuple[List[List[torch.Tensor]], List[List[List[int]]], List[int]]:
+ self, input_images: list[list[Image.Image | None]], vae_scale_factor: int
+ ) -> tuple[list[list[torch.Tensor]], list[list[list[int]]], list[int]]:
"""
Preprocesses input images for the VisualCloze pipeline.
@@ -52,7 +51,7 @@ def preprocess_image(
4. Tracking image sizes and positions of target images
Args:
- input_images (List[List[Optional[Image.Image]]]):
+ input_images (list[list[Image.Image | None]]):
A nested list of PIL Images where:
- Outer list represents different samples, including in-context examples and the query
- Inner list contains images for the task
@@ -61,17 +60,17 @@ def preprocess_image(
The scale factor used by the VAE for resizing images
Returns:
- Tuple containing:
- - List[List[torch.Tensor]]: Preprocessed images in tensor format
- - List[List[List[int]]]: Dimensions of each processed image [height, width]
- - List[int]: Target positions indicating which images are to be generated
+ tuple containing:
+ - list[list[torch.Tensor]]: Preprocessed images in tensor format
+ - list[list[list[int]]]: Dimensions of each processed image [height, width]
+ - list[int]: Target positions indicating which images are to be generated
"""
n_samples, n_task_images = len(input_images), len(input_images[0])
divisible = 2 * vae_scale_factor
- processed_images: List[List[Image.Image]] = [[] for _ in range(n_samples)]
- resize_size: List[Optional[Tuple[int, int]]] = [None for _ in range(n_samples)]
- target_position: List[int] = []
+ processed_images: list[list[Image.Image]] = [[] for _ in range(n_samples)]
+ resize_size: list[tuple[int, int] | None] = [None for _ in range(n_samples)]
+ target_position: list[int] = []
# Process each sample
for i in range(n_samples):
@@ -125,19 +124,19 @@ def preprocess_image(
return processed_images, image_sizes, target_position
def preprocess_mask(
- self, input_images: List[List[Image.Image]], target_position: List[int]
- ) -> List[List[torch.Tensor]]:
+ self, input_images: list[list[Image.Image]], target_position: list[int]
+ ) -> list[list[torch.Tensor]]:
"""
Generate masks for the VisualCloze pipeline.
Args:
- input_images (List[List[Image.Image]]):
+ input_images (list[list[Image.Image]]):
Processed images from preprocess_image
- target_position (List[int]):
+ target_position (list[int]):
Binary list marking the positions of target images (1 for target, 0 for condition)
Returns:
- List[List[torch.Tensor]]:
+ list[list[torch.Tensor]]:
A nested list of mask tensors (1 for target positions, 0 for condition images)
"""
mask = []
@@ -155,10 +154,10 @@ def preprocess_mask(
def preprocess_image_upsampling(
self,
- input_images: List[List[Image.Image]],
+ input_images: list[list[Image.Image]],
height: int,
width: int,
- ) -> Tuple[List[List[Image.Image]], List[List[List[int]]]]:
+ ) -> tuple[list[list[Image.Image]], list[list[list[int]]]]:
"""Process images for the upsampling stage in the VisualCloze pipeline.
Args:
@@ -167,7 +166,7 @@ def preprocess_image_upsampling(
width: Target width
Returns:
- Tuple of processed image and its size
+ tuple of processed image and its size
"""
image = self.resize(input_images[0][0], height, width)
image = self.pil_to_numpy(image) # to np
@@ -178,10 +177,10 @@ def preprocess_image_upsampling(
image_sizes = [[[height, width]]]
return input_images, image_sizes
- def preprocess_mask_upsampling(self, input_images: List[List[Image.Image]]) -> List[List[torch.Tensor]]:
+ def preprocess_mask_upsampling(self, input_images: list[list[Image.Image]]) -> list[list[torch.Tensor]]:
return [[torch.ones((1, 1, input_images[0][0].shape[2], input_images[0][0].shape[3]))]]
- def get_layout_prompt(self, size: Tuple[int, int]) -> str:
+ def get_layout_prompt(self, size: tuple[int, int]) -> str:
layout_instruction = (
f"A grid layout with {size[0]} rows and {size[1]} columns, displaying {size[0] * size[1]} images arranged side by side.",
)
@@ -189,26 +188,26 @@ def get_layout_prompt(self, size: Tuple[int, int]) -> str:
def preprocess(
self,
- task_prompt: Union[str, List[str]],
- content_prompt: Union[str, List[str]],
- input_images: Optional[List[List[List[Optional[str]]]]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ task_prompt: str | list[str],
+ content_prompt: str | list[str],
+ input_images: list[list[list[str | None]]] | None = None,
+ height: int | None = None,
+ width: int | None = None,
upsampling: bool = False,
vae_scale_factor: int = 16,
- ) -> Dict:
+ ) -> dict:
"""Process visual cloze inputs.
Args:
task_prompt: Task description(s)
content_prompt: Content description(s)
- input_images: List of images or None for the target images
+ input_images: list of images or None for the target images
height: Optional target height for upsampling stage
width: Optional target width for upsampling stage
upsampling: Whether this is in the upsampling processing stage
Returns:
- Dictionary containing processed images, masks, prompts and metadata
+ dictionary containing processed images, masks, prompts and metadata
"""
if isinstance(task_prompt, str):
task_prompt = [task_prompt]
diff --git a/src/diffusers/pipelines/wan/image_processor.py b/src/diffusers/pipelines/wan/image_processor.py
index b1594d08630f..fa18150fcc6e 100644
--- a/src/diffusers/pipelines/wan/image_processor.py
+++ b/src/diffusers/pipelines/wan/image_processor.py
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Optional, Tuple, Union
import numpy as np
import PIL.Image
@@ -36,7 +35,7 @@ class WanAnimateImageProcessor(VaeImageProcessor):
this factor.
vae_latent_channels (`int`, *optional*, defaults to `16`):
VAE latent channels.
- spatial_patch_size (`Tuple[int, int]`, *optional*, defaults to `(2, 2)`):
+ spatial_patch_size (`tuple[int, int]`, *optional*, defaults to `(2, 2)`):
The spatial patch size used by the diffusion transformer. For Wan models, this is typically (2, 2).
resample (`str`, *optional*, defaults to `lanczos`):
Resampling filter to use when resizing the image.
@@ -48,7 +47,7 @@ class WanAnimateImageProcessor(VaeImageProcessor):
Whether to convert the images to RGB format.
do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
Whether to convert the images to grayscale format.
- fill_color (`str` or `float` or `Tuple[float, ...]`, *optional*, defaults to `None`):
+ fill_color (`str` or `float` or `tuple[float, ...]`, *optional*, defaults to `None`):
An optional fill color when `resize_mode` is set to `"fill"`. This will fill the empty space with that
color instead of filling with data from the image. Any valid `color` argument to `PIL.Image.new` is valid;
if `None`, will default to filling with data from `image`.
@@ -60,14 +59,14 @@ def __init__(
do_resize: bool = True,
vae_scale_factor: int = 8,
vae_latent_channels: int = 16,
- spatial_patch_size: Tuple[int, int] = (2, 2),
+ spatial_patch_size: tuple[int, int] = (2, 2),
resample: str = "lanczos",
reducing_gap: int = None,
do_normalize: bool = True,
do_binarize: bool = False,
do_convert_rgb: bool = False,
do_convert_grayscale: bool = False,
- fill_color: Optional[Union[str, float, Tuple[float, ...]]] = 0,
+ fill_color: str | float | tuple[float, ...] | None = 0,
):
super().__init__()
if do_convert_rgb and do_convert_grayscale:
@@ -134,25 +133,25 @@ def _resize_and_fill(
def get_default_height_width(
self,
- image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
- height: Optional[int] = None,
- width: Optional[int] = None,
- ) -> Tuple[int, int]:
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
+ height: int | None = None,
+ width: int | None = None,
+ ) -> tuple[int, int]:
r"""
Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`.
Args:
- image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`):
+ image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it
should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch
tensor, it should have shape `[batch, channels, height, width]`.
- height (`Optional[int]`, *optional*, defaults to `None`):
+ height (`int | None`, *optional*, defaults to `None`):
The height of the preprocessed image. If `None`, the height of the `image` input will be used.
- width (`Optional[int]`, *optional*, defaults to `None`):
+ width (`int | None`, *optional*, defaults to `None`):
The width of the preprocessed image. If `None`, the width of the `image` input will be used.
Returns:
- `Tuple[int, int]`:
+ `tuple[int, int]`:
A tuple containing the height and width, both resized to the nearest integer multiple of
`vae_scale_factor * spatial_patch_size`.
"""
diff --git a/src/diffusers/pipelines/wan/pipeline_output.py b/src/diffusers/pipelines/wan/pipeline_output.py
index 88907ad0f0a1..7c7f8a4f0ad9 100644
--- a/src/diffusers/pipelines/wan/pipeline_output.py
+++ b/src/diffusers/pipelines/wan/pipeline_output.py
@@ -11,8 +11,8 @@ class WanPipelineOutput(BaseOutput):
Output class for Wan pipelines.
Args:
- frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
- List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]):
+ list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
diff --git a/src/diffusers/pipelines/wan/pipeline_wan.py b/src/diffusers/pipelines/wan/pipeline_wan.py
index e77a3356b883..d4edff01ad66 100644
--- a/src/diffusers/pipelines/wan/pipeline_wan.py
+++ b/src/diffusers/pipelines/wan/pipeline_wan.py
@@ -13,7 +13,7 @@
# limitations under the License.
import html
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import regex as re
import torch
@@ -134,9 +134,9 @@ def __init__(
text_encoder: UMT5EncoderModel,
vae: AutoencoderKLWan,
scheduler: FlowMatchEulerDiscreteScheduler,
- transformer: Optional[WanTransformer3DModel] = None,
- transformer_2: Optional[WanTransformer3DModel] = None,
- boundary_ratio: Optional[float] = None,
+ transformer: WanTransformer3DModel | None = None,
+ transformer_2: WanTransformer3DModel | None = None,
+ boundary_ratio: float | None = None,
expand_timesteps: bool = False, # Wan2.2 ti2v
):
super().__init__()
@@ -157,11 +157,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -198,23 +198,23 @@ def _get_t5_prompt_embeds(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -329,10 +329,10 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -382,35 +382,33 @@ def attention_kwargs(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 480,
width: int = 832,
num_frames: int = 81,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- guidance_scale_2: Optional[float] = None,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ guidance_scale_2: float | None = None,
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds`
instead. Ignored when not using guidance (`guidance_scale` < `1`).
height (`int`, defaults to `480`):
@@ -434,7 +432,7 @@ def __call__(
and the pipeline's `boundary_ratio` are not None.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -457,7 +455,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate.py b/src/diffusers/pipelines/wan/pipeline_wan_animate.py
index c7c983b2f7d4..5806032c0142 100644
--- a/src/diffusers/pipelines/wan/pipeline_wan_animate.py
+++ b/src/diffusers/pipelines/wan/pipeline_wan_animate.py
@@ -14,7 +14,7 @@
import html
from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL
import regex as re
@@ -135,7 +135,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -234,11 +234,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -277,7 +277,7 @@ def _get_t5_prompt_embeds(
def encode_image(
self,
image: PipelineImageInput,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
image = self.image_processor(images=image, return_tensors="pt").to(device)
@@ -287,23 +287,23 @@ def encode_image(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -463,9 +463,9 @@ def get_i2v_mask(
latent_h: int,
latent_w: int,
mask_len: int = 1,
- mask_pixel_values: Optional[torch.Tensor] = None,
- dtype: Optional[torch.dtype] = None,
- device: Union[str, torch.device] = "cuda",
+ mask_pixel_values: torch.Tensor | None = None,
+ dtype: torch.dtype | None = None,
+ device: str | torch.device = "cuda",
) -> torch.Tensor:
# mask_pixel_values shape (if supplied): [B, C = 1, T, latent_h, latent_w]
if mask_pixel_values is None:
@@ -490,9 +490,9 @@ def prepare_reference_image_latents(
image: torch.Tensor,
batch_size: int = 1,
sample_mode: int = "argmax",
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
) -> torch.Tensor:
# image shape: (B, C, H, W) or (B, C, T, H, W)
dtype = dtype or self.vae.dtype
@@ -537,9 +537,9 @@ def prepare_reference_image_latents(
def prepare_prev_segment_cond_latents(
self,
- prev_segment_cond_video: Optional[torch.Tensor] = None,
- background_video: Optional[torch.Tensor] = None,
- mask_video: Optional[torch.Tensor] = None,
+ prev_segment_cond_video: torch.Tensor | None = None,
+ background_video: torch.Tensor | None = None,
+ mask_video: torch.Tensor | None = None,
batch_size: int = 1,
segment_frame_length: int = 77,
start_frame: int = 0,
@@ -549,9 +549,9 @@ def prepare_prev_segment_cond_latents(
task: str = "animate",
interpolation_mode: str = "bicubic",
sample_mode: str = "argmax",
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
) -> torch.Tensor:
# prev_segment_cond_video shape: (B, C, T, H, W) in pixel space if supplied
# background_video shape: (B, C, T, H, W) (same as prev_segment_cond_video shape)
@@ -655,9 +655,9 @@ def prepare_pose_latents(
pose_video: torch.Tensor,
batch_size: int = 1,
sample_mode: int = "argmax",
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
) -> torch.Tensor:
# pose_video shape: (B, C, T, H, W)
pose_video = pose_video.to(device=device, dtype=dtype if dtype is not None else self.vae.dtype)
@@ -689,11 +689,11 @@ def prepare_latents(
height: int = 720,
width: int = 1280,
num_frames: int = 77,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
latent_width = width // self.vae_scale_factor_spatial
@@ -712,7 +712,7 @@ def prepare_latents(
return latents
- def pad_video_frames(self, frames: List[Any], num_target_frames: int) -> List[Any]:
+ def pad_video_frames(self, frames: list[Any], num_target_frames: int) -> list[Any]:
"""
Pads an array-like video `frames` to `num_target_frames` using a "reflect"-like strategy. The frame dimension
is assumed to be the first dimension. In the 1D case, we can visualize this strategy as follows:
@@ -762,33 +762,31 @@ def attention_kwargs(self):
def __call__(
self,
image: PipelineImageInput,
- pose_video: List[PIL.Image.Image],
- face_video: List[PIL.Image.Image],
- background_video: Optional[List[PIL.Image.Image]] = None,
- mask_video: Optional[List[PIL.Image.Image]] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ pose_video: list[PIL.Image.Image],
+ face_video: list[PIL.Image.Image],
+ background_video: list[PIL.Image.Image] | None = None,
+ mask_video: list[PIL.Image.Image] | None = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 720,
width: int = 1280,
segment_frame_length: int = 77,
num_inference_steps: int = 20,
mode: str = "animate",
prev_segment_conditioning_frames: int = 1,
- motion_encode_batch_size: Optional[int] = None,
+ motion_encode_batch_size: int | None = None,
guidance_scale: float = 1.0,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- image_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ image_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
@@ -798,20 +796,20 @@ def __call__(
image (`PipelineImageInput`):
The input character image to condition the generation on. Must be an image, a list of images or a
`torch.Tensor`.
- pose_video (`List[PIL.Image.Image]`):
+ pose_video (`list[PIL.Image.Image]`):
The input pose video to condition the generation on. Must be a list of PIL images.
- face_video (`List[PIL.Image.Image]`):
+ face_video (`list[PIL.Image.Image]`):
The input face video to condition the generation on. Must be a list of PIL images.
- background_video (`List[PIL.Image.Image]`, *optional*):
+ background_video (`list[PIL.Image.Image]`, *optional*):
When mode is `"replace"`, the input background video to condition the generation on. Must be a list of
PIL images.
- mask_video (`List[PIL.Image.Image]`, *optional*):
+ mask_video (`list[PIL.Image.Image]`, *optional*):
When mode is `"replace"`, the input mask video to condition the generation on. Must be a list of PIL
images.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -844,7 +842,7 @@ def __call__(
Animate inference.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
diff --git a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py
index c1c4a92c2c39..f669e9b1d0ec 100644
--- a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py
+++ b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py
@@ -13,7 +13,7 @@
# limitations under the License.
import html
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
import PIL
import regex as re
@@ -112,7 +112,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -174,7 +174,7 @@ def __init__(
image_encoder: CLIPVisionModel = None,
transformer: WanTransformer3DModel = None,
transformer_2: WanTransformer3DModel = None,
- boundary_ratio: Optional[float] = None,
+ boundary_ratio: float | None = None,
expand_timesteps: bool = False,
):
super().__init__()
@@ -198,11 +198,11 @@ def __init__(
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 512,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -240,7 +240,7 @@ def _get_t5_prompt_embeds(
def encode_image(
self,
image: PipelineImageInput,
- device: Optional[torch.device] = None,
+ device: torch.device | None = None,
):
device = device or self._execution_device
image = self.image_processor(images=image, return_tensors="pt").to(device)
@@ -250,23 +250,23 @@ def encode_image(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -398,12 +398,12 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- last_image: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ last_image: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
latent_width = width // self.vae_scale_factor_spatial
@@ -509,28 +509,26 @@ def attention_kwargs(self):
def __call__(
self,
image: PipelineImageInput,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 480,
width: int = 832,
num_frames: int = 81,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- guidance_scale_2: Optional[float] = None,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- image_embeds: Optional[torch.Tensor] = None,
- last_image: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ guidance_scale_2: float | None = None,
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ image_embeds: torch.Tensor | None = None,
+ last_image: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
@@ -539,10 +537,10 @@ def __call__(
Args:
image (`PipelineImageInput`):
The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -567,7 +565,7 @@ def __call__(
and the pipeline's `boundary_ratio` are not None.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -596,7 +594,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/wan/pipeline_wan_vace.py b/src/diffusers/pipelines/wan/pipeline_wan_vace.py
index 351ae2e70563..c016eec1b535 100644
--- a/src/diffusers/pipelines/wan/pipeline_wan_vace.py
+++ b/src/diffusers/pipelines/wan/pipeline_wan_vace.py
@@ -13,7 +13,7 @@
# limitations under the License.
import html
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL.Image
import regex as re
@@ -126,7 +126,7 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -183,7 +183,7 @@ def __init__(
scheduler: FlowMatchEulerDiscreteScheduler,
transformer: WanVACETransformer3DModel = None,
transformer_2: WanVACETransformer3DModel = None,
- boundary_ratio: Optional[float] = None,
+ boundary_ratio: float | None = None,
):
super().__init__()
@@ -203,11 +203,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -245,23 +245,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -413,15 +413,15 @@ def check_inputs(
def preprocess_conditions(
self,
- video: Optional[List[PipelineImageInput]] = None,
- mask: Optional[List[PipelineImageInput]] = None,
- reference_images: Optional[Union[PIL.Image.Image, List[PIL.Image.Image], List[List[PIL.Image.Image]]]] = None,
+ video: list[PipelineImageInput] | None = None,
+ mask: list[PipelineImageInput] | None = None,
+ reference_images: PIL.Image.Image | list[PIL.Image.Image] | list[list[PIL.Image.Image]] | None = None,
batch_size: int = 1,
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
):
if video is not None:
base = self.vae_scale_factor_spatial * (
@@ -515,9 +515,9 @@ def prepare_video_latents(
self,
video: torch.Tensor,
mask: torch.Tensor,
- reference_images: Optional[List[List[torch.Tensor]]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- device: Optional[torch.device] = None,
+ reference_images: list[list[torch.Tensor]] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ device: torch.device | None = None,
) -> torch.Tensor:
device = device or self._execution_device
@@ -581,8 +581,8 @@ def prepare_video_latents(
def prepare_masks(
self,
mask: torch.Tensor,
- reference_images: Optional[List[torch.Tensor]] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ reference_images: list[torch.Tensor] | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
) -> torch.Tensor:
if isinstance(generator, list):
# TODO: support this
@@ -637,10 +637,10 @@ def prepare_latents(
height: int = 480,
width: int = 832,
num_frames: int = 81,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
@@ -690,60 +690,58 @@ def attention_kwargs(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
- video: Optional[List[PipelineImageInput]] = None,
- mask: Optional[List[PipelineImageInput]] = None,
- reference_images: Optional[List[PipelineImageInput]] = None,
- conditioning_scale: Union[float, List[float], torch.Tensor] = 1.0,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
+ video: list[PipelineImageInput] | None = None,
+ mask: list[PipelineImageInput] | None = None,
+ reference_images: list[PipelineImageInput] | None = None,
+ conditioning_scale: float | list[float] | torch.Tensor = 1.0,
height: int = 480,
width: int = 832,
num_frames: int = 81,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
- guidance_scale_2: Optional[float] = None,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ guidance_scale_2: float | None = None,
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`
instead.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
- video (`List[PIL.Image.Image]`, *optional*):
+ video (`list[PIL.Image.Image]`, *optional*):
The input video or videos to be used as a starting point for the generation. The video should be a list
of PIL images, a numpy array, or a torch tensor. Currently, the pipeline only supports generating one
video at a time.
- mask (`List[PIL.Image.Image]`, *optional*):
+ mask (`list[PIL.Image.Image]`, *optional*):
The input mask defines which video regions to condition on and which to generate. Black areas in the
mask indicate conditioning regions, while white areas indicate regions for generation. The mask should
be a list of PIL images, a numpy array, or a torch tensor. Currently supports generating a single video
at a time.
- reference_images (`List[PIL.Image.Image]`, *optional*):
+ reference_images (`list[PIL.Image.Image]`, *optional*):
A list of one or more reference images as extra conditioning for the generation. For example, if you
are trying to inpaint a video to change the character, you can pass reference images of the new
character here. Refer to the Diffusers [examples](https://github.com/huggingface/diffusers/pull/11582)
and original [user
guide](https://github.com/ali-vilab/VACE/blob/0897c6d055d7d9ea9e191dce763006664d9780f8/UserGuide.md)
for a full list of supported tasks and use cases.
- conditioning_scale (`float`, `List[float]`, `torch.Tensor`, defaults to `1.0`):
+ conditioning_scale (`float`, `list[float]`, `torch.Tensor`, defaults to `1.0`):
The conditioning scale to be applied when adding the control conditioning latent stream to the
denoising latent stream in each control layer of the model. If a float is provided, it will be applied
uniformly to all layers. If a list or tensor is provided, it should have the same length as the number
@@ -769,7 +767,7 @@ def __call__(
and the pipeline's `boundary_ratio` are not None.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -792,7 +790,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/wan/pipeline_wan_video2video.py b/src/diffusers/pipelines/wan/pipeline_wan_video2video.py
index 5475b6e8b479..3d7c5297f4c4 100644
--- a/src/diffusers/pipelines/wan/pipeline_wan_video2video.py
+++ b/src/diffusers/pipelines/wan/pipeline_wan_video2video.py
@@ -14,7 +14,7 @@
import html
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import regex as re
import torch
@@ -100,10 +100,10 @@ def prompt_clean(text):
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -118,15 +118,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -159,7 +159,7 @@ def retrieve_timesteps(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -221,11 +221,11 @@ def __init__(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
def _get_t5_prompt_embeds(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
@@ -263,23 +263,23 @@ def _get_t5_prompt_embeds(
# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
max_sequence_length: int = 226,
- device: Optional[torch.device] = None,
- dtype: Optional[torch.dtype] = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
@@ -390,16 +390,16 @@ def check_inputs(
def prepare_latents(
self,
- video: Optional[torch.Tensor] = None,
+ video: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 16,
height: int = 480,
width: int = 832,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- generator: Optional[torch.Generator] = None,
- latents: Optional[torch.Tensor] = None,
- timestep: Optional[torch.Tensor] = None,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | None = None,
+ latents: torch.Tensor | None = None,
+ timestep: torch.Tensor | None = None,
):
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -480,34 +480,32 @@ def attention_kwargs(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- video: List[Image.Image] = None,
- prompt: Union[str, List[str]] = None,
- negative_prompt: Union[str, List[str]] = None,
+ video: list[Image.Image] = None,
+ prompt: str | list[str] = None,
+ negative_prompt: str | list[str] = None,
height: int = 480,
width: int = 832,
num_inference_steps: int = 50,
- timesteps: Optional[List[int]] = None,
+ timesteps: list[int] | None = None,
guidance_scale: float = 5.0,
strength: float = 0.8,
- num_videos_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "np",
+ num_videos_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ output_type: str | None = "np",
return_dict: bool = True,
- attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[
- Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
- ] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
The call function to the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`
instead.
height (`int`, defaults to `480`):
@@ -529,7 +527,7 @@ def __call__(
Higher strength leads to more differences between original image and generated video.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -552,7 +550,7 @@ def __call__(
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py b/src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py
index 5ab206b15176..932c7ac618f6 100644
--- a/src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py
+++ b/src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py
@@ -13,7 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Union
import torch
import torch.nn as nn
@@ -142,7 +141,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput:
@apply_forward_hook
def decode(
self, h: torch.Tensor, force_not_quantize: bool = True, return_dict: bool = True
- ) -> Union[DecoderOutput, torch.Tensor]:
+ ) -> DecoderOutput | torch.Tensor:
if not force_not_quantize:
quant, _, _ = self.vquantizer(h)
else:
@@ -155,7 +154,7 @@ def decode(
return DecoderOutput(sample=dec)
- def forward(self, sample: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ def forward(self, sample: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
r"""
Args:
sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py
index bbdb60471fd1..cce05c189201 100644
--- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py
+++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -156,7 +156,7 @@ def encode_prompt(
uncond_text_encoder_hidden_states = None
if do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -217,33 +217,33 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image_embeddings: Union[torch.Tensor, List[torch.Tensor]],
- prompt: Union[str, List[str]] = None,
+ image_embeddings: torch.Tensor | list[torch.Tensor],
+ prompt: str | list[str] = None,
num_inference_steps: int = 12,
- timesteps: Optional[List[float]] = None,
+ timesteps: list[float] | None = None,
guidance_scale: float = 0.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: str | list[str] | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- image_embedding (`torch.Tensor` or `List[torch.Tensor]`):
+ image_embedding (`torch.Tensor` or `list[torch.Tensor]`):
Image Embeddings either extracted from an image or generated by a Prior Model.
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
num_inference_steps (`int`, *optional*, defaults to 12):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 0.0):
@@ -252,12 +252,12 @@ def __call__(
equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by
setting `decoder_guidance_scale > 1`. Higher guidance scale encourages to generate images that are
closely linked to the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `decoder_guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -274,7 +274,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py
index c54c1fefe8fe..16300a7c71d2 100644
--- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py
+++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py
@@ -11,7 +11,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.
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import torch
from transformers import CLIPTextModel, CLIPTokenizer
@@ -110,10 +110,10 @@ def __init__(
vqgan=vqgan,
)
- def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
+ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)
- def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
@@ -123,7 +123,7 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t
self.prior_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device)
self.decoder_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device)
- def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None):
+ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
r"""
Offloads all models (`unet`, `text_encoder`, `vae`, and `safety checker` state dicts) to CPU using 🤗
Accelerate, significantly reducing memory usage. Models are moved to a `torch.device('meta')` and loaded on a
@@ -145,36 +145,36 @@ def set_progress_bar_config(self, **kwargs):
@replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] | None = None,
height: int = 512,
width: int = 512,
prior_num_inference_steps: int = 60,
- prior_timesteps: Optional[List[float]] = None,
+ prior_timesteps: list[float] | None = None,
prior_guidance_scale: float = 4.0,
num_inference_steps: int = 12,
- decoder_timesteps: Optional[List[float]] = None,
+ decoder_timesteps: list[float] | None = None,
decoder_guidance_scale: float = 0.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
num_images_per_prompt: int = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pil",
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ prior_callback_on_step_end: Callable[[int, int], None] | None = None,
+ prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation for the prior and decoder.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `guidance_scale` is less than `1`).
prompt_embeds (`torch.Tensor`, *optional*):
@@ -196,7 +196,7 @@ def __call__(
equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by
setting `prior_guidance_scale > 1`. Higher guidance scale encourages to generate images that are
closely linked to the text `prompt`, usually at the expense of lower image quality.
- prior_num_inference_steps (`Union[int, Dict[float, int]]`, *optional*, defaults to 60):
+ prior_num_inference_steps (`int | dict[float, int]`, *optional*, defaults to 60):
The number of prior denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference. For more specific timestep spacing, you can pass customized
`prior_timesteps`
@@ -204,10 +204,10 @@ def __call__(
The number of decoder denoising steps. More denoising steps usually lead to a higher quality image at
the expense of slower inference. For more specific timestep spacing, you can pass customized
`timesteps`
- prior_timesteps (`List[float]`, *optional*):
+ prior_timesteps (`list[float]`, *optional*):
Custom timesteps to use for the denoising process for the prior. If not defined, equal spaced
`prior_num_inference_steps` timesteps are used. Must be in descending order.
- decoder_timesteps (`List[float]`, *optional*):
+ decoder_timesteps (`list[float]`, *optional*):
Custom timesteps to use for the denoising process for the decoder. If not defined, equal spaced
`num_inference_steps` timesteps are used. Must be in descending order.
decoder_guidance_scale (`float`, *optional*, defaults to 0.0):
@@ -216,7 +216,7 @@ def __call__(
of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
`guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
the text `prompt`, usually at the expense of lower image quality.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -232,7 +232,7 @@ def __call__(
A function that calls at the end of each denoising steps during the inference. The function is called
with the following arguments: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep:
int, callback_kwargs: Dict)`.
- prior_callback_on_step_end_tensor_inputs (`List`, *optional*):
+ prior_callback_on_step_end_tensor_inputs (`list`, *optional*):
The list of tensor inputs for the `prior_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.
@@ -241,7 +241,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py
index e138b6e805c8..e79fcf8378aa 100644
--- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py
+++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py
@@ -14,7 +14,7 @@
from dataclasses import dataclass
from math import ceil
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable
import numpy as np
import torch
@@ -67,7 +67,7 @@ class WuerstchenPriorPipelineOutput(BaseOutput):
"""
- image_embeddings: Union[torch.Tensor, np.ndarray]
+ image_embeddings: torch.Tensor | np.ndarray
class WuerstchenPriorPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin):
@@ -145,8 +145,8 @@ def encode_prompt(
do_classifier_free_guidance,
prompt=None,
negative_prompt=None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
):
if prompt is not None and isinstance(prompt, str):
batch_size = 1
@@ -191,7 +191,7 @@ def encode_prompt(
prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
if negative_prompt_embeds is None and do_classifier_free_guidance:
- uncond_tokens: List[str]
+ uncond_tokens: list[str]
if negative_prompt is None:
uncond_tokens = [""] * batch_size
elif type(prompt) is not type(negative_prompt):
@@ -290,29 +290,29 @@ def num_timesteps(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Optional[Union[str, List[str]]] = None,
+ prompt: str | list[str] | None = None,
height: int = 1024,
width: int = 1024,
num_inference_steps: int = 60,
- timesteps: List[float] = None,
+ timesteps: list[float] = None,
guidance_scale: float = 8.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[torch.Tensor] = None,
- negative_prompt_embeds: Optional[torch.Tensor] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.Tensor] = None,
- output_type: Optional[str] = "pt",
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ output_type: str | None = "pt",
return_dict: bool = True,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
**kwargs,
):
"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`):
+ prompt (`str` or `list[str]`):
The prompt or prompts to guide the image generation.
height (`int`, *optional*, defaults to 1024):
The height in pixels of the generated image.
@@ -321,7 +321,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 60):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 8.0):
@@ -330,7 +330,7 @@ def __call__(
equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by
setting `decoder_guidance_scale > 1`. Higher guidance scale encourages to generate images that are
closely linked to the text `prompt`, usually at the expense of lower image quality.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
if `decoder_guidance_scale` is less than `1`).
prompt_embeds (`torch.Tensor`, *optional*):
@@ -342,7 +342,7 @@ def __call__(
argument.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.Tensor`, *optional*):
@@ -359,7 +359,7 @@ def __call__(
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
+ 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.
diff --git a/src/diffusers/pipelines/z_image/pipeline_output.py b/src/diffusers/pipelines/z_image/pipeline_output.py
index 69a320fc036a..5cbcbee4bfab 100644
--- a/src/diffusers/pipelines/z_image/pipeline_output.py
+++ b/src/diffusers/pipelines/z_image/pipeline_output.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -27,9 +26,9 @@ class ZImagePipelineOutput(BaseOutput):
Output class for Z-Image pipelines.
Args:
- images (`List[PIL.Image.Image]` or `np.ndarray`)
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image, np.ndarray]
diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py
index 82bdd7d361b7..31dac0e5c283 100644
--- a/src/diffusers/pipelines/z_image/pipeline_z_image.py
+++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import AutoTokenizer, PreTrainedModel
@@ -77,10 +77,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -95,15 +95,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -163,12 +163,12 @@ def __init__(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -197,11 +197,11 @@ def encode_prompt(
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
max_sequence_length: int = 512,
- ) -> List[torch.FloatTensor]:
+ ) -> list[torch.FloatTensor]:
device = device or self._execution_device
if prompt_embeds is not None:
@@ -294,32 +294,32 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
cfg_normalization: bool = False,
cfg_truncation: float = 1.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: list[torch.FloatTensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to 1024):
@@ -329,7 +329,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -343,23 +343,23 @@ def __call__(
Whether to apply configuration normalization.
cfg_truncation (`float`, *optional*, defaults to 1.0):
The truncation value for configuration.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will be generated by sampling using the supplied random `generator`.
- prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py
index 3c8db4a0f748..288aee039ed8 100644
--- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py
+++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import AutoTokenizer, PreTrainedModel
@@ -113,7 +113,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -128,10 +128,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -146,15 +146,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -217,12 +217,12 @@ def __init__(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -251,11 +251,11 @@ def encode_prompt(
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
max_sequence_length: int = 512,
- ) -> List[torch.FloatTensor]:
+ ) -> list[torch.FloatTensor]:
device = device or self._execution_device
if prompt_embeds is not None:
@@ -383,34 +383,34 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
control_image: PipelineImageInput = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 0.75,
+ controlnet_conditioning_scale: float | list[float] = 0.75,
cfg_normalization: bool = False,
cfg_truncation: float = 1.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: list[torch.FloatTensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to 1024):
@@ -420,7 +420,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -434,23 +434,23 @@ def __call__(
Whether to apply configuration normalization.
cfg_truncation (`float`, *optional*, defaults to 1.0):
The truncation value for configuration.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will be generated by sampling using the supplied random `generator`.
- prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py
index cdc60eaf4dd3..d278da912b5a 100644
--- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py
+++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
import torch.nn.functional as F
@@ -113,7 +113,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -128,10 +128,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -146,15 +146,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -224,12 +224,12 @@ def __init__(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -258,11 +258,11 @@ def encode_prompt(
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
max_sequence_length: int = 512,
- ) -> List[torch.FloatTensor]:
+ ) -> list[torch.FloatTensor]:
device = device or self._execution_device
if prompt_embeds is not None:
@@ -390,36 +390,36 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
image: PipelineImageInput = None,
mask_image: PipelineImageInput = None,
control_image: PipelineImageInput = None,
- controlnet_conditioning_scale: Union[float, List[float]] = 0.75,
+ controlnet_conditioning_scale: float | list[float] = 0.75,
cfg_normalization: bool = False,
cfg_truncation: float = 1.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: list[torch.FloatTensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to 1024):
@@ -429,7 +429,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -443,23 +443,23 @@ def __call__(
Whether to apply configuration normalization.
cfg_truncation (`float`, *optional*, defaults to 1.0):
The truncation value for configuration.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will be generated by sampling using the supplied random `generator`.
- prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py b/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py
index 2b3e80a2082b..94b98a5f8580 100644
--- a/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py
+++ b/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import torch
from transformers import AutoTokenizer, PreTrainedModel
@@ -74,7 +74,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -89,10 +89,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -107,15 +107,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -192,12 +192,12 @@ def __init__(
# Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -227,11 +227,11 @@ def encode_prompt(
# Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline._encode_prompt
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
max_sequence_length: int = 512,
- ) -> List[torch.FloatTensor]:
+ ) -> list[torch.FloatTensor]:
device = device or self._execution_device
if prompt_embeds is not None:
@@ -365,37 +365,37 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- prompt: Union[str, List[str]] = None,
+ prompt: str | list[str] = None,
image: PipelineImageInput = None,
strength: float = 0.6,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
cfg_normalization: bool = False,
cfg_truncation: float = 1.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: list[torch.FloatTensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for image-to-image generation.
Args:
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]`. If it's a tensor or a
list of tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or
@@ -413,7 +413,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -427,23 +427,23 @@ def __call__(
Whether to apply configuration normalization.
cfg_truncation (`float`, *optional*, defaults to 1.0):
The truncation value for configuration.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will be generated by sampling using the supplied random `generator`.
- prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py
index de9fb3c82f8c..2f842182edc6 100644
--- a/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py
+++ b/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py
@@ -90,7 +90,7 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
- encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
@@ -105,10 +105,10 @@ def retrieve_latents(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -123,15 +123,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -214,12 +214,12 @@ def __init__(
# Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -249,11 +249,11 @@ def encode_prompt(
# Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline._encode_prompt
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
max_sequence_length: int = 512,
- ) -> List[torch.FloatTensor]:
+ ) -> list[torch.FloatTensor]:
device = device or self._execution_device
if prompt_embeds is not None:
@@ -536,20 +536,20 @@ def __call__(
mask_image: PipelineImageInput = None,
masked_image_latents: Optional[torch.FloatTensor] = None,
strength: float = 1.0,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
cfg_normalization: bool = False,
cfg_truncation: float = 1.0,
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[List[torch.FloatTensor]] = None,
negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ output_type: str = "pil",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py b/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py
index 26848bea0a9e..a0c3ec03ef80 100644
--- a/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py
+++ b/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py
@@ -13,7 +13,7 @@
# limitations under the License.
import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
import PIL
import torch
@@ -78,10 +78,10 @@ def calculate_shift(
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
**kwargs,
):
r"""
@@ -96,15 +96,15 @@ def retrieve_timesteps(
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
@@ -169,12 +169,12 @@ def __init__(
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
do_classifier_free_guidance: bool = True,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt: str | list[str] | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: torch.FloatTensor | None = None,
max_sequence_length: int = 512,
num_condition_images: int = 0,
):
@@ -206,12 +206,12 @@ def encode_prompt(
def _encode_prompt(
self,
- prompt: Union[str, List[str]],
- device: Optional[torch.device] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
max_sequence_length: int = 512,
num_condition_images: int = 0,
- ) -> List[torch.FloatTensor]:
+ ) -> list[torch.FloatTensor]:
device = device or self._execution_device
if prompt_embeds is not None:
@@ -292,7 +292,7 @@ def prepare_latents(
def prepare_image_latents(
self,
- images: List[torch.Tensor],
+ images: list[torch.Tensor],
batch_size,
device,
dtype,
@@ -313,7 +313,7 @@ def prepare_image_latents(
def prepare_siglip_embeds(
self,
- images: List[torch.Tensor],
+ images: list[torch.Tensor],
batch_size,
device,
dtype,
@@ -357,39 +357,39 @@ def interrupt(self):
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
- image: Optional[Union[List[PIL.Image.Image], PIL.Image.Image]] = None,
- prompt: Union[str, List[str]] = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ image: list[PIL.Image.Image, PIL.Image.Image] | None = None,
+ prompt: str | list[str] = None,
+ height: int | None = None,
+ width: int | None = None,
num_inference_steps: int = 50,
- sigmas: Optional[List[float]] = None,
+ sigmas: list[float] | None = None,
guidance_scale: float = 5.0,
cfg_normalization: bool = False,
cfg_truncation: float = 1.0,
- negative_prompt: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
+ negative_prompt: str | list[str] | None = None,
+ num_images_per_prompt: int | None = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.FloatTensor | None = None,
+ prompt_embeds: list[torch.FloatTensor] | None = None,
+ negative_prompt_embeds: list[torch.FloatTensor] | None = None,
+ output_type: str | None = "pil",
return_dict: bool = True,
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
- image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`):
`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
latents as `image`, but if passing latents directly it is not encoded again.
- prompt (`str` or `List[str]`, *optional*):
+ prompt (`str` or `list[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to 1024):
@@ -399,7 +399,7 @@ def __call__(
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
@@ -413,23 +413,23 @@ def __call__(
Whether to apply configuration normalization.
cfg_truncation (`float`, *optional*, defaults to 1.0):
The truncation value for configuration.
- negative_prompt (`str` or `List[str]`, *optional*):
+ negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will be generated by sampling using the supplied random `generator`.
- prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+ negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py
index 070bcd0b2151..6cd24c459c9d 100644
--- a/src/diffusers/quantizers/auto.py
+++ b/src/diffusers/quantizers/auto.py
@@ -17,7 +17,6 @@
"""
import warnings
-from typing import Dict, Optional, Union
from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer
from .gguf import GGUFQuantizer
@@ -61,7 +60,7 @@ class DiffusersAutoQuantizer:
"""
@classmethod
- def from_dict(cls, quantization_config_dict: Dict):
+ def from_dict(cls, quantization_config_dict: dict):
quant_method = quantization_config_dict.get("quant_method", None)
# We need a special care for bnb models to make sure everything is BC ..
if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False):
@@ -82,7 +81,7 @@ def from_dict(cls, quantization_config_dict: Dict):
return target_cls.from_dict(quantization_config_dict)
@classmethod
- def from_config(cls, quantization_config: Union[QuantizationConfigMixin, Dict], **kwargs):
+ def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwargs):
# Convert it to a QuantizationConfig if the q_config is a dict
if isinstance(quantization_config, dict):
quantization_config = cls.from_dict(quantization_config)
@@ -123,8 +122,8 @@ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
@classmethod
def merge_quantization_configs(
cls,
- quantization_config: Union[dict, QuantizationConfigMixin],
- quantization_config_from_args: Optional[QuantizationConfigMixin],
+ quantization_config: dict | QuantizationConfigMixin,
+ quantization_config_from_args: QuantizationConfigMixin | None,
):
"""
handles situations where both quantization_config from args and quantization_config from model config are
diff --git a/src/diffusers/quantizers/base.py b/src/diffusers/quantizers/base.py
index 24fc724b4c88..b0988284b648 100644
--- a/src/diffusers/quantizers/base.py
+++ b/src/diffusers/quantizers/base.py
@@ -18,7 +18,7 @@
"""
from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any
from ..utils import is_torch_available
from .quantization_config import QuantizationConfigMixin
@@ -40,9 +40,9 @@ class DiffusersQuantizer(ABC):
Attributes
quantization_config (`diffusers.quantizers.quantization_config.QuantizationConfigMixin`):
The quantization config that defines the quantization parameters of your model that you want to quantize.
- modules_to_not_convert (`List[str]`, *optional*):
+ modules_to_not_convert (`list[str]`, *optional*):
The list of module names to not convert when quantizing the model.
- required_packages (`List[str]`, *optional*):
+ required_packages (`list[str]`, *optional*):
The list of required pip packages to install prior to using the quantizer
requires_calibration (`bool`):
Whether the quantization method requires to calibrate the model before using it.
@@ -76,14 +76,14 @@ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
"""
return torch_dtype
- def update_device_map(self, device_map: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+ def update_device_map(self, device_map: dict[str, Any] | None) -> dict[str, Any] | None:
"""
Override this method if you want to pass a override the existing device map with a new one. E.g. for
bitsandbytes, since `accelerate` is a hard requirement, if no device_map is passed, the device_map is set to
`"auto"``
Args:
- device_map (`Union[dict, str]`, *optional*):
+ device_map (`dict | str`, *optional*):
The device_map that is passed through the `from_pretrained` method.
"""
return device_map
@@ -100,17 +100,17 @@ def adjust_target_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
"""
return torch_dtype
- def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]:
+ def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
"""
Override this method if you want to adjust the `missing_keys`.
Args:
- missing_keys (`List[str]`, *optional*):
+ missing_keys (`list[str]`, *optional*):
The list of missing keys in the checkpoint compared to the state dict of the model
"""
return missing_keys
- def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> Dict[str, "torch.dtype"]:
+ def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> dict[str, "torch.dtype"]:
"""
returns dtypes for modules that are not quantized - used for the computation of the device_map in case one
passes a str as a device_map. The method will use the `modules_to_not_convert` that is modified in
@@ -130,7 +130,7 @@ def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> Dict[s
if any(m in name for m in self.modules_to_not_convert)
}
- def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
"""adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization"""
return max_memory
@@ -139,7 +139,7 @@ def check_if_quantized_param(
model: "ModelMixin",
param_value: "torch.Tensor",
param_name: str,
- state_dict: Dict[str, Any],
+ state_dict: dict[str, Any],
**kwargs,
) -> bool:
"""
diff --git a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py
index 0dfdff019b79..16ff0d83b8c4 100644
--- a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py
+++ b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py
@@ -16,7 +16,7 @@
https://github.com/huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339dbc020fb3b/src/transformers/quantizers/quantizer_bnb_4bit.py
"""
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any
from ...utils import get_module_from_name
from ..base import DiffusersQuantizer
@@ -111,7 +111,7 @@ def check_if_quantized_param(
model: "ModelMixin",
param_value: "torch.Tensor",
param_name: str,
- state_dict: Dict[str, Any],
+ state_dict: dict[str, Any],
**kwargs,
) -> bool:
import bitsandbytes as bnb
@@ -133,8 +133,8 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: Dict[str, Any],
- unexpected_keys: Optional[List[str]] = None,
+ state_dict: dict[str, Any],
+ unexpected_keys: list[str] | None = None,
**kwargs,
):
import bitsandbytes as bnb
@@ -218,7 +218,7 @@ def check_quantized_param_shape(self, param_name, current_param, loaded_param):
else:
return True
- def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
# need more space for buffers that are created during quantization
max_memory = {key: val * 0.90 for key, val in max_memory.items()}
return max_memory
@@ -255,7 +255,7 @@ def _process_model_before_weight_loading(
self,
model: "ModelMixin",
device_map,
- keep_in_fp32_modules: List[str] = [],
+ keep_in_fp32_modules: list[str] = [],
**kwargs,
):
from .utils import replace_with_bnb_linear
@@ -387,7 +387,7 @@ def validate_environment(self, *args, **kwargs):
)
# Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer.adjust_max_memory
- def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
# need more space for buffers that are created during quantization
max_memory = {key: val * 0.90 for key, val in max_memory.items()}
return max_memory
@@ -432,7 +432,7 @@ def check_if_quantized_param(
model: "ModelMixin",
param_value: "torch.Tensor",
param_name: str,
- state_dict: Dict[str, Any],
+ state_dict: dict[str, Any],
**kwargs,
):
import bitsandbytes as bnb
@@ -455,8 +455,8 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: Dict[str, Any],
- unexpected_keys: Optional[List[str]] = None,
+ state_dict: dict[str, Any],
+ unexpected_keys: list[str] | None = None,
**kwargs,
):
import bitsandbytes as bnb
@@ -513,7 +513,7 @@ def _process_model_before_weight_loading(
self,
model: "ModelMixin",
device_map,
- keep_in_fp32_modules: List[str] = [],
+ keep_in_fp32_modules: list[str] = [],
**kwargs,
):
from .utils import replace_with_bnb_linear
diff --git a/src/diffusers/quantizers/bitsandbytes/utils.py b/src/diffusers/quantizers/bitsandbytes/utils.py
index 429aabb8fae6..aea4e7dda57f 100644
--- a/src/diffusers/quantizers/bitsandbytes/utils.py
+++ b/src/diffusers/quantizers/bitsandbytes/utils.py
@@ -18,7 +18,6 @@
import inspect
from inspect import signature
-from typing import Union
from ...utils import is_accelerate_available, is_bitsandbytes_available, is_torch_available, logging
from ..quantization_config import QuantizationMethod
@@ -128,10 +127,10 @@ def replace_with_bnb_linear(model, modules_to_not_convert=None, current_key_name
Parameters:
model (`torch.nn.Module`):
Input model or `torch.nn.Module` as the function is run recursively.
- modules_to_not_convert (`List[`str`]`, *optional*, defaults to `[]`):
+ modules_to_not_convert (`list[`str`]`, *optional*, defaults to `[]`):
Names of the modules to not convert in `Linear8bitLt`. In practice we keep the `modules_to_not_convert` in
full precision for numerical stability reasons.
- current_key_name (`List[`str`]`, *optional*):
+ current_key_name (`list[`str`]`, *optional*):
An array to track the current key of the recursion. This is used to check whether the current key (part of
it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or
`disk`).
@@ -305,7 +304,7 @@ def dequantize_and_replace(
return model
-def _check_bnb_status(module) -> Union[bool, bool]:
+def _check_bnb_status(module) -> bool | bool:
is_loaded_in_4bit_bnb = (
hasattr(module, "is_loaded_in_4bit")
and module.is_loaded_in_4bit
diff --git a/src/diffusers/quantizers/gguf/gguf_quantizer.py b/src/diffusers/quantizers/gguf/gguf_quantizer.py
index aa5ebf5711a3..15f39dd9605e 100644
--- a/src/diffusers/quantizers/gguf/gguf_quantizer.py
+++ b/src/diffusers/quantizers/gguf/gguf_quantizer.py
@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
from ..base import DiffusersQuantizer
@@ -41,7 +43,7 @@ def __init__(self, quantization_config, **kwargs):
self.compute_dtype = quantization_config.compute_dtype
self.pre_quantized = quantization_config.pre_quantized
- self.modules_to_not_convert = quantization_config.modules_to_not_convert
+ self.modules_to_not_convert = quantization_config.modules_to_not_convert or []
if not isinstance(self.modules_to_not_convert, list):
self.modules_to_not_convert = [self.modules_to_not_convert]
@@ -57,7 +59,7 @@ def validate_environment(self, *args, **kwargs):
)
# Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer.adjust_max_memory
- def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
# need more space for buffers that are created during quantization
max_memory = {key: val * 0.90 for key, val in max_memory.items()}
return max_memory
@@ -90,9 +92,9 @@ def check_quantized_param_shape(self, param_name, current_param, loaded_param):
def check_if_quantized_param(
self,
model: "ModelMixin",
- param_value: Union["GGUFParameter", "torch.Tensor"],
+ param_value: "GGUFParameter" | "torch.Tensor",
param_name: str,
- state_dict: Dict[str, Any],
+ state_dict: dict[str, Any],
**kwargs,
) -> bool:
if isinstance(param_value, GGUFParameter):
@@ -103,11 +105,11 @@ def check_if_quantized_param(
def create_quantized_param(
self,
model: "ModelMixin",
- param_value: Union["GGUFParameter", "torch.Tensor"],
+ param_value: "GGUFParameter" | "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: Optional[Dict[str, Any]] = None,
- unexpected_keys: Optional[List[str]] = None,
+ state_dict: dict[str, Any] | None = None,
+ unexpected_keys: list[str] | None = None,
**kwargs,
):
module, tensor_name = get_module_from_name(model, param_name)
@@ -123,7 +125,7 @@ def _process_model_before_weight_loading(
self,
model: "ModelMixin",
device_map,
- keep_in_fp32_modules: List[str] = [],
+ keep_in_fp32_modules: list[str] = [],
**kwargs,
):
state_dict = kwargs.get("state_dict", None)
diff --git a/src/diffusers/quantizers/gguf/utils.py b/src/diffusers/quantizers/gguf/utils.py
index adb429688723..e0ad0e1cce42 100644
--- a/src/diffusers/quantizers/gguf/utils.py
+++ b/src/diffusers/quantizers/gguf/utils.py
@@ -516,6 +516,9 @@ def dequantize_gguf_tensor(tensor):
block_size, type_size = GGML_QUANT_SIZES[quant_type]
+ # Conver to plain tensor to avoid unnecessary __torch_function__ overhead.
+ tensor = tensor.as_tensor()
+
tensor = tensor.view(torch.uint8)
shape = _quant_shape_from_byte_shape(tensor.shape, type_size, block_size)
@@ -525,7 +528,7 @@ def dequantize_gguf_tensor(tensor):
dequant = dequant_fn(blocks, block_size, type_size)
dequant = dequant.reshape(shape)
- return dequant.as_tensor()
+ return dequant
class GGUFParameter(torch.nn.Parameter):
diff --git a/src/diffusers/quantizers/modelopt/modelopt_quantizer.py b/src/diffusers/quantizers/modelopt/modelopt_quantizer.py
index 7312036f52d0..5c9e198c82b2 100644
--- a/src/diffusers/quantizers/modelopt/modelopt_quantizer.py
+++ b/src/diffusers/quantizers/modelopt/modelopt_quantizer.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict, List, Union
+from typing import TYPE_CHECKING, Any
from ...utils import (
get_module_from_name,
@@ -61,7 +61,7 @@ def check_if_quantized_param(
model: "ModelMixin",
param_value: "torch.Tensor",
param_name: str,
- state_dict: Dict[str, Any],
+ state_dict: dict[str, Any],
**kwargs,
):
# ModelOpt imports diffusers internally. This is here to prevent circular imports
@@ -101,7 +101,7 @@ def create_quantized_param(
mtq.compress(module)
module.weight.requires_grad = False
- def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
max_memory = {key: val * 0.90 for key, val in max_memory.items()}
return max_memory
@@ -116,7 +116,7 @@ def update_torch_dtype(self, torch_dtype: "torch.dtype" = None) -> "torch.dtype"
torch_dtype = torch.float32
return torch_dtype
- def get_conv_param_names(self, model: "ModelMixin") -> List[str]:
+ def get_conv_param_names(self, model: "ModelMixin") -> list[str]:
"""
Get parameter names for all convolutional layers in a HuggingFace ModelMixin. Includes Conv1d/2d/3d and
ConvTranspose1d/2d/3d.
@@ -142,7 +142,7 @@ def _process_model_before_weight_loading(
self,
model: "ModelMixin",
device_map,
- keep_in_fp32_modules: List[str] = [],
+ keep_in_fp32_modules: list[str] = [],
**kwargs,
):
# ModelOpt imports diffusers internally. This is here to prevent circular imports
diff --git a/src/diffusers/quantizers/pipe_quant_config.py b/src/diffusers/quantizers/pipe_quant_config.py
index f75a337341a9..90be7d796438 100644
--- a/src/diffusers/quantizers/pipe_quant_config.py
+++ b/src/diffusers/quantizers/pipe_quant_config.py
@@ -12,8 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import inspect
-from typing import Dict, List, Optional, Union
from ..utils import is_transformers_available, logging
from .quantization_config import QuantizationConfigMixin as DiffQuantConfigMixin
@@ -47,9 +48,9 @@ class PipelineQuantizationConfig:
def __init__(
self,
quant_backend: str = None,
- quant_kwargs: Dict[str, Union[str, float, int, dict]] = None,
- components_to_quantize: Optional[Union[List[str], str]] = None,
- quant_mapping: Dict[str, Union[DiffQuantConfigMixin, "TransformersQuantConfigMixin"]] = None,
+ quant_kwargs: dict[str, str | float | int | dict] = None,
+ components_to_quantize: list[str] | str | None = None,
+ quant_mapping: dict[str, DiffQuantConfigMixin | "TransformersQuantConfigMixin"] = None,
):
self.quant_backend = quant_backend
# Initialize kwargs to be {} to set to the defaults.
diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py
index c905a928c79d..138ec7b7e989 100644
--- a/src/diffusers/quantizers/quantization_config.py
+++ b/src/diffusers/quantizers/quantization_config.py
@@ -20,6 +20,8 @@
https://github.com/huggingface/transformers/blob/52cb4034ada381fe1ffe8d428a1076e5411a8026/src/transformers/utils/quantization_config.py
"""
+from __future__ import annotations
+
import copy
import dataclasses
import importlib.metadata
@@ -30,7 +32,7 @@
from dataclasses import dataclass, is_dataclass
from enum import Enum
from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
from packaging import version
@@ -76,12 +78,12 @@ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
Instantiates a [`QuantizationConfigMixin`] from a Python dictionary of parameters.
Args:
- config_dict (`Dict[str, Any]`):
+ config_dict (`dict[str, Any]`):
Dictionary that will be used to instantiate the configuration object.
return_unused_kwargs (`bool`, *optional*, defaults to `False`):
Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in
`PreTrainedModel`.
- kwargs (`Dict[str, Any]`):
+ kwargs (`dict[str, Any]`):
Additional parameters from which to initialize the configuration object.
Returns:
@@ -103,7 +105,7 @@ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
else:
return config
- def to_json_file(self, json_file_path: Union[str, os.PathLike]):
+ def to_json_file(self, json_file_path: str | os.PathLike):
"""
Save this instance to a JSON file.
@@ -120,10 +122,10 @@ def to_json_file(self, json_file_path: Union[str, os.PathLike]):
writer.write(json_string)
- def to_dict(self) -> Dict[str, Any]:
+ def to_dict(self) -> dict[str, Any]:
"""
Serializes this instance to a Python dictionary. Returns:
- `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
"""
return copy.deepcopy(self.__dict__)
@@ -159,11 +161,11 @@ def update(self, **kwargs):
returning all the unused kwargs.
Args:
- kwargs (`Dict[str, Any]`):
+ kwargs (`dict[str, Any]`):
Dictionary of attributes to tentatively update this class.
Returns:
- `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance.
+ `dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance.
"""
to_remove = []
for key, value in kwargs.items():
@@ -202,7 +204,7 @@ class BitsAndBytesConfig(QuantizationConfigMixin):
These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of
magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6,
but a lower threshold might be needed for more unstable models (small models, fine-tuning).
- llm_int8_skip_modules (`List[str]`, *optional*):
+ llm_int8_skip_modules (`list[str]`, *optional*):
An explicit list of the modules that we do not want to convert in 8-bit. This is useful for models such as
Jukebox that has several heads in different places and not necessarily at the last position. For example
for `CausalLM` models, the last `lm_head` is typically kept in its original `dtype`.
@@ -225,7 +227,7 @@ class BitsAndBytesConfig(QuantizationConfigMixin):
quantized again.
bnb_4bit_quant_storage (`torch.dtype` or str, *optional*, defaults to `torch.uint8`):
This sets the storage type to pack the quanitzed 4-bit prarams.
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional parameters from which to initialize the configuration object.
"""
@@ -376,10 +378,10 @@ def quantization_method(self):
else:
return None
- def to_dict(self) -> Dict[str, Any]:
+ def to_dict(self) -> dict[str, Any]:
"""
Serializes this instance to a Python dictionary. Returns:
- `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
"""
output = copy.deepcopy(self.__dict__)
output["bnb_4bit_compute_dtype"] = str(output["bnb_4bit_compute_dtype"]).split(".")[1]
@@ -393,13 +395,13 @@ def __repr__(self):
config_dict = self.to_dict()
return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"
- def to_diff_dict(self) -> Dict[str, Any]:
+ def to_diff_dict(self) -> dict[str, Any]:
"""
Removes all attributes from config which correspond to the default config attributes for better readability and
serializes to a Python dictionary.
Returns:
- `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
"""
config_dict = self.to_dict()
@@ -427,7 +429,7 @@ class GGUFQuantizationConfig(QuantizationConfigMixin):
"""
- def __init__(self, compute_dtype: Optional["torch.dtype"] = None):
+ def __init__(self, compute_dtype: "torch.dtype" | None = None):
self.quant_method = QuantizationMethod.GGUF
self.compute_dtype = compute_dtype
self.pre_quantized = True
@@ -444,7 +446,7 @@ class TorchAoConfig(QuantizationConfigMixin):
"""This is a config class for torchao quantization/sparsity techniques.
Args:
- quant_type (Union[`str`, AOBaseConfig]):
+ quant_type (`str` | AOBaseConfig):
The type of quantization we want to use, currently supporting:
- **Integer quantization:**
- Full function names: `int4_weight_only`, `int8_dynamic_activation_int4_weight`,
@@ -467,10 +469,10 @@ class TorchAoConfig(QuantizationConfigMixin):
- Full function names: `uintx_weight_only`
- Shorthands: `uint1wo`, `uint2wo`, `uint3wo`, `uint4wo`, `uint5wo`, `uint6wo`, `uint7wo`
- An AOBaseConfig instance: for more advanced configuration options.
- modules_to_not_convert (`List[str]`, *optional*, default to `None`):
+ modules_to_not_convert (`list[str]`, *optional*, default to `None`):
The list of modules to not quantize, useful for quantizing models that explicitly require to have some
modules left in their original precision.
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization
supports two keyword arguments `group_size` and `inner_k_tiles` currently. More API examples and
documentation of arguments can be found in
@@ -498,8 +500,8 @@ class TorchAoConfig(QuantizationConfigMixin):
def __init__(
self,
- quant_type: Union[str, "AOBaseConfig"], # noqa: F821
- modules_to_not_convert: Optional[List[str]] = None,
+ quant_type: str | "AOBaseConfig", # noqa: F821
+ modules_to_not_convert: list[str] | None = None,
**kwargs,
) -> None:
self.quant_method = QuantizationMethod.TORCHAO
@@ -839,7 +841,7 @@ class QuantoConfig(QuantizationConfigMixin):
def __init__(
self,
weights_dtype: str = "int8",
- modules_to_not_convert: Optional[List[str]] = None,
+ modules_to_not_convert: list[str] | None = None,
**kwargs,
):
self.quant_method = QuantizationMethod.QUANTO
@@ -871,7 +873,7 @@ class NVIDIAModelOptConfig(QuantizationConfigMixin):
- INT4
- NF4
- NVFP4
- modules_to_not_convert (`List[str]`, *optional*, default to `None`):
+ modules_to_not_convert (`list[str]`, *optional*, default to `None`):
The list of modules to not quantize, useful for quantizing models that explicitly require to have some
weight_only (`bool`, *optional*, default to `False`):
If set to `True`, the quantization will be applied only to the weights of the model.
@@ -891,7 +893,7 @@ class NVIDIAModelOptConfig(QuantizationConfigMixin):
The modelopt config, useful for passing custom configs to modelopt.
disable_conv_quantization (`bool`, *optional*, default to `False`):
If set to `True`, the quantization will be disabled for convolutional layers.
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional parameters which are to be used for calibration.
"""
@@ -910,15 +912,15 @@ class NVIDIAModelOptConfig(QuantizationConfigMixin):
def __init__(
self,
quant_type: str,
- modules_to_not_convert: Optional[List[str]] = None,
+ modules_to_not_convert: list[str] | None = None,
weight_only: bool = True,
- channel_quantize: Optional[int] = None,
- block_quantize: Optional[int] = None,
- scale_channel_quantize: Optional[int] = None,
- scale_block_quantize: Optional[int] = None,
+ channel_quantize: int | None = None,
+ block_quantize: int | None = None,
+ scale_channel_quantize: int | None = None,
+ scale_block_quantize: int | None = None,
algorithm: str = "max",
- forward_loop: Optional[Callable] = None,
- modelopt_config: Optional[dict] = None,
+ forward_loop: Callable | None = None,
+ modelopt_config: dict | None = None,
disable_conv_quantization: bool = False,
**kwargs,
) -> None:
@@ -980,7 +982,7 @@ def _normalize_quant_type(self, quant_type: str) -> str:
act_type = None
self.quant_type = w_type + ("_" + act_type if act_type is not None else "")
- def get_config_from_quant_type(self) -> Dict[str, Any]:
+ def get_config_from_quant_type(self) -> dict[str, Any]:
"""
Get the config from the quantization type.
"""
diff --git a/src/diffusers/quantizers/quanto/quanto_quantizer.py b/src/diffusers/quantizers/quanto/quanto_quantizer.py
index c5f71f816fc3..a036dabfe6f4 100644
--- a/src/diffusers/quantizers/quanto/quanto_quantizer.py
+++ b/src/diffusers/quantizers/quanto/quanto_quantizer.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict, List, Union
+from typing import TYPE_CHECKING, Any
from diffusers.utils.import_utils import is_optimum_quanto_version
@@ -68,7 +68,7 @@ def check_if_quantized_param(
model: "ModelMixin",
param_value: "torch.Tensor",
param_name: str,
- state_dict: Dict[str, Any],
+ state_dict: dict[str, Any],
**kwargs,
):
# Quanto imports diffusers internally. This is here to prevent circular imports
@@ -105,7 +105,7 @@ def create_quantized_param(
module.freeze()
module.weight.requires_grad = False
- def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
max_memory = {key: val * 0.90 for key, val in max_memory.items()}
return max_memory
@@ -127,7 +127,7 @@ def update_torch_dtype(self, torch_dtype: "torch.dtype" = None) -> "torch.dtype"
torch_dtype = torch.float32
return torch_dtype
- def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]:
+ def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
# Quanto imports diffusers internally. This is here to prevent circular imports
from optimum.quanto import QModuleMixin
@@ -147,7 +147,7 @@ def _process_model_before_weight_loading(
self,
model: "ModelMixin",
device_map,
- keep_in_fp32_modules: List[str] = [],
+ keep_in_fp32_modules: list[str] = [],
**kwargs,
):
self.modules_to_not_convert = self.quantization_config.modules_to_not_convert
diff --git a/src/diffusers/quantizers/torchao/torchao_quantizer.py b/src/diffusers/quantizers/torchao/torchao_quantizer.py
index 11435b85eb4d..1679ed26a104 100644
--- a/src/diffusers/quantizers/torchao/torchao_quantizer.py
+++ b/src/diffusers/quantizers/torchao/torchao_quantizer.py
@@ -21,7 +21,7 @@
import re
import types
from fnmatch import fnmatch
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any
from packaging import version
@@ -119,7 +119,7 @@ def _update_torch_safe_globals():
_update_torch_safe_globals()
-def fuzzy_match_size(config_name: str) -> Optional[str]:
+def fuzzy_match_size(config_name: str) -> str | None:
"""
Extract the size digit from strings like "4weight", "8weight". Returns the digit as an integer if found, otherwise
None.
@@ -271,7 +271,7 @@ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
f"dtype you are using should be supported, please open an issue at https://github.com/huggingface/diffusers/issues."
)
- def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
max_memory = {key: val * 0.9 for key, val in max_memory.items()}
return max_memory
@@ -280,7 +280,7 @@ def check_if_quantized_param(
model: "ModelMixin",
param_value: "torch.Tensor",
param_name: str,
- state_dict: Dict[str, Any],
+ state_dict: dict[str, Any],
**kwargs,
) -> bool:
param_device = kwargs.pop("param_device", None)
@@ -301,8 +301,8 @@ def create_quantized_param(
param_value: "torch.Tensor",
param_name: str,
target_device: "torch.device",
- state_dict: Dict[str, Any],
- unexpected_keys: List[str],
+ state_dict: dict[str, Any],
+ unexpected_keys: list[str],
**kwargs,
):
r"""
@@ -365,7 +365,7 @@ def _process_model_before_weight_loading(
self,
model: "ModelMixin",
device_map,
- keep_in_fp32_modules: List[str] = [],
+ keep_in_fp32_modules: list[str] = [],
**kwargs,
):
self.modules_to_not_convert = self.quantization_config.modules_to_not_convert
diff --git a/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py b/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py
index 9206ee80a6b6..b6b905b0f088 100644
--- a/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py
+++ b/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py
@@ -14,7 +14,6 @@
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -43,7 +42,7 @@ class KarrasVeOutput(BaseOutput):
prev_sample: torch.Tensor
derivative: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
class KarrasVeScheduler(SchedulerMixin, ConfigMixin):
@@ -93,7 +92,7 @@ def __init__(
self.timesteps: np.IntTensor = None
self.schedule: torch.Tensor = None # sigma(t_i)
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -110,7 +109,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
"""
return sample
- def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -133,8 +132,8 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
self.schedule = torch.tensor(schedule, dtype=torch.float32, device=device)
def add_noise_to_input(
- self, sample: torch.Tensor, sigma: float, generator: Optional[torch.Generator] = None
- ) -> Tuple[torch.Tensor, float]:
+ self, sample: torch.Tensor, sigma: float, generator: torch.Generator | None = None
+ ) -> tuple[torch.Tensor, float]:
"""
Explicit Langevin-like "churn" step of adding noise to the sample according to a `gamma_i ≥ 0` to reach a
higher noise level `sigma_hat = sigma_i + gamma_i*sigma_i`.
@@ -165,7 +164,7 @@ def step(
sigma_prev: float,
sample_hat: torch.Tensor,
return_dict: bool = True,
- ) -> Union[KarrasVeOutput, Tuple]:
+ ) -> KarrasVeOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -206,7 +205,7 @@ def step_correct(
sample_prev: torch.Tensor,
derivative: torch.Tensor,
return_dict: bool = True,
- ) -> Union[KarrasVeOutput, Tuple]:
+ ) -> KarrasVeOutput | tuple:
"""
Corrects the predicted sample based on the `model_output` of the network.
diff --git a/src/diffusers/schedulers/deprecated/scheduling_sde_vp.py b/src/diffusers/schedulers/deprecated/scheduling_sde_vp.py
index 5088bdb49761..99fa408e70c0 100644
--- a/src/diffusers/schedulers/deprecated/scheduling_sde_vp.py
+++ b/src/diffusers/schedulers/deprecated/scheduling_sde_vp.py
@@ -15,7 +15,6 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch
import math
-from typing import Union
import torch
@@ -48,7 +47,7 @@ def __init__(self, num_train_timesteps=2000, beta_min=0.1, beta_max=20, sampling
self.discrete_sigmas = None
self.timesteps = None
- def set_timesteps(self, num_inference_steps, device: Union[str, torch.device] = None):
+ def set_timesteps(self, num_inference_steps, device: str | torch.device = None):
"""
Sets the continuous timesteps used for the diffusion chain (to be run before inference).
diff --git a/src/diffusers/schedulers/scheduling_amused.py b/src/diffusers/schedulers/scheduling_amused.py
index a0b8fbc862b0..ef49d8db91b1 100644
--- a/src/diffusers/schedulers/scheduling_amused.py
+++ b/src/diffusers/schedulers/scheduling_amused.py
@@ -1,6 +1,6 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import torch
@@ -9,7 +9,7 @@
from .scheduling_utils import SchedulerMixin
-def gumbel_noise(t: torch.Tensor, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+def gumbel_noise(t: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
"""
Generate Gumbel noise for sampling.
@@ -32,7 +32,7 @@ def mask_by_random_topk(
mask_len: torch.Tensor,
probs: torch.Tensor,
temperature: float = 1.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
) -> torch.Tensor:
"""
Mask tokens by selecting the top-k lowest confidence scores with temperature-based randomness.
@@ -73,7 +73,7 @@ class AmusedSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Generator | None = None
class AmusedScheduler(SchedulerMixin, ConfigMixin):
@@ -96,8 +96,8 @@ class AmusedScheduler(SchedulerMixin, ConfigMixin):
order = 1
- temperatures: Optional[torch.Tensor]
- timesteps: Optional[torch.Tensor]
+ temperatures: torch.Generator | None
+ timesteps: torch.Generator | None
@register_to_config
def __init__(
@@ -111,23 +111,9 @@ def __init__(
def set_timesteps(
self,
num_inference_steps: int,
- temperature: Union[float, Tuple[float, float], List[float]] = (2, 0),
- device: Optional[Union[str, torch.device]] = None,
- ) -> None:
- """
- Set the discrete timesteps used for the diffusion chain (to be run before inference).
-
- Args:
- num_inference_steps (`int`):
- The number of diffusion steps used when generating samples with a pre-trained model.
- temperature (`Union[float, Tuple[float, float], List[float]]`, *optional*, defaults to `(2, 0)`):
- Temperature parameter(s) for controlling the randomness of sampling. If a tuple or list is provided,
- temperatures will be linearly interpolated between the first and second values across all timesteps. If
- a single value is provided, temperatures will be linearly interpolated from that value to 0.01.
- device (`str` or `torch.device`, *optional*):
- The device to which the timesteps and temperatures should be moved to. If `None`, the timesteps are not
- moved.
- """
+ temperature: int | tuple[int, int] | list[int] = (2, 0),
+ device: str | torch.device = None,
+ ):
self.timesteps = torch.arange(num_inference_steps, device=device).flip(0)
if isinstance(temperature, (tuple, list)):
@@ -141,35 +127,9 @@ def step(
timestep: int,
sample: torch.LongTensor,
starting_mask_ratio: float = 1.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[AmusedSchedulerOutput, Tuple[torch.Tensor, torch.Tensor]]:
- """
- Predict the sample at the previous timestep by masking tokens based on confidence scores.
-
- Args:
- model_output (`torch.Tensor`):
- The direct output from the learned diffusion model. Typically of shape `(batch_size, num_tokens,
- codebook_size)` or `(batch_size, codebook_size, height, width)` for 2D inputs.
- timestep (`int`):
- The current discrete timestep in the diffusion chain.
- sample (`torch.LongTensor`):
- A current instance of a sample created by the diffusion process. Contains token IDs, with masked
- positions indicated by `mask_token_id`.
- starting_mask_ratio (`float`, *optional*, defaults to 1.0):
- A multiplier applied to the mask ratio schedule. Values less than 1.0 will result in fewer tokens being
- masked at each step.
- generator (`torch.Generator`, *optional*):
- A random number generator for reproducible sampling.
- return_dict (`bool`, *optional*, defaults to `True`):
- Whether to return an [`~schedulers.scheduling_amused.AmusedSchedulerOutput`] or a plain tuple.
-
- Returns:
- [`~schedulers.scheduling_amused.AmusedSchedulerOutput`] or `tuple`:
- If `return_dict` is `True`, [`~schedulers.scheduling_amused.AmusedSchedulerOutput`] is returned,
- otherwise a tuple is returned where the first element is the sample tensor (`prev_sample`) and the
- second element is the predicted original sample tensor (`pred_original_sample`).
- """
+ ) -> AmusedSchedulerOutput | tuple:
two_dim_input = sample.ndim == 3 and model_output.ndim == 4
if two_dim_input:
@@ -234,7 +194,7 @@ def add_noise(
self,
sample: torch.LongTensor,
timesteps: int,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
) -> torch.LongTensor:
"""
Add noise to a sample by randomly masking tokens according to the masking schedule.
diff --git a/src/diffusers/schedulers/scheduling_consistency_decoder.py b/src/diffusers/schedulers/scheduling_consistency_decoder.py
index 23c0e138c4ce..7f5e0cad22b9 100644
--- a/src/diffusers/schedulers/scheduling_consistency_decoder.py
+++ b/src/diffusers/schedulers/scheduling_consistency_decoder.py
@@ -1,6 +1,6 @@
import math
from dataclasses import dataclass
-from typing import Literal, Optional, Tuple, Union
+from typing import Literal
import torch
@@ -120,19 +120,9 @@ def __init__(
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- ) -> None:
- """
- Sets the discrete timesteps used for the diffusion chain (to be run before inference).
-
- Args:
- num_inference_steps (`int`, *optional*):
- The number of diffusion steps used when generating samples with a pre-trained model. Currently, only
- `2` inference steps are supported.
- device (`str` or `torch.device`, *optional*):
- The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- """
+ num_inference_steps: int | None = None,
+ device: str | torch.device = None,
+ ):
if num_inference_steps != 2:
raise ValueError("Currently more than 2 inference steps are not supported.")
@@ -155,7 +145,7 @@ def init_noise_sigma(self) -> torch.Tensor:
"""
return self.sqrt_one_minus_alphas_cumprod[self.timesteps[0]]
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -175,11 +165,11 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
def step(
self,
model_output: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[ConsistencyDecoderSchedulerOutput, Tuple]:
+ ) -> ConsistencyDecoderSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_consistency_models.py b/src/diffusers/schedulers/scheduling_consistency_models.py
index 195ff81b4c91..d6bdf73ff9ce 100644
--- a/src/diffusers/schedulers/scheduling_consistency_models.py
+++ b/src/diffusers/schedulers/scheduling_consistency_models.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
import numpy as np
import torch
@@ -66,7 +65,7 @@ class CMStochasticIterativeScheduler(SchedulerMixin, ConfigMixin):
[paper](https://huggingface.co/papers/2206.00364). Defaults to 7.0 from the original implementation.
clip_denoised (`bool`, defaults to `True`):
Whether to clip the denoised outputs to `(-1, 1)`.
- timesteps (`List` or `np.ndarray` or `torch.Tensor`, *optional*):
+ timesteps (`list` or `np.ndarray` or `torch.Tensor`, *optional*):
An explicit timestep schedule that can be optionally specified. The timesteps are expected to be in
increasing order.
"""
@@ -102,7 +101,7 @@ def __init__(
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
@@ -113,7 +112,7 @@ def step_index(self) -> Optional[int]:
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
@@ -134,7 +133,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
"""
self._begin_index = begin_index
- def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
"""
Scales the consistency model input by `(sigma**2 + sigma_data**2) ** 0.5`.
@@ -159,7 +158,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T
self.is_scale_input_called = True
return sample
- def sigma_to_t(self, sigmas: Union[float, np.ndarray]) -> np.ndarray:
+ def sigma_to_t(self, sigmas: float | np.ndarray):
"""
Gets scaled timesteps from the Karras sigmas for input to the consistency model.
@@ -180,10 +179,10 @@ def sigma_to_t(self, sigmas: Union[float, np.ndarray]) -> np.ndarray:
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- ) -> None:
+ num_inference_steps: int | None = None,
+ device: str | torch.device = None,
+ timesteps: list[int] | None = None,
+ ):
"""
Sets the timesteps used for the diffusion chain (to be run before inference).
@@ -192,7 +191,7 @@ def set_timesteps(
The number of diffusion steps used when generating samples with a pre-trained model.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed,
`num_inference_steps` must be `None`.
@@ -274,7 +273,7 @@ def _convert_to_karras(self, ramp: np.ndarray) -> np.ndarray:
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
return sigmas
- def get_scalings(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def get_scalings(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Computes the scaling factors for the consistency model output.
@@ -283,7 +282,7 @@ def get_scalings(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]
The current sigma value in the noise schedule.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing `c_skip` (scaling for the input sample) and `c_out` (scaling for the model output).
"""
sigma_data = self.config.sigma_data
@@ -292,7 +291,7 @@ def get_scalings(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]
c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
return c_skip, c_out
- def get_scalings_for_boundary_condition(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def get_scalings_for_boundary_condition(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Gets the scalings used in the consistency model parameterization (from Appendix C of the
[paper](https://huggingface.co/papers/2303.01469)) to enforce boundary condition.
@@ -304,7 +303,7 @@ def get_scalings_for_boundary_condition(self, sigma: torch.Tensor) -> Tuple[torc
The current sigma in the Karras sigma schedule.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A two-element tuple where `c_skip` (which weights the current sample) is the first element and `c_out`
(which weights the consistency model output) is the second element.
"""
@@ -317,7 +316,7 @@ def get_scalings_for_boundary_condition(self, sigma: torch.Tensor) -> Tuple[torc
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -347,7 +346,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -365,11 +364,11 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[CMStochasticIterativeSchedulerOutput, Tuple]:
+ ) -> CMStochasticIterativeSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py
index 796d62e5b6a0..18ee272ef619 100644
--- a/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py
+++ b/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py
@@ -15,7 +15,7 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver and https://github.com/NVlabs/edm
import math
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -128,7 +128,7 @@ def init_noise_sigma(self) -> float:
return (self.config.sigma_max**2 + 1) ** 0.5
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
@@ -139,7 +139,7 @@ def step_index(self) -> Optional[int]:
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
@@ -161,7 +161,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
self._begin_index = begin_index
# Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_inputs
- def precondition_inputs(self, sample: torch.Tensor, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
+ def precondition_inputs(self, sample: torch.Tensor, sigma: float | torch.Tensor) -> torch.Tensor:
"""
Precondition the input sample by scaling it according to the EDM formulation.
@@ -179,7 +179,7 @@ def precondition_inputs(self, sample: torch.Tensor, sigma: Union[float, torch.Te
scaled_sample = sample * c_in
return scaled_sample
- def precondition_noise(self, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
+ def precondition_noise(self, sigma: float | torch.Tensor) -> torch.Tensor:
"""
Precondition the noise level by computing a normalized timestep representation.
@@ -201,7 +201,7 @@ def precondition_outputs(
self,
sample: torch.Tensor,
model_output: torch.Tensor,
- sigma: Union[float, torch.Tensor],
+ sigma: float | torch.Tensor,
) -> torch.Tensor:
"""
Precondition the model outputs according to the EDM formulation.
@@ -233,7 +233,7 @@ def precondition_outputs(
return denoised
# Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.scale_model_input
- def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
"""
Scale the denoising model input to match the Euler algorithm. Ensures interchangeability with schedulers that
need to scale the denoising model input depending on the current timestep.
@@ -257,9 +257,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T
self.is_scale_input_called = True
return sample
- def set_timesteps(
- self, num_inference_steps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None
- ) -> None:
+ def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device | None = None) -> None:
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -309,8 +307,8 @@ def set_timesteps(
def _compute_karras_sigmas(
self,
ramp: torch.Tensor,
- sigma_min: Optional[float] = None,
- sigma_max: Optional[float] = None,
+ sigma_min: float | None = None,
+ sigma_max: float | None = None,
) -> torch.Tensor:
"""
Construct the noise schedule of [Karras et al. (2022)](https://huggingface.co/papers/2206.00364).
@@ -340,8 +338,8 @@ def _compute_karras_sigmas(
def _compute_exponential_sigmas(
self,
ramp: torch.Tensor,
- sigma_min: Optional[float] = None,
- sigma_max: Optional[float] = None,
+ sigma_min: float | None = None,
+ sigma_max: float | None = None,
) -> torch.Tensor:
"""
Compute the exponential sigma schedule. Implementation closely follows k-diffusion:
@@ -401,7 +399,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:
t = t.reshape(sigma.shape)
return t
- def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Convert sigma to alpha and sigma_t values for the diffusion process.
@@ -410,7 +408,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to
The sigma (noise level) value.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing `alpha_t` (always 1 since inputs are pre-scaled) and `sigma_t` (same as input
sigma).
"""
@@ -451,7 +449,7 @@ def dpm_solver_first_order_update(
self,
model_output: torch.Tensor,
sample: torch.Tensor = None,
- noise: Optional[torch.Tensor] = None,
+ noise: torch.Tensor | None = None,
) -> torch.Tensor:
"""
One step for the first-order DPMSolver (equivalent to DDIM).
@@ -487,15 +485,15 @@ def dpm_solver_first_order_update(
def multistep_dpm_solver_second_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
sample: torch.Tensor = None,
- noise: Optional[torch.Tensor] = None,
+ noise: torch.Tensor | None = None,
) -> torch.Tensor:
"""
One step for the second-order multistep DPMSolver.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`):
A current instance of a sample created by the diffusion process.
@@ -546,8 +544,8 @@ def multistep_dpm_solver_second_order_update(
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
def index_for_timestep(
self,
- timestep: Union[int, torch.Tensor],
- schedule_timesteps: Optional[torch.Tensor] = None,
+ timestep: int | torch.Tensor,
+ schedule_timesteps: torch.Tensor | None = None,
) -> int:
"""
Find the index for a given timestep in the schedule.
@@ -581,7 +579,7 @@ def index_for_timestep(
return step_index
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
- def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: int | torch.Tensor) -> None:
"""
Initialize the step_index counter for the scheduler.
@@ -600,11 +598,11 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[int, torch.Tensor],
+ timestep: int | torch.Tensor,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the multistep DPMSolver.
@@ -731,7 +729,7 @@ def add_noise(
return noisy_samples
# Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._get_conditioning_c_in
- def _get_conditioning_c_in(self, sigma: Union[float, torch.Tensor]) -> Union[float, torch.Tensor]:
+ def _get_conditioning_c_in(self, sigma: float | torch.Tensor) -> float | torch.Tensor:
"""
Compute the input conditioning factor for the EDM formulation.
diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py
index 92c3e20013dd..d29bb5b2593c 100644
--- a/src/diffusers/schedulers/scheduling_ddim.py
+++ b/src/diffusers/schedulers/scheduling_ddim.py
@@ -17,7 +17,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -44,7 +44,7 @@ class DDIMSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -196,8 +196,8 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.0001,
beta_end: float = 0.02,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
clip_sample: bool = True,
set_alpha_to_one: bool = True,
steps_offset: int = 0,
@@ -242,7 +242,7 @@ def __init__(
self.num_inference_steps = None
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -331,14 +331,14 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
return sample
- def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None) -> None:
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
Args:
num_inference_steps (`int`):
The number of diffusion steps used when generating samples with a pre-trained model.
- device (`Union[str, torch.device]`, *optional*):
+ device (`str | torch.device`, *optional*):
The device to use for the timesteps.
Raises:
@@ -388,10 +388,10 @@ def step(
sample: torch.Tensor,
eta: float = 0.0,
use_clipped_model_output: bool = False,
- generator: Optional[torch.Generator] = None,
- variance_noise: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ variance_noise: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[DDIMSchedulerOutput, Tuple]:
+ ) -> DDIMSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py
index 1a77a652786d..11cbb8a6b9f7 100644
--- a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py
+++ b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py
@@ -18,7 +18,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -44,7 +44,7 @@ class DDIMSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -192,7 +192,7 @@ def __init__(
beta_start: float = 0.00085,
beta_end: float = 0.0120,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "scaled_linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
clip_sample: bool = True,
set_alpha_to_one: bool = True,
steps_offset: int = 0,
@@ -257,7 +257,7 @@ def _get_variance(self, timestep: int, prev_timestep: int) -> torch.Tensor:
return variance
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -277,7 +277,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
def set_timesteps(
self,
num_inference_steps: int,
- device: Optional[Union[str, torch.device]] = None,
+ device: str | torch.device | None = None,
) -> None:
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -330,10 +330,10 @@ def step(
sample: torch.Tensor,
eta: float = 0.0,
use_clipped_model_output: bool = False,
- generator: Optional[torch.Generator] = None,
- variance_noise: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ variance_noise: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[DDIMSchedulerOutput, Tuple]:
+ ) -> DDIMSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_ddim_flax.py b/src/diffusers/schedulers/scheduling_ddim_flax.py
index 476f741bcdde..3bb44caef2bf 100644
--- a/src/diffusers/schedulers/scheduling_ddim_flax.py
+++ b/src/diffusers/schedulers/scheduling_ddim_flax.py
@@ -16,7 +16,6 @@
# and https://github.com/hojonathanho/diffusion
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import flax
import jax.numpy as jnp
@@ -44,7 +43,7 @@ class DDIMSchedulerState:
# setable values
init_noise_sigma: jnp.ndarray
timesteps: jnp.ndarray
- num_inference_steps: Optional[int] = None
+ num_inference_steps: int = None
@classmethod
def create(
@@ -121,7 +120,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: str = "linear",
- trained_betas: Optional[jnp.ndarray] = None,
+ trained_betas: jnp.ndarray | None = None,
clip_sample: bool = True,
clip_sample_range: float = 1.0,
set_alpha_to_one: bool = True,
@@ -135,7 +134,7 @@ def __init__(
)
self.dtype = dtype
- def create_state(self, common: Optional[CommonSchedulerState] = None) -> DDIMSchedulerState:
+ def create_state(self, common: CommonSchedulerState | None = None) -> DDIMSchedulerState:
if common is None:
common = CommonSchedulerState.create(self)
@@ -163,7 +162,7 @@ def scale_model_input(
self,
state: DDIMSchedulerState,
sample: jnp.ndarray,
- timestep: Optional[int] = None,
+ timestep: int | None = None,
) -> jnp.ndarray:
"""
Args:
@@ -177,7 +176,7 @@ def scale_model_input(
return sample
def set_timesteps(
- self, state: DDIMSchedulerState, num_inference_steps: int, shape: Tuple = ()
+ self, state: DDIMSchedulerState, num_inference_steps: int, shape: tuple = ()
) -> DDIMSchedulerState:
"""
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -220,7 +219,7 @@ def step(
sample: jnp.ndarray,
eta: float = 0.0,
return_dict: bool = True,
- ) -> Union[FlaxDDIMSchedulerOutput, Tuple]:
+ ) -> FlaxDDIMSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_ddim_inverse.py b/src/diffusers/schedulers/scheduling_ddim_inverse.py
index a3c9ed1f6258..4d130266841b 100644
--- a/src/diffusers/schedulers/scheduling_ddim_inverse.py
+++ b/src/diffusers/schedulers/scheduling_ddim_inverse.py
@@ -16,7 +16,7 @@
# and https://github.com/hojonathanho/diffusion
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -42,7 +42,7 @@ class DDIMSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -188,7 +188,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
clip_sample: bool = True,
set_alpha_to_one: bool = True,
steps_offset: int = 0,
@@ -247,7 +247,7 @@ def __init__(
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps).copy().astype(np.int64))
# Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.scale_model_input
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -267,7 +267,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
def set_timesteps(
self,
num_inference_steps: int,
- device: Optional[Union[str, torch.device]] = None,
+ device: str | torch.device | None = None,
) -> None:
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -312,7 +312,7 @@ def step(
timestep: int,
sample: torch.Tensor,
return_dict: bool = True,
- ) -> Union[DDIMSchedulerOutput, Tuple]:
+ ) -> DDIMSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_ddim_parallel.py b/src/diffusers/schedulers/scheduling_ddim_parallel.py
index 76f0636fbf6c..a5bbac60275c 100644
--- a/src/diffusers/schedulers/scheduling_ddim_parallel.py
+++ b/src/diffusers/schedulers/scheduling_ddim_parallel.py
@@ -17,7 +17,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -44,7 +44,7 @@ class DDIMParallelSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -202,8 +202,8 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.0001,
beta_end: float = 0.02,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
clip_sample: bool = True,
set_alpha_to_one: bool = True,
steps_offset: int = 0,
@@ -249,7 +249,7 @@ def __init__(
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
# Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.scale_model_input
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -266,7 +266,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
"""
return sample
- def _get_variance(self, timestep: int, prev_timestep: Optional[int] = None) -> torch.Tensor:
+ def _get_variance(self, timestep: int, prev_timestep: int | None = None) -> torch.Tensor:
if prev_timestep is None:
prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps
@@ -335,14 +335,14 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
return sample
# Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.set_timesteps
- def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None) -> None:
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None) -> None:
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
Args:
num_inference_steps (`int`):
The number of diffusion steps used when generating samples with a pre-trained model.
- device (`Union[str, torch.device]`, *optional*):
+ device (`str | torch.device`, *optional*):
The device to use for the timesteps.
Raises:
@@ -392,10 +392,10 @@ def step(
sample: torch.Tensor,
eta: float = 0.0,
use_clipped_model_output: bool = False,
- generator: Optional[torch.Generator] = None,
- variance_noise: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ variance_noise: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[DDIMParallelSchedulerOutput, Tuple]:
+ ) -> DDIMParallelSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_ddpm.py b/src/diffusers/schedulers/scheduling_ddpm.py
index 2e2816bbf3ba..972c46c6e930 100644
--- a/src/diffusers/schedulers/scheduling_ddpm.py
+++ b/src/diffusers/schedulers/scheduling_ddpm.py
@@ -16,7 +16,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -42,7 +42,7 @@ class DDPMSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
def betas_for_alpha_bar(
@@ -190,7 +190,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2", "sigmoid"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
variance_type: Literal[
"fixed_small",
"fixed_small_log",
@@ -254,7 +254,7 @@ def __init__(
self.variance_type = variance_type
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -273,9 +273,9 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device] = None,
- timesteps: Optional[List[int]] = None,
+ num_inference_steps: int = None,
+ device: str | torch.device = None,
+ timesteps: list[int] | None = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -286,7 +286,7 @@ def set_timesteps(
`timesteps` must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed,
`num_inference_steps` must be `None`.
@@ -348,17 +348,16 @@ def set_timesteps(
def _get_variance(
self,
t: int,
- predicted_variance: Optional[torch.Tensor] = None,
- variance_type: Optional[
- Literal[
- "fixed_small",
- "fixed_small_log",
- "fixed_large",
- "fixed_large_log",
- "learned",
- "learned_range",
- ]
- ] = None,
+ predicted_variance: torch.Tensor | None = None,
+ variance_type: Literal[
+ "fixed_small",
+ "fixed_small_log",
+ "fixed_large",
+ "fixed_large_log",
+ "learned",
+ "learned_range",
+ ]
+ | None = None,
) -> torch.Tensor:
"""
Compute the variance for a given timestep according to the specified variance type.
@@ -464,9 +463,9 @@ def step(
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[DDPMSchedulerOutput, Tuple]:
+ ) -> DDPMSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -646,7 +645,7 @@ def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: tor
def __len__(self) -> int:
return self.config.num_train_timesteps
- def previous_timestep(self, timestep: int) -> Union[int, torch.Tensor]:
+ def previous_timestep(self, timestep: int) -> int | torch.Tensor:
"""
Compute the previous timestep in the diffusion chain.
diff --git a/src/diffusers/schedulers/scheduling_ddpm_flax.py b/src/diffusers/schedulers/scheduling_ddpm_flax.py
index e02b7ea0c0f3..7840ff729488 100644
--- a/src/diffusers/schedulers/scheduling_ddpm_flax.py
+++ b/src/diffusers/schedulers/scheduling_ddpm_flax.py
@@ -15,7 +15,6 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import flax
import jax
@@ -43,7 +42,7 @@ class DDPMSchedulerState:
# setable values
init_noise_sigma: jnp.ndarray
timesteps: jnp.ndarray
- num_inference_steps: Optional[int] = None
+ num_inference_steps: int = None
@classmethod
def create(
@@ -108,7 +107,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: str = "linear",
- trained_betas: Optional[jnp.ndarray] = None,
+ trained_betas: jnp.ndarray | None = None,
variance_type: str = "fixed_small",
clip_sample: bool = True,
prediction_type: str = "epsilon",
@@ -120,7 +119,7 @@ def __init__(
)
self.dtype = dtype
- def create_state(self, common: Optional[CommonSchedulerState] = None) -> DDPMSchedulerState:
+ def create_state(self, common: CommonSchedulerState | None = None) -> DDPMSchedulerState:
if common is None:
common = CommonSchedulerState.create(self)
@@ -139,7 +138,7 @@ def scale_model_input(
self,
state: DDPMSchedulerState,
sample: jnp.ndarray,
- timestep: Optional[int] = None,
+ timestep: int | None = None,
) -> jnp.ndarray:
"""
Args:
@@ -153,7 +152,7 @@ def scale_model_input(
return sample
def set_timesteps(
- self, state: DDPMSchedulerState, num_inference_steps: int, shape: Tuple = ()
+ self, state: DDPMSchedulerState, num_inference_steps: int, shape: tuple = ()
) -> DDPMSchedulerState:
"""
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -214,9 +213,9 @@ def step(
model_output: jnp.ndarray,
timestep: int,
sample: jnp.ndarray,
- key: Optional[jax.Array] = None,
+ key: jax.Array | None = None,
return_dict: bool = True,
- ) -> Union[FlaxDDPMSchedulerOutput, Tuple]:
+ ) -> FlaxDDPMSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_ddpm_parallel.py b/src/diffusers/schedulers/scheduling_ddpm_parallel.py
index b02c5376f2c6..725958b01a44 100644
--- a/src/diffusers/schedulers/scheduling_ddpm_parallel.py
+++ b/src/diffusers/schedulers/scheduling_ddpm_parallel.py
@@ -16,7 +16,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -43,7 +43,7 @@ class DDPMParallelSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -203,7 +203,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2", "sigmoid"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
variance_type: Literal[
"fixed_small",
"fixed_small_log",
@@ -268,7 +268,7 @@ def __init__(
self.variance_type = variance_type
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.scale_model_input
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -288,9 +288,9 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.set_timesteps
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device] = None,
- timesteps: Optional[List[int]] = None,
+ num_inference_steps: int = None,
+ device: str | torch.device = None,
+ timesteps: list[int] | None = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -301,7 +301,7 @@ def set_timesteps(
`timesteps` must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed,
`num_inference_steps` must be `None`.
@@ -364,17 +364,16 @@ def set_timesteps(
def _get_variance(
self,
t: int,
- predicted_variance: Optional[torch.Tensor] = None,
- variance_type: Optional[
- Literal[
- "fixed_small",
- "fixed_small_log",
- "fixed_large",
- "fixed_large_log",
- "learned",
- "learned_range",
- ]
- ] = None,
+ predicted_variance: torch.Tensor | None = None,
+ variance_type: Literal[
+ "fixed_small",
+ "fixed_small_log",
+ "fixed_large",
+ "fixed_large_log",
+ "learned",
+ "learned_range",
+ ]
+ | None = None,
) -> torch.Tensor:
"""
Compute the variance for a given timestep according to the specified variance type.
@@ -481,9 +480,9 @@ def step(
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[DDPMParallelSchedulerOutput, Tuple]:
+ ) -> DDPMParallelSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -747,7 +746,7 @@ def __len__(self):
return self.config.num_train_timesteps
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep
- def previous_timestep(self, timestep: int) -> Union[int, torch.Tensor]:
+ def previous_timestep(self, timestep: int) -> int | torch.Tensor:
"""
Compute the previous timestep in the diffusion chain.
diff --git a/src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py b/src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py
index 71f08277ebd7..e3b87b789a03 100644
--- a/src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py
+++ b/src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py
@@ -17,7 +17,6 @@
import math
from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
import torch
@@ -125,7 +124,7 @@ def _alpha_cumprod(self, t, device):
) ** 2 / self._init_alpha_cumprod.to(device)
return alpha_cumprod.clamp(0.0001, 0.9999)
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -142,14 +141,14 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
def set_timesteps(
self,
num_inference_steps: int = None,
- timesteps: Optional[List[int]] = None,
- device: Union[str, torch.device] = None,
+ timesteps: list[int] | None = None,
+ device: str | torch.device = None,
):
"""
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
Args:
- num_inference_steps (`Dict[float, int]`):
+ num_inference_steps (`dict[float, int]`):
the number of diffusion steps used when generating samples with a pre-trained model. If passed, then
`timesteps` must be `None`.
device (`str` or `torch.device`, optional):
@@ -168,7 +167,7 @@ def step(
sample: torch.Tensor,
generator=None,
return_dict: bool = True,
- ) -> Union[DDPMWuerstchenSchedulerOutput, Tuple]:
+ ) -> DDPMWuerstchenSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_deis_multistep.py b/src/diffusers/schedulers/scheduling_deis_multistep.py
index 58086c13b5dc..c22d2b5b439d 100644
--- a/src/diffusers/schedulers/scheduling_deis_multistep.py
+++ b/src/diffusers/schedulers/scheduling_deis_multistep.py
@@ -16,7 +16,7 @@
# The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
import math
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -100,7 +100,7 @@ class DEISMultistepScheduler(SchedulerMixin, ConfigMixin):
beta_schedule (`"linear"`, `"scaled_linear"`, or `"squaredcos_cap_v2"`, defaults to `"linear"`):
The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
`linear`, `scaled_linear`, or `squaredcos_cap_v2`.
- trained_betas (`np.ndarray` or `List[float]`, *optional*):
+ trained_betas (`np.ndarray` or `list[float]`, *optional*):
Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
solver_order (`int`, defaults to `2`):
The DEIS order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided
@@ -155,7 +155,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
solver_order: int = 2,
prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon",
thresholding: bool = False,
@@ -164,11 +164,11 @@ def __init__(
algorithm_type: Literal["deis"] = "deis",
solver_type: Literal["logrho"] = "logrho",
lower_order_final: bool = True,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- use_flow_sigmas: Optional[bool] = False,
- flow_shift: Optional[float] = 1.0,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ use_flow_sigmas: bool = False,
+ flow_shift: float = 1.0,
timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
steps_offset: int = 0,
use_dynamic_shifting: bool = False,
@@ -245,14 +245,14 @@ def __init__(
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
"""
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
"""
@@ -269,12 +269,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
"""
self._begin_index = begin_index
- def set_timesteps(
- self,
- num_inference_steps: int,
- device: Union[str, torch.device] = None,
- mu: Optional[float] = None,
- ) -> None:
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -440,7 +435,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:
return t
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t
- def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Convert sigma values to alpha_t and sigma_t values.
@@ -449,7 +444,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to
The sigma value(s) to convert.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing (alpha_t, sigma_t) values.
"""
if self.config.use_flow_sigmas:
@@ -706,7 +701,7 @@ def deis_first_order_update(
def multistep_deis_second_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
sample: torch.Tensor = None,
**kwargs,
@@ -715,7 +710,7 @@ def multistep_deis_second_order_update(
One step for the second-order multistep DEIS.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`):
A current instance of a sample created by the diffusion process.
@@ -779,7 +774,7 @@ def ind_fn(t, b, c):
def multistep_deis_third_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
sample: torch.Tensor = None,
**kwargs,
@@ -788,7 +783,7 @@ def multistep_deis_third_order_update(
One step for the third-order multistep DEIS.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`):
A current instance of a sample created by diffusion process.
@@ -868,8 +863,8 @@ def ind_fn(t, b, c, d):
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
def index_for_timestep(
self,
- timestep: Union[int, torch.Tensor],
- schedule_timesteps: Optional[torch.Tensor] = None,
+ timestep: int | torch.Tensor,
+ schedule_timesteps: torch.Tensor | None = None,
) -> int:
"""
Find the index for a given timestep in the schedule.
@@ -903,7 +898,7 @@ def index_for_timestep(
return step_index
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
- def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: int | torch.Tensor) -> None:
"""
Initialize the step_index counter for the scheduler.
@@ -922,10 +917,10 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[int, torch.Tensor],
+ timestep: int | torch.Tensor,
sample: torch.Tensor,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the multistep DEIS.
diff --git a/src/diffusers/schedulers/scheduling_dpm_cogvideox.py b/src/diffusers/schedulers/scheduling_dpm_cogvideox.py
index 0c576467a19a..4a238a5408ae 100644
--- a/src/diffusers/schedulers/scheduling_dpm_cogvideox.py
+++ b/src/diffusers/schedulers/scheduling_dpm_cogvideox.py
@@ -18,7 +18,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -45,7 +45,7 @@ class DDIMSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -194,7 +194,7 @@ def __init__(
beta_start: float = 0.00085,
beta_end: float = 0.0120,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "scaled_linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
clip_sample: bool = True,
set_alpha_to_one: bool = True,
steps_offset: int = 0,
@@ -259,7 +259,7 @@ def _get_variance(self, timestep, prev_timestep):
return variance
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -279,7 +279,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
def set_timesteps(
self,
num_inference_steps: int,
- device: Optional[Union[str, torch.device]] = None,
+ device: str | torch.device | None = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -332,8 +332,8 @@ def get_variables(
self,
alpha_prod_t: torch.Tensor,
alpha_prod_t_prev: torch.Tensor,
- alpha_prod_t_back: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor]:
+ alpha_prod_t_back: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Compute the variables used for DPM-Solver++ (2M) referencing the original implementation.
@@ -364,14 +364,11 @@ def get_variables(
def get_mult(
self,
h: torch.Tensor,
- r: Optional[torch.Tensor],
+ r: torch.Tensor,
alpha_prod_t: torch.Tensor,
alpha_prod_t_prev: torch.Tensor,
- alpha_prod_t_back: Optional[torch.Tensor] = None,
- ) -> Union[
- Tuple[torch.Tensor, torch.Tensor],
- Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
- ]:
+ alpha_prod_t_back: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Compute the multipliers for the previous sample and the predicted original sample.
@@ -404,16 +401,16 @@ def get_mult(
def step(
self,
model_output: torch.Tensor,
- old_pred_original_sample: Optional[torch.Tensor],
+ old_pred_original_sample: torch.Tensor,
timestep: int,
timestep_back: int,
sample: torch.Tensor,
eta: float = 0.0,
use_clipped_model_output: bool = False,
- generator: Optional[torch.Generator] = None,
- variance_noise: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ variance_noise: torch.Tensor | None = None,
return_dict: bool = False,
- ) -> Union[DDIMSchedulerOutput, Tuple]:
+ ) -> DDIMSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
index a1b118b2adcc..9c15df4569ca 100644
--- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
+++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
@@ -15,7 +15,7 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver
import math
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -217,8 +217,8 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.0001,
beta_end: float = 0.02,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
solver_order: int = 2,
prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon",
thresholding: bool = False,
@@ -228,15 +228,15 @@ def __init__(
solver_type: Literal["midpoint", "heun"] = "midpoint",
lower_order_final: bool = True,
euler_at_final: bool = False,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- use_lu_lambdas: Optional[bool] = False,
- use_flow_sigmas: Optional[bool] = False,
- flow_shift: Optional[float] = 1.0,
- final_sigmas_type: Optional[Literal["zero", "sigma_min"]] = "zero",
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ use_lu_lambdas: bool = False,
+ use_flow_sigmas: bool = False,
+ flow_shift: float = 1.0,
+ final_sigmas_type: Literal["zero", "sigma_min"] = "zero",
lambda_min_clipped: float = -float("inf"),
- variance_type: Optional[Literal["learned", "learned_range"]] = None,
+ variance_type: Literal["learned", "learned_range"] | None = None,
timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
steps_offset: int = 0,
rescale_betas_zero_snr: bool = False,
@@ -366,11 +366,11 @@ def set_begin_index(self, begin_index: int = 0):
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- mu: Optional[float] = None,
- timesteps: Optional[List[int]] = None,
- ) -> None:
+ num_inference_steps: int = None,
+ device: str | torch.device = None,
+ mu: float | None = None,
+ timesteps: list[int] | None = None,
+ ):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -379,10 +379,7 @@ def set_timesteps(
The number of diffusion steps used when generating samples with a pre-trained model.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- mu (`float`, *optional*):
- The mu parameter for dynamic shifting. If provided, requires `use_dynamic_shifting=True` and
- `time_shift_type="exponential"`.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated
based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas`
must be `None`, and `timestep_spacing` attribute will be ignored.
@@ -577,7 +574,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:
t = t.reshape(sigma.shape)
return t
- def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Convert sigma values to alpha_t and sigma_t values.
@@ -586,7 +583,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to
The sigma value(s) to convert.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing (alpha_t, sigma_t) values.
"""
if self.config.use_flow_sigmas:
@@ -750,7 +747,7 @@ def convert_model_output(
self,
model_output: torch.Tensor,
*args,
- sample: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -848,8 +845,8 @@ def dpm_solver_first_order_update(
self,
model_output: torch.Tensor,
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -920,17 +917,17 @@ def dpm_solver_first_order_update(
def multistep_dpm_solver_second_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
One step for the second-order multistep DPMSolver.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`, *optional*):
A current instance of a sample created by the diffusion process.
@@ -1043,17 +1040,17 @@ def multistep_dpm_solver_second_order_update(
def multistep_dpm_solver_third_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
One step for the third-order multistep DPMSolver.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`, *optional*):
A current instance of a sample created by diffusion process.
@@ -1140,8 +1137,8 @@ def multistep_dpm_solver_third_order_update(
def index_for_timestep(
self,
- timestep: Union[int, torch.Tensor],
- schedule_timesteps: Optional[torch.Tensor] = None,
+ timestep: int | torch.Tensor,
+ schedule_timesteps: torch.Tensor | None = None,
) -> int:
"""
Find the index for a given timestep in the schedule.
@@ -1174,7 +1171,7 @@ def index_for_timestep(
return step_index
- def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: int | torch.Tensor) -> None:
"""
Initialize the step_index counter for the scheduler.
@@ -1193,12 +1190,12 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[int, torch.Tensor],
+ timestep: int | torch.Tensor,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
- variance_noise: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ variance_noise: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the multistep DPMSolver.
diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py
index 66398073b29e..dd579e84d609 100644
--- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py
+++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py
@@ -15,7 +15,6 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver
from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
import flax
import jax
@@ -45,13 +44,13 @@ class DPMSolverMultistepSchedulerState:
# setable values
init_noise_sigma: jnp.ndarray
timesteps: jnp.ndarray
- num_inference_steps: Optional[int] = None
+ num_inference_steps: int = None
# running values
- model_outputs: Optional[jnp.ndarray] = None
- lower_order_nums: Optional[jnp.int32] = None
- prev_timestep: Optional[jnp.int32] = None
- cur_sample: Optional[jnp.ndarray] = None
+ model_outputs: jnp.ndarray | None = None
+ lower_order_nums: jnp.int32 | None = None
+ prev_timestep: jnp.int32 | None = None
+ cur_sample: jnp.ndarray | None = None
@classmethod
def create(
@@ -163,7 +162,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: str = "linear",
- trained_betas: Optional[jnp.ndarray] = None,
+ trained_betas: jnp.ndarray | None = None,
solver_order: int = 2,
prediction_type: str = "epsilon",
thresholding: bool = False,
@@ -181,7 +180,7 @@ def __init__(
)
self.dtype = dtype
- def create_state(self, common: Optional[CommonSchedulerState] = None) -> DPMSolverMultistepSchedulerState:
+ def create_state(self, common: CommonSchedulerState | None = None) -> DPMSolverMultistepSchedulerState:
if common is None:
common = CommonSchedulerState.create(self)
@@ -214,7 +213,7 @@ def set_timesteps(
self,
state: DPMSolverMultistepSchedulerState,
num_inference_steps: int,
- shape: Tuple,
+ shape: tuple,
) -> DPMSolverMultistepSchedulerState:
"""
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -224,7 +223,7 @@ def set_timesteps(
the `FlaxDPMSolverMultistepScheduler` state data class instance.
num_inference_steps (`int`):
the number of diffusion steps used when generating samples with a pre-trained model.
- shape (`Tuple`):
+ shape (`tuple`):
the shape of the samples to be generated.
"""
last_timestep = self.config.num_train_timesteps
@@ -379,7 +378,7 @@ def multistep_dpm_solver_second_order_update(
self,
state: DPMSolverMultistepSchedulerState,
model_output_list: jnp.ndarray,
- timestep_list: List[int],
+ timestep_list: list[int],
prev_timestep: int,
sample: jnp.ndarray,
) -> jnp.ndarray:
@@ -387,7 +386,7 @@ def multistep_dpm_solver_second_order_update(
One step for the second-order multistep DPM-Solver.
Args:
- model_output_list (`List[jnp.ndarray]`):
+ model_output_list (`list[jnp.ndarray]`):
direct outputs from learned diffusion model at current and latter timesteps.
timestep (`int`): current and latter discrete timestep in the diffusion chain.
prev_timestep (`int`): previous discrete timestep in the diffusion chain.
@@ -443,7 +442,7 @@ def multistep_dpm_solver_third_order_update(
self,
state: DPMSolverMultistepSchedulerState,
model_output_list: jnp.ndarray,
- timestep_list: List[int],
+ timestep_list: list[int],
prev_timestep: int,
sample: jnp.ndarray,
) -> jnp.ndarray:
@@ -451,7 +450,7 @@ def multistep_dpm_solver_third_order_update(
One step for the third-order multistep DPM-Solver.
Args:
- model_output_list (`List[jnp.ndarray]`):
+ model_output_list (`list[jnp.ndarray]`):
direct outputs from learned diffusion model at current and latter timesteps.
timestep (`int`): current and latter discrete timestep in the diffusion chain.
prev_timestep (`int`): previous discrete timestep in the diffusion chain.
@@ -507,7 +506,7 @@ def step(
timestep: int,
sample: jnp.ndarray,
return_dict: bool = True,
- ) -> Union[FlaxDPMSolverMultistepSchedulerOutput, Tuple]:
+ ) -> FlaxDPMSolverMultistepSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by DPM-Solver. Core function to propagate the diffusion process
from the learned model outputs (most often the predicted noise).
@@ -641,7 +640,7 @@ def scale_model_input(
self,
state: DPMSolverMultistepSchedulerState,
sample: jnp.ndarray,
- timestep: Optional[int] = None,
+ timestep: int | None = None,
) -> jnp.ndarray:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py
index e7f0d1a13a9c..3386514f64e9 100644
--- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py
+++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py
@@ -15,7 +15,7 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver
import math
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -168,7 +168,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
solver_order: int = 2,
prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon",
thresholding: bool = False,
@@ -178,13 +178,13 @@ def __init__(
solver_type: Literal["midpoint", "heun"] = "midpoint",
lower_order_final: bool = True,
euler_at_final: bool = False,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- use_flow_sigmas: Optional[bool] = False,
- flow_shift: Optional[float] = 1.0,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ use_flow_sigmas: bool = False,
+ flow_shift: float = 1.0,
lambda_min_clipped: float = -float("inf"),
- variance_type: Optional[Literal["learned", "learned_range"]] = None,
+ variance_type: Literal["learned", "learned_range"] | None = None,
timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
steps_offset: int = 0,
):
@@ -282,8 +282,8 @@ def step_index(self):
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -453,7 +453,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:
return t
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t
- def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Convert sigma values to alpha_t and sigma_t values.
@@ -462,7 +462,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to
The sigma value(s) to convert.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing (alpha_t, sigma_t) values.
"""
if self.config.use_flow_sigmas:
@@ -601,7 +601,7 @@ def convert_model_output(
self,
model_output: torch.Tensor,
*args,
- sample: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -700,8 +700,8 @@ def dpm_solver_first_order_update(
self,
model_output: torch.Tensor,
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -773,17 +773,17 @@ def dpm_solver_first_order_update(
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_second_order_update
def multistep_dpm_solver_second_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
One step for the second-order multistep DPMSolver.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`, *optional*):
A current instance of a sample created by the diffusion process.
@@ -897,17 +897,17 @@ def multistep_dpm_solver_second_order_update(
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_third_order_update
def multistep_dpm_solver_third_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
One step for the third-order multistep DPMSolver.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`, *optional*):
A current instance of a sample created by diffusion process.
@@ -992,7 +992,7 @@ def multistep_dpm_solver_third_order_update(
)
return x_t
- def _init_step_index(self, timestep: Union[int, torch.Tensor]):
+ def _init_step_index(self, timestep: int | torch.Tensor):
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
@@ -1014,12 +1014,12 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]):
def step(
self,
model_output: torch.Tensor,
- timestep: Union[int, torch.Tensor],
+ timestep: int | torch.Tensor,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
- variance_noise: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ variance_noise: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the multistep DPMSolver.
diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_sde.py b/src/diffusers/schedulers/scheduling_dpmsolver_sde.py
index a1d4a997d126..b288d9f61b14 100644
--- a/src/diffusers/schedulers/scheduling_dpmsolver_sde.py
+++ b/src/diffusers/schedulers/scheduling_dpmsolver_sde.py
@@ -45,7 +45,7 @@ class DPMSolverSDESchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
class BatchedBrownianTree:
@@ -243,12 +243,12 @@ def __init__(
beta_start: float = 0.00085, # sensible defaults
beta_end: float = 0.012,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon",
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- noise_sampler_seed: Optional[int] = None,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ noise_sampler_seed: int | None = None,
timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
steps_offset: int = 0,
):
@@ -302,7 +302,7 @@ def __init__(
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -332,7 +332,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -383,7 +383,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
def scale_model_input(
self,
sample: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
@@ -410,8 +410,8 @@ def scale_model_input(
def set_timesteps(
self,
num_inference_steps: int,
- device: Union[str, torch.device] = None,
- num_train_timesteps: Optional[int] = None,
+ device: str | torch.device = None,
+ num_train_timesteps: int | None = None,
) -> None:
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -655,11 +655,11 @@ def state_in_first_order(self) -> bool:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
sample: torch.Tensor,
return_dict: bool = True,
s_noise: float = 1.0,
- ) -> Union[DPMSolverSDESchedulerOutput, Tuple]:
+ ) -> DPMSolverSDESchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py
index 9ab362da3006..002f2b844a52 100644
--- a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py
+++ b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py
@@ -15,7 +15,7 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver
import math
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -102,7 +102,7 @@ class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin):
beta_schedule (`"linear"`, `"scaled_linear"`, or `"squaredcos_cap_v2"`, defaults to `"linear"`):
The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
`linear`, `scaled_linear`, or `squaredcos_cap_v2`.
- trained_betas (`np.ndarray` or `List[float]`, *optional*):
+ trained_betas (`np.ndarray` or `list[float]`, *optional*):
Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
solver_order (`int`, defaults to `2`):
The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided
@@ -168,7 +168,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
solver_order: int = 2,
prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon",
thresholding: bool = False,
@@ -177,14 +177,14 @@ def __init__(
algorithm_type: Literal["dpmsolver", "dpmsolver++", "sde-dpmsolver++"] = "dpmsolver++",
solver_type: Literal["midpoint", "heun"] = "midpoint",
lower_order_final: bool = False,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- use_flow_sigmas: Optional[bool] = False,
- flow_shift: Optional[float] = 1.0,
- final_sigmas_type: Optional[Literal["zero", "sigma_min"]] = "zero",
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ use_flow_sigmas: bool = False,
+ flow_shift: float = 1.0,
+ final_sigmas_type: Literal["zero", "sigma_min"] = "zero",
lambda_min_clipped: float = -float("inf"),
- variance_type: Optional[Literal["learned", "learned_range"]] = None,
+ variance_type: Literal["learned", "learned_range"] | None = None,
use_dynamic_shifting: bool = False,
time_shift_type: Literal["exponential"] = "exponential",
) -> None:
@@ -250,7 +250,7 @@ def __init__(
self._begin_index = None
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
- def get_order_list(self, num_inference_steps: int) -> List[int]:
+ def get_order_list(self, num_inference_steps: int) -> list[int]:
"""
Computes the solver order at each time step.
@@ -259,7 +259,7 @@ def get_order_list(self, num_inference_steps: int) -> List[int]:
The number of diffusion steps used when generating samples with a pre-trained model.
Returns:
- `List[int]`:
+ `list[int]`:
The list of solver orders for each timestep.
"""
steps = num_inference_steps
@@ -295,7 +295,7 @@ def get_order_list(self, num_inference_steps: int) -> List[int]:
return orders
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
@@ -306,7 +306,7 @@ def step_index(self) -> Optional[int]:
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
@@ -329,11 +329,11 @@ def set_begin_index(self, begin_index: int = 0) -> None:
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- mu: Optional[float] = None,
- timesteps: Optional[List[int]] = None,
- ) -> None:
+ num_inference_steps: int = None,
+ device: str | torch.device = None,
+ mu: float | None = None,
+ timesteps: list[int] | None = None,
+ ):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -342,9 +342,7 @@ def set_timesteps(
The number of diffusion steps used when generating samples with a pre-trained model.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- mu (`float`, *optional*):
- The mu parameter for dynamic shifting.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
timestep spacing strategy of equal spacing between timesteps schedule is used. If `timesteps` is
passed, `num_inference_steps` must be `None`.
@@ -519,7 +517,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:
return t
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t
- def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Convert sigma values to alpha_t and sigma_t values.
@@ -528,7 +526,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to
The sigma value(s) to convert.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing (alpha_t, sigma_t) values.
"""
if self.config.use_flow_sigmas:
@@ -666,7 +664,7 @@ def convert_model_output(
self,
model_output: torch.Tensor,
*args,
- sample: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -762,8 +760,8 @@ def dpm_solver_first_order_update(
self,
model_output: torch.Tensor,
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -824,10 +822,10 @@ def dpm_solver_first_order_update(
def singlestep_dpm_solver_second_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -835,7 +833,7 @@ def singlestep_dpm_solver_second_order_update(
time `timestep_list[-2]`.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
timestep (`int`):
The current and latter discrete timestep in the diffusion chain.
@@ -935,10 +933,10 @@ def singlestep_dpm_solver_second_order_update(
def singlestep_dpm_solver_third_order_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
- sample: Optional[torch.Tensor] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
@@ -946,7 +944,7 @@ def singlestep_dpm_solver_third_order_update(
time `timestep_list[-3]`.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
timestep (`int`):
The current and latter discrete timestep in the diffusion chain.
@@ -1057,18 +1055,18 @@ def singlestep_dpm_solver_third_order_update(
def singlestep_dpm_solver_update(
self,
- model_output_list: List[torch.Tensor],
+ model_output_list: list[torch.Tensor],
*args,
- sample: Optional[torch.Tensor] = None,
- order: Optional[int] = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor | None = None,
+ order: int = None,
+ noise: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
One step for the singlestep DPMSolver.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
timestep (`int`):
The current and latter discrete timestep in the diffusion chain.
@@ -1121,8 +1119,8 @@ def singlestep_dpm_solver_update(
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
def index_for_timestep(
self,
- timestep: Union[int, torch.Tensor],
- schedule_timesteps: Optional[torch.Tensor] = None,
+ timestep: int | torch.Tensor,
+ schedule_timesteps: torch.Tensor | None = None,
) -> int:
"""
Find the index for a given timestep in the schedule.
@@ -1156,7 +1154,7 @@ def index_for_timestep(
return step_index
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
- def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: int | torch.Tensor) -> None:
"""
Initialize the step_index counter for the scheduler.
@@ -1175,11 +1173,11 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[int, torch.Tensor],
+ timestep: int | torch.Tensor,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the singlestep DPMSolver.
diff --git a/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py
index 282540f5a23c..20ad17cfdd96 100644
--- a/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py
+++ b/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py
@@ -15,7 +15,7 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver and https://github.com/NVlabs/edm
import math
-from typing import List, Optional, Tuple, Union
+from typing import List, Literal
import numpy as np
import torch
@@ -51,13 +51,15 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
schedule was incorporated in this model: https://huggingface.co/stabilityai/cosxl.
num_train_timesteps (`int`, defaults to 1000):
The number of diffusion steps to train the model.
- solver_order (`int`, defaults to 2):
- The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided
- sampling, and `solver_order=3` for unconditional sampling.
prediction_type (`str`, defaults to `epsilon`, *optional*):
Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
`sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
Video](https://huggingface.co/papers/2210.02303) paper).
+ rho (`float`, *optional*, defaults to 7.0):
+ The rho parameter in the Karras sigma schedule. This was set to 7.0 in the EDM paper [1].
+ solver_order (`int`, defaults to 2):
+ The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided
+ sampling, and `solver_order=3` for unconditional sampling.
thresholding (`bool`, defaults to `False`):
Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
as Stable Diffusion.
@@ -94,19 +96,19 @@ def __init__(
sigma_min: float = 0.002,
sigma_max: float = 80.0,
sigma_data: float = 0.5,
- sigma_schedule: str = "karras",
+ sigma_schedule: Literal["karras", "exponential"] = "karras",
num_train_timesteps: int = 1000,
- prediction_type: str = "epsilon",
+ prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon",
rho: float = 7.0,
solver_order: int = 2,
thresholding: bool = False,
dynamic_thresholding_ratio: float = 0.995,
sample_max_value: float = 1.0,
- algorithm_type: str = "dpmsolver++",
- solver_type: str = "midpoint",
+ algorithm_type: Literal["dpmsolver++", "sde-dpmsolver++"] = "dpmsolver++",
+ solver_type: Literal["midpoint", "heun"] = "midpoint",
lower_order_final: bool = True,
euler_at_final: bool = False,
- final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min"
+ final_sigmas_type: Literal["zero", "sigma_min"] = "zero", # "zero", "sigma_min"
):
# settings for DPM-Solver
if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"]:
@@ -145,19 +147,19 @@ def __init__(
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
@property
- def init_noise_sigma(self):
+ def init_noise_sigma(self) -> float:
# standard deviation of the initial noise distribution
return (self.config.sigma_max**2 + 1) ** 0.5
@property
- def step_index(self):
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
"""
return self._step_index
@property
- def begin_index(self):
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
"""
@@ -175,7 +177,7 @@ def set_begin_index(self, begin_index: int = 0):
self._begin_index = begin_index
# Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_inputs
- def precondition_inputs(self, sample: torch.Tensor, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
+ def precondition_inputs(self, sample: torch.Tensor, sigma: float | torch.Tensor) -> torch.Tensor:
"""
Precondition the input sample by scaling it according to the EDM formulation.
@@ -194,7 +196,7 @@ def precondition_inputs(self, sample: torch.Tensor, sigma: Union[float, torch.Te
return scaled_sample
# Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_noise
- def precondition_noise(self, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
+ def precondition_noise(self, sigma: float | torch.Tensor) -> torch.Tensor:
"""
Precondition the noise level by applying a logarithmic transformation.
@@ -218,7 +220,7 @@ def precondition_outputs(
self,
sample: torch.Tensor,
model_output: torch.Tensor,
- sigma: Union[float, torch.Tensor],
+ sigma: float | torch.Tensor,
) -> torch.Tensor:
"""
Precondition the model outputs according to the EDM formulation.
@@ -250,7 +252,7 @@ def precondition_outputs(
return denoised
# Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.scale_model_input
- def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
"""
Scale the denoising model input to match the Euler algorithm. Ensures interchangeability with schedulers that
need to scale the denoising model input depending on the current timestep.
@@ -274,7 +276,11 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T
self.is_scale_input_called = True
return sample
- def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None):
+ def set_timesteps(
+ self,
+ num_inference_steps: int = None,
+ device: str | torch.device | None = None,
+ ):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -321,8 +327,8 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc
def _compute_karras_sigmas(
self,
ramp: torch.Tensor,
- sigma_min: Optional[float] = None,
- sigma_max: Optional[float] = None,
+ sigma_min: float | None = None,
+ sigma_max: float | None = None,
) -> torch.Tensor:
"""
Construct the noise schedule of [Karras et al. (2022)](https://huggingface.co/papers/2206.00364).
@@ -352,8 +358,8 @@ def _compute_karras_sigmas(
def _compute_exponential_sigmas(
self,
ramp: torch.Tensor,
- sigma_min: Optional[float] = None,
- sigma_max: Optional[float] = None,
+ sigma_min: float | None = None,
+ sigma_max: float | None = None,
) -> torch.Tensor:
"""
Compute the exponential sigma schedule. Implementation closely follows k-diffusion:
@@ -460,13 +466,12 @@ def _sigma_to_t(self, sigma, log_sigmas):
def _sigma_to_alpha_sigma_t(self, sigma):
alpha_t = torch.tensor(1) # Inputs are pre-scaled before going into unet, so alpha_t = 1
sigma_t = sigma
-
return alpha_t, sigma_t
def convert_model_output(
self,
model_output: torch.Tensor,
- sample: torch.Tensor = None,
+ sample: torch.Tensor,
) -> torch.Tensor:
"""
Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is
@@ -497,8 +502,8 @@ def convert_model_output(
def dpm_solver_first_order_update(
self,
model_output: torch.Tensor,
- sample: torch.Tensor = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor,
+ noise: torch.Tensor | None = None,
) -> torch.Tensor:
"""
One step for the first-order DPMSolver (equivalent to DDIM).
@@ -508,6 +513,8 @@ def dpm_solver_first_order_update(
The direct output from the learned diffusion model.
sample (`torch.Tensor`):
A current instance of a sample created by the diffusion process.
+ noise (`torch.Tensor`, *optional*):
+ The noise tensor to add to the original samples.
Returns:
`torch.Tensor`:
@@ -538,17 +545,19 @@ def dpm_solver_first_order_update(
def multistep_dpm_solver_second_order_update(
self,
model_output_list: List[torch.Tensor],
- sample: torch.Tensor = None,
- noise: Optional[torch.Tensor] = None,
+ sample: torch.Tensor,
+ noise: torch.Tensor | None = None,
) -> torch.Tensor:
"""
One step for the second-order multistep DPMSolver.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`):
A current instance of a sample created by the diffusion process.
+ noise (`torch.Tensor`, *optional*):
+ The noise tensor to add to the original samples.
Returns:
`torch.Tensor`:
@@ -608,14 +617,14 @@ def multistep_dpm_solver_second_order_update(
def multistep_dpm_solver_third_order_update(
self,
- model_output_list: List[torch.Tensor],
- sample: torch.Tensor = None,
+ model_output_list: list[torch.Tensor],
+ sample: torch.Tensor,
) -> torch.Tensor:
"""
One step for the third-order multistep DPMSolver.
Args:
- model_output_list (`List[torch.Tensor]`):
+ model_output_list (`list[torch.Tensor]`):
The direct outputs from learned diffusion model at current and latter timesteps.
sample (`torch.Tensor`):
A current instance of a sample created by diffusion process.
@@ -663,8 +672,8 @@ def multistep_dpm_solver_third_order_update(
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
def index_for_timestep(
self,
- timestep: Union[int, torch.Tensor],
- schedule_timesteps: Optional[torch.Tensor] = None,
+ timestep: int | torch.Tensor,
+ schedule_timesteps: torch.Tensor | None = None,
) -> int:
"""
Find the index for a given timestep in the schedule.
@@ -698,7 +707,7 @@ def index_for_timestep(
return step_index
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
- def _init_step_index(self, timestep):
+ def _init_step_index(self, timestep: int | torch.Tensor) -> None:
"""
Initialize the step_index counter for the scheduler.
@@ -717,11 +726,11 @@ def _init_step_index(self, timestep):
def step(
self,
model_output: torch.Tensor,
- timestep: Union[int, torch.Tensor],
+ timestep: int | torch.Tensor,
sample: torch.Tensor,
- generator=None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the multistep DPMSolver.
@@ -845,7 +854,7 @@ def add_noise(
return noisy_samples
# Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._get_conditioning_c_in
- def _get_conditioning_c_in(self, sigma: Union[float, torch.Tensor]) -> Union[float, torch.Tensor]:
+ def _get_conditioning_c_in(self, sigma: float | torch.Tensor) -> float | torch.Tensor:
"""
Compute the input conditioning factor for the EDM formulation.
@@ -860,5 +869,5 @@ def _get_conditioning_c_in(self, sigma: Union[float, torch.Tensor]) -> Union[flo
c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
return c_in
- def __len__(self):
+ def __len__(self) -> int:
return self.config.num_train_timesteps
diff --git a/src/diffusers/schedulers/scheduling_edm_euler.py b/src/diffusers/schedulers/scheduling_edm_euler.py
index 604d8b3ea6fa..dfb70ce68b46 100644
--- a/src/diffusers/schedulers/scheduling_edm_euler.py
+++ b/src/diffusers/schedulers/scheduling_edm_euler.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import torch
@@ -43,7 +43,7 @@ class EDMEulerSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
class EDMEulerScheduler(SchedulerMixin, ConfigMixin):
@@ -141,7 +141,7 @@ def init_noise_sigma(self) -> float:
return (self.config.sigma_max**2 + 1) ** 0.5
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
Return the index counter for the current timestep. The index will increase by 1 after each scheduler step.
@@ -152,7 +152,7 @@ def step_index(self) -> Optional[int]:
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
Return the index for the first timestep. This should be set from the pipeline with the `set_begin_index`
method.
@@ -174,7 +174,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
"""
self._begin_index = begin_index
- def precondition_inputs(self, sample: torch.Tensor, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
+ def precondition_inputs(self, sample: torch.Tensor, sigma: float | torch.Tensor) -> torch.Tensor:
"""
Precondition the input sample by scaling it according to the EDM formulation.
@@ -192,7 +192,7 @@ def precondition_inputs(self, sample: torch.Tensor, sigma: Union[float, torch.Te
scaled_sample = sample * c_in
return scaled_sample
- def precondition_noise(self, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
+ def precondition_noise(self, sigma: float | torch.Tensor) -> torch.Tensor:
"""
Precondition the noise level by applying a logarithmic transformation.
@@ -215,7 +215,7 @@ def precondition_outputs(
self,
sample: torch.Tensor,
model_output: torch.Tensor,
- sigma: Union[float, torch.Tensor],
+ sigma: float | torch.Tensor,
) -> torch.Tensor:
"""
Precondition the model outputs according to the EDM formulation.
@@ -246,7 +246,7 @@ def precondition_outputs(
return denoised
- def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
"""
Scale the denoising model input to match the Euler algorithm. Ensures interchangeability with schedulers that
need to scale the denoising model input depending on the current timestep.
@@ -272,10 +272,10 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- sigmas: Optional[Union[torch.Tensor, List[float]]] = None,
- ) -> None:
+ num_inference_steps: int = None,
+ device: str | torch.device = None,
+ sigmas: torch.Tensor | list[float] | None = None,
+ ):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -284,7 +284,7 @@ def set_timesteps(
The number of diffusion steps used when generating samples with a pre-trained model.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- sigmas (`torch.Tensor` or `List[float]`, *optional*):
+ sigmas (`torch.Tensor | list[float]`, *optional*):
Custom sigmas to use for the denoising process. If not defined, the default behavior when
`num_inference_steps` is passed will be used.
"""
@@ -323,8 +323,8 @@ def set_timesteps(
def _compute_karras_sigmas(
self,
ramp: torch.Tensor,
- sigma_min: Optional[float] = None,
- sigma_max: Optional[float] = None,
+ sigma_min: float | None = None,
+ sigma_max: float | None = None,
) -> torch.Tensor:
"""
Construct the noise schedule of [Karras et al. (2022)](https://huggingface.co/papers/2206.00364).
@@ -353,8 +353,8 @@ def _compute_karras_sigmas(
def _compute_exponential_sigmas(
self,
ramp: torch.Tensor,
- sigma_min: Optional[float] = None,
- sigma_max: Optional[float] = None,
+ sigma_min: float | None = None,
+ sigma_max: float | None = None,
) -> torch.Tensor:
"""
Compute the exponential sigma schedule. Implementation closely follows k-diffusion:
@@ -379,7 +379,7 @@ def _compute_exponential_sigmas(
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -409,7 +409,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -427,16 +427,16 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
sample: torch.Tensor,
s_churn: float = 0.0,
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- pred_original_sample: Optional[torch.Tensor] = None,
- ) -> Union[EDMEulerSchedulerOutput, Tuple[torch.Tensor, torch.Tensor]]:
+ pred_original_sample: torch.Tensor | None = None,
+ ) -> EDMEulerSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -581,7 +581,7 @@ def add_noise(
noisy_samples = original_samples + noise * sigma
return noisy_samples
- def _get_conditioning_c_in(self, sigma: Union[float, torch.Tensor]) -> Union[float, torch.Tensor]:
+ def _get_conditioning_c_in(self, sigma: float | torch.Tensor) -> float | torch.Tensor:
"""
Compute the input conditioning factor for the EDM formulation.
diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py
index 11fec60c9c0c..8ad6abd90668 100644
--- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py
+++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -44,7 +44,7 @@ class EulerAncestralDiscreteSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -180,10 +180,10 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.0001,
beta_end: float = 0.02,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
- prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon",
- timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
+ prediction_type: str = "epsilon",
+ timestep_spacing: str = "linspace",
steps_offset: int = 0,
rescale_betas_zero_snr: bool = False,
) -> None:
@@ -234,14 +234,14 @@ def init_noise_sigma(self) -> torch.Tensor:
return (self.sigmas.max() ** 2 + 1) ** 0.5
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
"""
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
"""
@@ -258,7 +258,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
"""
self._begin_index = begin_index
- def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep. Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm.
@@ -282,7 +282,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T
self.is_scale_input_called = True
return sample
- def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None) -> None:
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -328,7 +328,7 @@ def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, to
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -358,7 +358,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -376,11 +376,11 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]:
+ ) -> EulerAncestralDiscreteSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py
index 8b141325fbd3..484f2ca58e1b 100644
--- a/src/diffusers/schedulers/scheduling_euler_discrete.py
+++ b/src/diffusers/schedulers/scheduling_euler_discrete.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -47,7 +47,7 @@ class EulerDiscreteSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -205,15 +205,15 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.0001,
beta_end: float = 0.02,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
- prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon",
- interpolation_type: Literal["linear", "log_linear"] = "linear",
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- sigma_min: Optional[float] = None,
- sigma_max: Optional[float] = None,
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
+ prediction_type: str = "epsilon",
+ interpolation_type: str = "linear",
+ use_karras_sigmas: bool | None = False,
+ use_exponential_sigmas: bool | None = False,
+ use_beta_sigmas: bool | None = False,
+ sigma_min: float | None = None,
+ sigma_max: float | None = None,
timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
timestep_type: Literal["discrete", "continuous"] = "discrete",
steps_offset: int = 0,
@@ -275,7 +275,7 @@ def __init__(
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
@property
- def init_noise_sigma(self) -> Union[float, torch.Tensor]:
+ def init_noise_sigma(self) -> float | torch.Tensor:
"""
The standard deviation of the initial noise distribution.
@@ -291,7 +291,7 @@ def init_noise_sigma(self) -> Union[float, torch.Tensor]:
return (max_sigma**2 + 1) ** 0.5
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase by 1 after each scheduler step.
@@ -302,7 +302,7 @@ def step_index(self) -> Optional[int]:
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
@@ -323,7 +323,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
"""
self._begin_index = begin_index
- def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep. Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm.
@@ -349,11 +349,11 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- sigmas: Optional[List[float]] = None,
- ) -> None:
+ num_inference_steps: int = None,
+ device: str | torch.device = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ ):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -363,14 +363,15 @@ def set_timesteps(
`timesteps` or `sigmas` must be provided.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated
based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas`
must be `None`, and `timestep_spacing` attribute will be ignored.
- sigmas (`List[float]`, *optional*):
- Custom sigmas used to support arbitrary timesteps schedule. If `None`, timesteps and sigmas will be
- generated based on the relevant scheduler attributes. If `sigmas` is passed, `num_inference_steps` and
- `timesteps` must be `None`, and the timesteps will be generated based on the custom sigmas schedule.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to support arbitrary timesteps schedule schedule. If `None`, timesteps and sigmas
+ will be generated based on the relevant scheduler attributes. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`, and the timesteps will be generated based on the
+ custom sigmas schedule.
"""
if timesteps is not None and sigmas is not None:
@@ -638,7 +639,7 @@ def _convert_to_beta(
return sigmas
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -667,7 +668,7 @@ def index_for_timestep(
return indices[pos].item()
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -685,15 +686,15 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
sample: torch.Tensor,
s_churn: float = 0.0,
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[EulerDiscreteSchedulerOutput, Tuple]:
+ ) -> EulerDiscreteSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_euler_discrete_flax.py b/src/diffusers/schedulers/scheduling_euler_discrete_flax.py
index 2bb6bf35585c..fc555dd21395 100644
--- a/src/diffusers/schedulers/scheduling_euler_discrete_flax.py
+++ b/src/diffusers/schedulers/scheduling_euler_discrete_flax.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import flax
import jax.numpy as jnp
@@ -40,7 +39,7 @@ class EulerDiscreteSchedulerState:
init_noise_sigma: jnp.ndarray
timesteps: jnp.ndarray
sigmas: jnp.ndarray
- num_inference_steps: Optional[int] = None
+ num_inference_steps: int = None
@classmethod
def create(
@@ -107,7 +106,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: str = "linear",
- trained_betas: Optional[jnp.ndarray] = None,
+ trained_betas: jnp.ndarray | None = None,
prediction_type: str = "epsilon",
timestep_spacing: str = "linspace",
dtype: jnp.dtype = jnp.float32,
@@ -118,7 +117,7 @@ def __init__(
)
self.dtype = dtype
- def create_state(self, common: Optional[CommonSchedulerState] = None) -> EulerDiscreteSchedulerState:
+ def create_state(self, common: CommonSchedulerState | None = None) -> EulerDiscreteSchedulerState:
if common is None:
common = CommonSchedulerState.create(self)
@@ -166,7 +165,7 @@ def set_timesteps(
self,
state: EulerDiscreteSchedulerState,
num_inference_steps: int,
- shape: Tuple = (),
+ shape: tuple = (),
) -> EulerDiscreteSchedulerState:
"""
Sets the timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -218,7 +217,7 @@ def step(
timestep: int,
sample: jnp.ndarray,
return_dict: bool = True,
- ) -> Union[FlaxEulerDiscreteSchedulerOutput, Tuple]:
+ ) -> FlaxEulerDiscreteSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py
index 378a62ca8aee..1021abf0f6f6 100644
--- a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py
+++ b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
+from typing import Literal, Optional, Union
import numpy as np
import torch
@@ -93,21 +93,30 @@ def __init__(
num_train_timesteps: int = 1000,
shift: float = 1.0,
use_dynamic_shifting: bool = False,
- base_shift: Optional[float] = 0.5,
- max_shift: Optional[float] = 1.15,
- base_image_seq_len: Optional[int] = 256,
- max_image_seq_len: Optional[int] = 4096,
+ base_shift: float | None = 0.5,
+ max_shift: float | None = 1.15,
+ base_image_seq_len: int = 256,
+ max_image_seq_len: int = 4096,
invert_sigmas: bool = False,
- shift_terminal: Optional[float] = None,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- time_shift_type: str = "exponential",
+ shift_terminal: float = None,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ time_shift_type: Literal["exponential", "linear"] = "exponential",
stochastic_sampling: bool = False,
):
if self.config.use_beta_sigmas and not is_scipy_available():
raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
- if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
+ if (
+ sum(
+ [
+ self.config.use_beta_sigmas,
+ self.config.use_exponential_sigmas,
+ self.config.use_karras_sigmas,
+ ]
+ )
+ > 1
+ ):
raise ValueError(
"Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
)
@@ -166,13 +175,20 @@ def set_begin_index(self, begin_index: int = 0):
self._begin_index = begin_index
def set_shift(self, shift: float):
+ """
+ Sets the shift value for the scheduler.
+
+ Args:
+ shift (`float`):
+ The shift value to be set.
+ """
self._shift = shift
def scale_noise(
self,
sample: torch.FloatTensor,
- timestep: torch.FloatTensor,
- noise: torch.FloatTensor,
+ timestep: float | torch.FloatTensor,
+ noise: torch.FloatTensor | None = None,
) -> torch.FloatTensor:
"""
Forward process in flow-matching
@@ -218,10 +234,25 @@ def scale_noise(
return sample
- def _sigma_to_t(self, sigma):
+ def _sigma_to_t(self, sigma) -> float:
return sigma * self.config.num_train_timesteps
- def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+ def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+ """
+ Apply time shifting to the sigmas.
+
+ Args:
+ mu (`float`):
+ The mu parameter for the time shift.
+ sigma (`float`):
+ The sigma parameter for the time shift.
+ t (`torch.Tensor`):
+ The input timesteps.
+
+ Returns:
+ `torch.Tensor`:
+ The time-shifted timesteps.
+ """
if self.config.time_shift_type == "exponential":
return self._time_shift_exponential(mu, sigma, t)
elif self.config.time_shift_type == "linear":
@@ -250,11 +281,11 @@ def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device] = None,
- sigmas: Optional[List[float]] = None,
- mu: Optional[float] = None,
- timesteps: Optional[List[float]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device = None,
+ sigmas: list[float] | None = None,
+ mu: float | None = None,
+ timesteps: list[float] | None = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -264,13 +295,13 @@ def set_timesteps(
The number of diffusion steps used when generating samples with a pre-trained model.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
automatically.
mu (`float`, *optional*):
Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
shifting.
- timesteps (`List[float]`, *optional*):
+ timesteps (`list[float]`, *optional*):
Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
automatically.
"""
@@ -302,7 +333,9 @@ def set_timesteps(
if sigmas is None:
if timesteps is None:
timesteps = np.linspace(
- self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+ self._sigma_to_t(self.sigma_max),
+ self._sigma_to_t(self.sigma_min),
+ num_inference_steps,
)
sigmas = timesteps / self.config.num_train_timesteps
else:
@@ -350,7 +383,24 @@ def set_timesteps(
self._step_index = None
self._begin_index = None
- def index_for_timestep(self, timestep, schedule_timesteps=None):
+ def index_for_timestep(
+ self,
+ timestep: Union[float, torch.FloatTensor],
+ schedule_timesteps: Optional[torch.FloatTensor] = None,
+ ) -> int:
+ """
+ Get the index for the given timestep.
+
+ Args:
+ timestep (`float` or `torch.FloatTensor`):
+ The timestep to find the index for.
+ schedule_timesteps (`torch.FloatTensor`, *optional*):
+ The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used.
+
+ Returns:
+ `int`:
+ The index of the timestep.
+ """
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
@@ -364,7 +414,7 @@ def index_for_timestep(self, timestep, schedule_timesteps=None):
return indices[pos].item()
- def _init_step_index(self, timestep):
+ def _init_step_index(self, timestep: Union[float, torch.FloatTensor]) -> None:
if self.begin_index is None:
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
@@ -375,16 +425,16 @@ def _init_step_index(self, timestep):
def step(
self,
model_output: torch.FloatTensor,
- timestep: Union[float, torch.FloatTensor],
+ timestep: float | torch.FloatTensor,
sample: torch.FloatTensor,
s_churn: float = 0.0,
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
- generator: Optional[torch.Generator] = None,
- per_token_timesteps: Optional[torch.Tensor] = None,
+ generator: torch.Generator | None = None,
+ per_token_timesteps: torch.Tensor | None = None,
return_dict: bool = True,
- ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
+ ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -405,7 +455,7 @@ def step(
A random number generator.
per_token_timesteps (`torch.Tensor`, *optional*):
The timesteps for each token in the sample.
- return_dict (`bool`):
+ return_dict (`bool`, defaults to `True`):
Whether or not to return a
[`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
@@ -474,7 +524,7 @@ def step(
return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
- def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
+ def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
"""
Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative
Models](https://huggingface.co/papers/2206.00364).
@@ -595,11 +645,11 @@ def _convert_to_beta(
)
return sigmas
- def _time_shift_exponential(self, mu, sigma, t):
+ def _time_shift_exponential(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
- def _time_shift_linear(self, mu, sigma, t):
+ def _time_shift_linear(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
return mu / (mu + (1 / t - 1) ** sigma)
- def __len__(self):
+ def __len__(self) -> int:
return self.config.num_train_timesteps
diff --git a/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py
index 6b85194f8b5e..49892e043abd 100644
--- a/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py
+++ b/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -51,9 +50,6 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
Args:
num_train_timesteps (`int`, defaults to 1000):
The number of diffusion steps to train the model.
- timestep_spacing (`str`, defaults to `"linspace"`):
- The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
- Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
shift (`float`, defaults to 1.0):
The shift value for the timestep schedule.
"""
@@ -110,7 +106,7 @@ def set_begin_index(self, begin_index: int = 0):
def scale_noise(
self,
sample: torch.FloatTensor,
- timestep: torch.FloatTensor,
+ timestep: float | torch.FloatTensor,
noise: torch.FloatTensor,
) -> torch.FloatTensor:
"""
@@ -119,7 +115,7 @@ def scale_noise(
Args:
sample (`torch.FloatTensor`):
The input sample.
- timestep (`torch.FloatTensor`):
+ timestep (`float` or `torch.FloatTensor`):
The current timestep in the diffusion chain.
noise (`torch.FloatTensor`):
The noise tensor.
@@ -137,10 +133,14 @@ def scale_noise(
return sample
- def _sigma_to_t(self, sigma):
+ def _sigma_to_t(self, sigma: float) -> float:
return sigma * self.config.num_train_timesteps
- def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+ def set_timesteps(
+ self,
+ num_inference_steps: int,
+ device: str | torch.device = None,
+ ) -> None:
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -153,7 +153,9 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
self.num_inference_steps = num_inference_steps
timesteps = np.linspace(
- self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+ self._sigma_to_t(self.sigma_max),
+ self._sigma_to_t(self.sigma_min),
+ num_inference_steps,
)
sigmas = timesteps / self.config.num_train_timesteps
@@ -174,7 +176,24 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
self._step_index = None
self._begin_index = None
- def index_for_timestep(self, timestep, schedule_timesteps=None):
+ def index_for_timestep(
+ self,
+ timestep: float | torch.FloatTensor,
+ schedule_timesteps: torch.FloatTensor | None = None,
+ ) -> int:
+ """
+ Find the index of a given timestep in the timestep schedule.
+
+ Args:
+ timestep (`float` or `torch.FloatTensor`):
+ The timestep value to find in the schedule.
+ schedule_timesteps (`torch.FloatTensor`, *optional*):
+ The timestep schedule to search in. If `None`, uses `self.timesteps`.
+
+ Returns:
+ `int`:
+ The index of the timestep in the schedule.
+ """
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
@@ -188,7 +207,7 @@ def index_for_timestep(self, timestep, schedule_timesteps=None):
return indices[pos].item()
- def _init_step_index(self, timestep):
+ def _init_step_index(self, timestep: float | torch.FloatTensor) -> None:
if self.begin_index is None:
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
@@ -197,21 +216,24 @@ def _init_step_index(self, timestep):
self._step_index = self._begin_index
@property
- def state_in_first_order(self):
+ def state_in_first_order(self) -> bool:
+ """
+ Returns whether the scheduler is in the first-order state.
+ """
return self.dt is None
def step(
self,
model_output: torch.FloatTensor,
- timestep: Union[float, torch.FloatTensor],
+ timestep: float | torch.FloatTensor,
sample: torch.FloatTensor,
s_churn: float = 0.0,
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]:
+ ) -> FlowMatchHeunDiscreteSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -219,13 +241,19 @@ def step(
Args:
model_output (`torch.FloatTensor`):
The direct output from learned diffusion model.
- timestep (`float`):
+ timestep (`float` or `torch.FloatTensor`):
The current discrete timestep in the diffusion chain.
sample (`torch.FloatTensor`):
A current instance of a sample created by the diffusion process.
s_churn (`float`):
- s_tmin (`float`):
- s_tmax (`float`):
+ Stochasticity parameter that controls the amount of noise added during sampling. Higher values increase
+ randomness.
+ s_tmin (`float`):
+ Minimum timestep threshold for applying stochasticity. Only timesteps above this value will have noise
+ added.
+ s_tmax (`float`):
+ Maximum timestep threshold for applying stochasticity. Only timesteps below this value will have noise
+ added.
s_noise (`float`, defaults to 1.0):
Scaling factor for noise added to the sample.
generator (`torch.Generator`, *optional*):
@@ -274,7 +302,10 @@ def step(
if gamma > 0:
noise = randn_tensor(
- model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
+ model_output.shape,
+ dtype=model_output.dtype,
+ device=model_output.device,
+ generator=generator,
)
eps = noise * s_noise
sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
@@ -320,5 +351,5 @@ def step(
return FlowMatchHeunDiscreteSchedulerOutput(prev_sample=prev_sample)
- def __len__(self):
+ def __len__(self) -> int:
return self.config.num_train_timesteps
diff --git a/src/diffusers/schedulers/scheduling_flow_match_lcm.py b/src/diffusers/schedulers/scheduling_flow_match_lcm.py
index 8ef0e2ec8175..97d4ebbc8e42 100644
--- a/src/diffusers/schedulers/scheduling_flow_match_lcm.py
+++ b/src/diffusers/schedulers/scheduling_flow_match_lcm.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -42,7 +42,7 @@ class FlowMatchLCMSchedulerOutput(BaseOutput):
denoising loop.
"""
- prev_sample: torch.FloatTensor
+ prev_sample: torch.Tensor
class FlowMatchLCMScheduler(SchedulerMixin, ConfigMixin):
@@ -80,11 +80,11 @@ class FlowMatchLCMScheduler(SchedulerMixin, ConfigMixin):
use_beta_sigmas (`bool`, defaults to False):
Whether to use beta sigmas for step sizes in the noise schedule during sampling.
time_shift_type (`str`, defaults to "exponential"):
- The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
- scale_factors ('list', defaults to None)
+ The type of dynamic resolution-dependent timestep shifting to apply.
+ scale_factors (`list[float]`, *optional*, defaults to `None`):
It defines how to scale the latents at which predictions are made.
- upscale_mode ('str', defaults to 'bicubic')
- Upscaling method, applied if scale-wise generation is considered
+ upscale_mode (`str`, *optional*, defaults to "bicubic"):
+ Upscaling method, applied if scale-wise generation is considered.
"""
_compatibles = []
@@ -96,22 +96,39 @@ def __init__(
num_train_timesteps: int = 1000,
shift: float = 1.0,
use_dynamic_shifting: bool = False,
- base_shift: Optional[float] = 0.5,
- max_shift: Optional[float] = 1.15,
- base_image_seq_len: Optional[int] = 256,
- max_image_seq_len: Optional[int] = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+ base_image_seq_len: int = 256,
+ max_image_seq_len: int = 4096,
invert_sigmas: bool = False,
- shift_terminal: Optional[float] = None,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- time_shift_type: str = "exponential",
- scale_factors: Optional[List[float]] = None,
- upscale_mode: Optional[str] = "bicubic",
+ shift_terminal: float | None = None,
+ use_karras_sigmas: bool | None = False,
+ use_exponential_sigmas: bool | None = False,
+ use_beta_sigmas: bool | None = False,
+ time_shift_type: Literal["exponential", "linear"] = "exponential",
+ scale_factors: list[float] | None = None,
+ upscale_mode: Literal[
+ "nearest",
+ "linear",
+ "bilinear",
+ "bicubic",
+ "trilinear",
+ "area",
+ "nearest-exact",
+ ] = "bicubic",
):
if self.config.use_beta_sigmas and not is_scipy_available():
raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
- if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
+ if (
+ sum(
+ [
+ self.config.use_beta_sigmas,
+ self.config.use_exponential_sigmas,
+ self.config.use_karras_sigmas,
+ ]
+ )
+ > 1
+ ):
raise ValueError(
"Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
)
@@ -163,7 +180,7 @@ def begin_index(self):
return self._begin_index
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
- def set_begin_index(self, begin_index: int = 0):
+ def set_begin_index(self, begin_index: int = 0) -> None:
"""
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
@@ -173,18 +190,18 @@ def set_begin_index(self, begin_index: int = 0):
"""
self._begin_index = begin_index
- def set_shift(self, shift: float):
+ def set_shift(self, shift: float) -> None:
self._shift = shift
- def set_scale_factors(self, scale_factors: list, upscale_mode):
+ def set_scale_factors(self, scale_factors: list[float], upscale_mode: str) -> None:
"""
Sets scale factors for a scale-wise generation regime.
Args:
- scale_factors (`list`):
- The scale factors for each step
+ scale_factors (`list[float]`):
+ The scale factors for each step.
upscale_mode (`str`):
- Upscaling method
+ Upscaling method.
"""
self._scale_factors = scale_factors
self._upscale_mode = upscale_mode
@@ -192,8 +209,8 @@ def set_scale_factors(self, scale_factors: list, upscale_mode):
def scale_noise(
self,
sample: torch.FloatTensor,
- timestep: torch.FloatTensor,
- noise: torch.FloatTensor,
+ timestep: float | torch.FloatTensor,
+ noise: torch.FloatTensor | None = None,
) -> torch.FloatTensor:
"""
Forward process in flow-matching
@@ -239,16 +256,18 @@ def scale_noise(
return sample
- def _sigma_to_t(self, sigma):
+ def _sigma_to_t(self, sigma: float | torch.FloatTensor) -> float | torch.FloatTensor:
return sigma * self.config.num_train_timesteps
- def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+ def time_shift(
+ self, mu: float, sigma: float, t: float | np.ndarray | torch.Tensor
+ ) -> float | np.ndarray | torch.Tensor:
if self.config.time_shift_type == "exponential":
return self._time_shift_exponential(mu, sigma, t)
elif self.config.time_shift_type == "linear":
return self._time_shift_linear(mu, sigma, t)
- def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
+ def stretch_shift_to_terminal(self, t: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
r"""
Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config
value.
@@ -257,12 +276,13 @@ def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51
Args:
- t (`torch.Tensor`):
- A tensor of timesteps to be stretched and shifted.
+ t (`torch.Tensor` or `np.ndarray`):
+ A tensor or numpy array of timesteps to be stretched and shifted.
Returns:
- `torch.Tensor`:
- A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.
+ `torch.Tensor` or `np.ndarray`:
+ A tensor or numpy array of adjusted timesteps such that the final value equals
+ `self.config.shift_terminal`.
"""
one_minus_z = 1 - t
scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
@@ -271,12 +291,12 @@ def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device] = None,
- sigmas: Optional[List[float]] = None,
- mu: Optional[float] = None,
- timesteps: Optional[List[float]] = None,
- ):
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ sigmas: list[float] | None = None,
+ mu: float | None = None,
+ timesteps: list[float] | None = None,
+ ) -> None:
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -285,13 +305,13 @@ def set_timesteps(
The number of diffusion steps used when generating samples with a pre-trained model.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- sigmas (`List[float]`, *optional*):
+ sigmas (`list[float]`, *optional*):
Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
automatically.
mu (`float`, *optional*):
Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
shifting.
- timesteps (`List[float]`, *optional*):
+ timesteps (`list[float]`, *optional*):
Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
automatically.
"""
@@ -318,43 +338,45 @@ def set_timesteps(
is_timesteps_provided = timesteps is not None
if is_timesteps_provided:
- timesteps = np.array(timesteps).astype(np.float32)
+ timesteps = np.array(timesteps).astype(np.float32) # type: ignore
if sigmas is None:
if timesteps is None:
- timesteps = np.linspace(
- self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+ timesteps = np.linspace( # type: ignore
+ self._sigma_to_t(self.sigma_max),
+ self._sigma_to_t(self.sigma_min),
+ num_inference_steps,
)
- sigmas = timesteps / self.config.num_train_timesteps
+ sigmas = timesteps / self.config.num_train_timesteps # type: ignore
else:
- sigmas = np.array(sigmas).astype(np.float32)
+ sigmas = np.array(sigmas).astype(np.float32) # type: ignore
num_inference_steps = len(sigmas)
# 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
# "exponential" or "linear" type is applied
if self.config.use_dynamic_shifting:
- sigmas = self.time_shift(mu, 1.0, sigmas)
+ sigmas = self.time_shift(mu, 1.0, sigmas) # type: ignore
else:
- sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)
+ sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) # type: ignore
# 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
if self.config.shift_terminal:
- sigmas = self.stretch_shift_to_terminal(sigmas)
+ sigmas = self.stretch_shift_to_terminal(sigmas) # type: ignore
# 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
if self.config.use_karras_sigmas:
- sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+ sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) # type: ignore
elif self.config.use_exponential_sigmas:
- sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+ sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) # type: ignore
elif self.config.use_beta_sigmas:
- sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+ sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) # type: ignore
# 5. Convert sigmas and timesteps to tensors and move to specified device
- sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+ sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) # type: ignore
if not is_timesteps_provided:
- timesteps = sigmas * self.config.num_train_timesteps
+ timesteps = sigmas * self.config.num_train_timesteps # type: ignore
else:
- timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device)
+ timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device) # type: ignore
# 6. Append the terminal sigma value.
# If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
@@ -371,7 +393,11 @@ def set_timesteps(
self._step_index = None
self._begin_index = None
- def index_for_timestep(self, timestep, schedule_timesteps=None):
+ def index_for_timestep(
+ self,
+ timestep: float | torch.Tensor,
+ schedule_timesteps: torch.Tensor | None = None,
+ ) -> int:
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
@@ -383,9 +409,9 @@ def index_for_timestep(self, timestep, schedule_timesteps=None):
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
pos = 1 if len(indices) > 1 else 0
- return indices[pos].item()
+ return int(indices[pos].item())
- def _init_step_index(self, timestep):
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
if self.begin_index is None:
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
@@ -396,11 +422,11 @@ def _init_step_index(self, timestep):
def step(
self,
model_output: torch.FloatTensor,
- timestep: Union[float, torch.FloatTensor],
+ timestep: float | torch.FloatTensor,
sample: torch.FloatTensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[FlowMatchLCMSchedulerOutput, Tuple]:
+ ) -> FlowMatchLCMSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -460,7 +486,12 @@ def step(
size = [round(self._scale_factors[self._step_index] * size) for size in self._init_size]
x0_pred = torch.nn.functional.interpolate(x0_pred, size=size, mode=self._upscale_mode)
- noise = randn_tensor(x0_pred.shape, generator=generator, device=x0_pred.device, dtype=x0_pred.dtype)
+ noise = randn_tensor(
+ x0_pred.shape,
+ generator=generator,
+ device=x0_pred.device,
+ dtype=x0_pred.dtype,
+ )
prev_sample = (1 - sigma_next) * x0_pred + sigma_next * noise
# upon completion increase step index by one
@@ -474,7 +505,7 @@ def step(
return FlowMatchLCMSchedulerOutput(prev_sample=prev_sample)
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
- def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
+ def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
"""
Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative
Models](https://huggingface.co/papers/2206.00364).
@@ -595,11 +626,15 @@ def _convert_to_beta(
)
return sigmas
- def _time_shift_exponential(self, mu, sigma, t):
+ def _time_shift_exponential(
+ self, mu: float, sigma: float, t: float | np.ndarray | torch.Tensor
+ ) -> float | np.ndarray | torch.Tensor:
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
- def _time_shift_linear(self, mu, sigma, t):
+ def _time_shift_linear(
+ self, mu: float, sigma: float, t: float | np.ndarray | torch.Tensor
+ ) -> float | np.ndarray | torch.Tensor:
return mu / (mu + (1 / t - 1) ** sigma)
- def __len__(self):
+ def __len__(self) -> int:
return self.config.num_train_timesteps
diff --git a/src/diffusers/schedulers/scheduling_heun_discrete.py b/src/diffusers/schedulers/scheduling_heun_discrete.py
index 0c5e28ad067d..df2593493aeb 100644
--- a/src/diffusers/schedulers/scheduling_heun_discrete.py
+++ b/src/diffusers/schedulers/scheduling_heun_discrete.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -44,7 +44,7 @@ class HeunDiscreteSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -151,13 +151,13 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.00085, # sensible defaults
beta_end: float = 0.012,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2", "exp"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
- prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon",
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- clip_sample: Optional[bool] = False,
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
+ prediction_type: str = "epsilon",
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ clip_sample: bool = False,
clip_sample_range: float = 1.0,
timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
steps_offset: int = 0,
@@ -196,7 +196,7 @@ def __init__(
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -261,7 +261,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
def scale_model_input(
self,
sample: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
@@ -286,11 +286,11 @@ def scale_model_input(
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device] = None,
- num_train_timesteps: Optional[int] = None,
- timesteps: Optional[List[int]] = None,
- ) -> None:
+ num_inference_steps: int | None = None,
+ device: str | torch.device = None,
+ num_train_timesteps: int | None = None,
+ timesteps: list[int] | None = None,
+ ):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -302,7 +302,7 @@ def set_timesteps(
num_train_timesteps (`int`, *optional*, defaults to `None`):
The number of diffusion steps used when training the model. If `None`, the default
`num_train_timesteps` attribute is used.
- timesteps (`List[int]`, *optional*, defaults to `None`):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, timesteps will be
generated based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps`
must be `None`, and `timestep_spacing` attribute will be ignored.
@@ -540,7 +540,7 @@ def state_in_first_order(self):
return self.dt is None
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -557,11 +557,11 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
def step(
self,
- model_output: Union[torch.Tensor, np.ndarray],
- timestep: Union[float, torch.Tensor],
- sample: Union[torch.Tensor, np.ndarray],
+ model_output: torch.Tensor | np.ndarray,
+ timestep: float | torch.Tensor,
+ sample: torch.Tensor | np.ndarray,
return_dict: bool = True,
- ) -> Union[HeunDiscreteSchedulerOutput, Tuple]:
+ ) -> HeunDiscreteSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_ipndm.py b/src/diffusers/schedulers/scheduling_ipndm.py
index da188fe8297c..0a02311ba9b6 100644
--- a/src/diffusers/schedulers/scheduling_ipndm.py
+++ b/src/diffusers/schedulers/scheduling_ipndm.py
@@ -13,7 +13,6 @@
# limitations under the License.
import math
-from typing import List, Optional, Tuple, Union
import numpy as np
import torch
@@ -39,9 +38,7 @@ class IPNDMScheduler(SchedulerMixin, ConfigMixin):
order = 1
@register_to_config
- def __init__(
- self, num_train_timesteps: int = 1000, trained_betas: Optional[Union[np.ndarray, List[float]]] = None
- ):
+ def __init__(self, num_train_timesteps: int = 1000, trained_betas: np.ndarray | list[float] | None = None):
# set `betas`, `alphas`, `timesteps`
self.set_timesteps(num_train_timesteps)
@@ -83,7 +80,7 @@ def set_begin_index(self, begin_index: int = 0):
"""
self._begin_index = begin_index
- def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device | None = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -113,7 +110,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -143,7 +140,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -161,10 +158,10 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[int, torch.Tensor],
+ timestep: int | torch.Tensor,
sample: torch.Tensor,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the linear multistep method. It performs one forward pass multiple times to approximate the solution.
diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py
index ee49ae67b9cb..706190e291a4 100644
--- a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py
+++ b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -45,7 +45,7 @@ class KDPM2AncestralDiscreteSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -150,10 +150,10 @@ def __init__(
beta_start: float = 0.00085, # sensible defaults
beta_end: float = 0.012,
beta_schedule: str = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
+ trained_betas: np.ndarray | list[float] | None = None,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
prediction_type: str = "epsilon",
timestep_spacing: str = "linspace",
steps_offset: int = 0,
@@ -222,7 +222,7 @@ def set_begin_index(self, begin_index: int = 0):
def scale_model_input(
self,
sample: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
@@ -252,8 +252,8 @@ def scale_model_input(
def set_timesteps(
self,
num_inference_steps: int,
- device: Union[str, torch.device] = None,
- num_train_timesteps: Optional[int] = None,
+ device: str | torch.device | None = None,
+ num_train_timesteps: int | None = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -513,7 +513,7 @@ def state_in_first_order(self):
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -543,7 +543,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -560,12 +560,12 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
def step(
self,
- model_output: Union[torch.Tensor, np.ndarray],
- timestep: Union[float, torch.Tensor],
- sample: Union[torch.Tensor, np.ndarray],
- generator: Optional[torch.Generator] = None,
+ model_output: torch.Tensor | np.ndarray,
+ timestep: float | torch.Tensor,
+ sample: torch.Tensor | np.ndarray,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[KDPM2AncestralDiscreteSchedulerOutput, Tuple]:
+ ) -> KDPM2AncestralDiscreteSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py
index 6effb3699b9a..3092c610b46d 100644
--- a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py
+++ b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -44,7 +44,7 @@ class KDPM2DiscreteSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -149,10 +149,10 @@ def __init__(
beta_start: float = 0.00085, # sensible defaults
beta_end: float = 0.012,
beta_schedule: str = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
+ trained_betas: np.ndarray | list[float] | None = None,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
prediction_type: str = "epsilon",
timestep_spacing: str = "linspace",
steps_offset: int = 0,
@@ -222,7 +222,7 @@ def set_begin_index(self, begin_index: int = 0):
def scale_model_input(
self,
sample: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
@@ -252,8 +252,8 @@ def scale_model_input(
def set_timesteps(
self,
num_inference_steps: int,
- device: Union[str, torch.device] = None,
- num_train_timesteps: Optional[int] = None,
+ device: str | torch.device = None,
+ num_train_timesteps: int = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -339,7 +339,7 @@ def state_in_first_order(self):
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -369,7 +369,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -545,11 +545,11 @@ def _convert_to_beta(
def step(
self,
- model_output: Union[torch.Tensor, np.ndarray],
- timestep: Union[float, torch.Tensor],
- sample: Union[torch.Tensor, np.ndarray],
+ model_output: torch.Tensor | np.ndarray,
+ timestep: float | torch.Tensor,
+ sample: torch.Tensor | np.ndarray,
return_dict: bool = True,
- ) -> Union[KDPM2DiscreteSchedulerOutput, Tuple]:
+ ) -> KDPM2DiscreteSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_karras_ve_flax.py b/src/diffusers/schedulers/scheduling_karras_ve_flax.py
index 3f43a5fa9963..04c46220fcca 100644
--- a/src/diffusers/schedulers/scheduling_karras_ve_flax.py
+++ b/src/diffusers/schedulers/scheduling_karras_ve_flax.py
@@ -14,7 +14,6 @@
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import flax
import jax
@@ -32,9 +31,9 @@
@flax.struct.dataclass
class KarrasVeSchedulerState:
# setable values
- num_inference_steps: Optional[int] = None
- timesteps: Optional[jnp.ndarray] = None
- schedule: Optional[jnp.ndarray] = None # sigma(t_i)
+ num_inference_steps: int = None
+ timesteps: jnp.ndarray | None = None
+ schedule: jnp.ndarray | None = None # sigma(t_i)
@classmethod
def create(cls):
@@ -114,7 +113,7 @@ def create_state(self):
return KarrasVeSchedulerState.create()
def set_timesteps(
- self, state: KarrasVeSchedulerState, num_inference_steps: int, shape: Tuple = ()
+ self, state: KarrasVeSchedulerState, num_inference_steps: int, shape: tuple = ()
) -> KarrasVeSchedulerState:
"""
Sets the continuous timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -147,7 +146,7 @@ def add_noise_to_input(
sample: jnp.ndarray,
sigma: float,
key: jax.Array,
- ) -> Tuple[jnp.ndarray, float]:
+ ) -> tuple[jnp.ndarray, float]:
"""
Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a
higher noise level sigma_hat = sigma_i + gamma_i*sigma_i.
@@ -175,7 +174,7 @@ def step(
sigma_prev: float,
sample_hat: jnp.ndarray,
return_dict: bool = True,
- ) -> Union[FlaxKarrasVeOutput, Tuple]:
+ ) -> FlaxKarrasVeOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -213,7 +212,7 @@ def step_correct(
sample_prev: jnp.ndarray,
derivative: jnp.ndarray,
return_dict: bool = True,
- ) -> Union[FlaxKarrasVeOutput, Tuple]:
+ ) -> FlaxKarrasVeOutput | tuple:
"""
Correct the predicted sample based on the output model_output of the network. TODO complete description
diff --git a/src/diffusers/schedulers/scheduling_lcm.py b/src/diffusers/schedulers/scheduling_lcm.py
index ada8806e8c73..e70bc4745a24 100644
--- a/src/diffusers/schedulers/scheduling_lcm.py
+++ b/src/diffusers/schedulers/scheduling_lcm.py
@@ -17,7 +17,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -46,7 +46,7 @@ class LCMSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- denoised: Optional[torch.Tensor] = None
+ denoised: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -207,7 +207,7 @@ def __init__(
beta_start: float = 0.00085,
beta_end: float = 0.012,
beta_schedule: str = "scaled_linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
original_inference_steps: int = 50,
clip_sample: bool = False,
clip_sample_range: float = 1.0,
@@ -260,7 +260,7 @@ def __init__(
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -290,7 +290,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -327,7 +327,7 @@ def set_begin_index(self, begin_index: int = 0):
"""
self._begin_index = begin_index
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -389,10 +389,10 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device] = None,
- original_inference_steps: Optional[int] = None,
- timesteps: Optional[List[int]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device = None,
+ original_inference_steps: int | None = None,
+ timesteps: list[int] | None = None,
strength: int = 1.0,
):
"""
@@ -409,7 +409,7 @@ def set_timesteps(
schedule (which is different from the standard `diffusers` implementation). We will then take
`num_inference_steps` timesteps from this schedule, evenly spaced in terms of indices, and use that as
our final timestep schedule. If not set, this will default to the `original_inference_steps` attribute.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
timestep spacing strategy of equal spacing between timesteps on the training/distillation timestep
schedule is used. If `timesteps` is passed, `num_inference_steps` must be `None`.
@@ -540,9 +540,9 @@ def step(
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[LCMSchedulerOutput, Tuple]:
+ ) -> LCMSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py
index a1f9d27fd938..e917adda9516 100644
--- a/src/diffusers/schedulers/scheduling_lms_discrete.py
+++ b/src/diffusers/schedulers/scheduling_lms_discrete.py
@@ -14,18 +14,21 @@
import math
import warnings
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
-import scipy.stats
import torch
-from scipy import integrate
from ..configuration_utils import ConfigMixin, register_to_config
-from ..utils import BaseOutput
+from ..utils import BaseOutput, is_scipy_available
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+if is_scipy_available():
+ import scipy.stats
+ from scipy import integrate
+
+
@dataclass
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete
class LMSDiscreteSchedulerOutput(BaseOutput):
@@ -42,7 +45,7 @@ class LMSDiscreteSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -144,11 +147,11 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.0001,
beta_end: float = 0.02,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon",
timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
steps_offset: int = 0,
@@ -189,7 +192,7 @@ def __init__(
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
@property
- def init_noise_sigma(self) -> Union[float, torch.Tensor]:
+ def init_noise_sigma(self) -> float | torch.Tensor:
"""
The standard deviation of the initial noise distribution.
@@ -205,7 +208,7 @@ def init_noise_sigma(self) -> Union[float, torch.Tensor]:
return (self.sigmas.max() ** 2 + 1) ** 0.5
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase by 1 after each scheduler step.
@@ -216,7 +219,7 @@ def step_index(self) -> Optional[int]:
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
@@ -237,7 +240,7 @@ def set_begin_index(self, begin_index: int = 0) -> None:
"""
self._begin_index = begin_index
- def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -290,7 +293,7 @@ def lms_derivative(tau):
return integrated_coeff
- def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None) -> None:
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -350,7 +353,7 @@ def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, to
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -380,7 +383,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -542,11 +545,11 @@ def _convert_to_beta(
def step(
self,
model_output: torch.Tensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
sample: torch.Tensor,
order: int = 4,
return_dict: bool = True,
- ) -> Union[LMSDiscreteSchedulerOutput, Tuple]:
+ ) -> LMSDiscreteSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_lms_discrete_flax.py b/src/diffusers/schedulers/scheduling_lms_discrete_flax.py
index 4edb091348c8..65902678e1d9 100644
--- a/src/diffusers/schedulers/scheduling_lms_discrete_flax.py
+++ b/src/diffusers/schedulers/scheduling_lms_discrete_flax.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import flax
import jax.numpy as jnp
@@ -41,10 +40,10 @@ class LMSDiscreteSchedulerState:
init_noise_sigma: jnp.ndarray
timesteps: jnp.ndarray
sigmas: jnp.ndarray
- num_inference_steps: Optional[int] = None
+ num_inference_steps: int = None
# running values
- derivatives: Optional[jnp.ndarray] = None
+ derivatives: jnp.ndarray | None = None
@classmethod
def create(
@@ -110,7 +109,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: str = "linear",
- trained_betas: Optional[jnp.ndarray] = None,
+ trained_betas: jnp.ndarray | None = None,
prediction_type: str = "epsilon",
dtype: jnp.dtype = jnp.float32,
):
@@ -120,7 +119,7 @@ def __init__(
)
self.dtype = dtype
- def create_state(self, common: Optional[CommonSchedulerState] = None) -> LMSDiscreteSchedulerState:
+ def create_state(self, common: CommonSchedulerState | None = None) -> LMSDiscreteSchedulerState:
if common is None:
common = CommonSchedulerState.create(self)
@@ -185,7 +184,7 @@ def set_timesteps(
self,
state: LMSDiscreteSchedulerState,
num_inference_steps: int,
- shape: Tuple = (),
+ shape: tuple = (),
) -> LMSDiscreteSchedulerState:
"""
Sets the timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -233,7 +232,7 @@ def step(
sample: jnp.ndarray,
order: int = 4,
return_dict: bool = True,
- ) -> Union[FlaxLMSSchedulerOutput, Tuple]:
+ ) -> FlaxLMSSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_ltx_euler_ancestral_rf.py b/src/diffusers/schedulers/scheduling_ltx_euler_ancestral_rf.py
index 6710254f4445..453c8515c301 100644
--- a/src/diffusers/schedulers/scheduling_ltx_euler_ancestral_rf.py
+++ b/src/diffusers/schedulers/scheduling_ltx_euler_ancestral_rf.py
@@ -23,7 +23,6 @@
"""
from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
import torch
@@ -84,18 +83,18 @@ def __init__(
s_noise: float = 1.0,
):
# Note: num_train_timesteps is kept only for config compatibility.
- self.num_inference_steps: Optional[int] = None
- self.sigmas: Optional[torch.Tensor] = None
- self.timesteps: Optional[torch.Tensor] = None
- self._step_index: Optional[int] = None
- self._begin_index: Optional[int] = None
+ self.num_inference_steps: int = None
+ self.sigmas: torch.Tensor | None = None
+ self.timesteps: torch.Tensor | None = None
+ self._step_index: int = None
+ self._begin_index: int = None
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It can be set from a pipeline with `set_begin_index` to support
image-to-image like workflows that start denoising part-way through the schedule.
@@ -110,7 +109,7 @@ def set_begin_index(self, begin_index: int = 0):
self._begin_index = begin_index
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Map a (continuous) `timestep` value to an index into `self.timesteps`.
@@ -142,7 +141,7 @@ def index_for_timestep(
return indices[pos].item()
- def _init_step_index(self, timestep: Union[float, torch.Tensor]):
+ def _init_step_index(self, timestep: float | torch.Tensor):
"""
Initialize the internal step index based on a given timestep.
"""
@@ -158,11 +157,11 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]):
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device, None] = None,
- sigmas: Optional[Union[List[float], torch.Tensor]] = None,
- timesteps: Optional[Union[List[float], torch.Tensor]] = None,
- mu: Optional[float] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ sigmas: list[float] | torch.Tensor | None = None,
+ timesteps: list[float] | torch.Tensor | None = None,
+ mu: float | None = None,
**kwargs,
):
"""
@@ -179,9 +178,9 @@ def set_timesteps(
to be consistent and are otherwise ignored with a warning.
device (`str` or `torch.device`, *optional*):
Device to move the internal tensors to.
- sigmas (`List[float]` or `torch.Tensor`, *optional*):
+ sigmas (`list[float]` or `torch.Tensor`, *optional*):
Explicit sigma schedule, e.g. `[1.0, 0.99, ..., 0.0]`.
- timesteps (`List[float]` or `torch.Tensor`, *optional*):
+ timesteps (`list[float]` or `torch.Tensor`, *optional*):
Optional alias for `sigmas`. If `sigmas` is None and `timesteps` is provided, timesteps are treated as
sigmas.
mu (`float`, *optional*):
@@ -278,11 +277,11 @@ def _sigma_broadcast(self, sigma: torch.Tensor, sample: torch.Tensor) -> torch.T
def step(
self,
model_output: torch.FloatTensor,
- timestep: Union[float, torch.Tensor],
+ timestep: float | torch.Tensor,
sample: torch.FloatTensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[LTXEulerAncestralRFSchedulerOutput, Tuple[torch.FloatTensor]]:
+ ) -> LTXEulerAncestralRFSchedulerOutput | tuple[torch.FloatTensor]:
"""
Perform a single Euler-Ancestral RF update step.
diff --git a/src/diffusers/schedulers/scheduling_pndm.py b/src/diffusers/schedulers/scheduling_pndm.py
index 0820f5baa871..f87dff14375d 100644
--- a/src/diffusers/schedulers/scheduling_pndm.py
+++ b/src/diffusers/schedulers/scheduling_pndm.py
@@ -15,7 +15,7 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim
import math
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -122,8 +122,8 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.0001,
beta_end: float = 0.02,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
skip_prk_steps: bool = False,
set_alpha_to_one: bool = False,
prediction_type: Literal["epsilon", "v_prediction"] = "epsilon",
@@ -169,7 +169,7 @@ def __init__(
self.plms_timesteps = None
self.timesteps = None
- def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None) -> None:
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -235,7 +235,7 @@ def step(
timestep: int,
sample: torch.Tensor,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise), and calls [`~PNDMScheduler.step_prk`]
@@ -268,7 +268,7 @@ def step_prk(
timestep: int,
sample: torch.Tensor,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the Runge-Kutta method. It performs four forward passes to approximate the solution to the differential
@@ -327,7 +327,7 @@ def step_plms(
timestep: int,
sample: torch.Tensor,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the linear multistep method. It performs one forward pass multiple times to approximate the solution.
diff --git a/src/diffusers/schedulers/scheduling_pndm_flax.py b/src/diffusers/schedulers/scheduling_pndm_flax.py
index bbef4649ecb5..377e2d35ba0a 100644
--- a/src/diffusers/schedulers/scheduling_pndm_flax.py
+++ b/src/diffusers/schedulers/scheduling_pndm_flax.py
@@ -15,7 +15,6 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import flax
import jax
@@ -43,15 +42,15 @@ class PNDMSchedulerState:
# setable values
init_noise_sigma: jnp.ndarray
timesteps: jnp.ndarray
- num_inference_steps: Optional[int] = None
- prk_timesteps: Optional[jnp.ndarray] = None
- plms_timesteps: Optional[jnp.ndarray] = None
+ num_inference_steps: int = None
+ prk_timesteps: jnp.ndarray | None = None
+ plms_timesteps: jnp.ndarray | None = None
# running values
- cur_model_output: Optional[jnp.ndarray] = None
- counter: Optional[jnp.int32] = None
- cur_sample: Optional[jnp.ndarray] = None
- ets: Optional[jnp.ndarray] = None
+ cur_model_output: jnp.ndarray | None = None
+ counter: jnp.int32 | None = None
+ cur_sample: jnp.ndarray | None = None
+ ets: jnp.ndarray | None = None
@classmethod
def create(
@@ -128,7 +127,7 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: str = "linear",
- trained_betas: Optional[jnp.ndarray] = None,
+ trained_betas: jnp.ndarray | None = None,
skip_prk_steps: bool = False,
set_alpha_to_one: bool = False,
steps_offset: int = 0,
@@ -146,7 +145,7 @@ def __init__(
# mainly at formula (9), (12), (13) and the Algorithm 2.
self.pndm_order = 4
- def create_state(self, common: Optional[CommonSchedulerState] = None) -> PNDMSchedulerState:
+ def create_state(self, common: CommonSchedulerState | None = None) -> PNDMSchedulerState:
if common is None:
common = CommonSchedulerState.create(self)
@@ -170,7 +169,7 @@ def create_state(self, common: Optional[CommonSchedulerState] = None) -> PNDMSch
timesteps=timesteps,
)
- def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, shape: Tuple) -> PNDMSchedulerState:
+ def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, shape: tuple) -> PNDMSchedulerState:
"""
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -179,7 +178,7 @@ def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, sha
the `FlaxPNDMScheduler` state data class instance.
num_inference_steps (`int`):
the number of diffusion steps used when generating samples with a pre-trained model.
- shape (`Tuple`):
+ shape (`tuple`):
the shape of the samples to be generated.
"""
@@ -232,7 +231,7 @@ def scale_model_input(
self,
state: PNDMSchedulerState,
sample: jnp.ndarray,
- timestep: Optional[int] = None,
+ timestep: int | None = None,
) -> jnp.ndarray:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
@@ -255,7 +254,7 @@ def step(
timestep: int,
sample: jnp.ndarray,
return_dict: bool = True,
- ) -> Union[FlaxPNDMSchedulerOutput, Tuple]:
+ ) -> FlaxPNDMSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -309,7 +308,7 @@ def step_prk(
model_output: jnp.ndarray,
timestep: int,
sample: jnp.ndarray,
- ) -> Union[FlaxPNDMSchedulerOutput, Tuple]:
+ ) -> FlaxPNDMSchedulerOutput | tuple:
"""
Step function propagating the sample with the Runge-Kutta method. RK takes 4 forward passes to approximate the
solution to the differential equation.
@@ -379,7 +378,7 @@ def step_plms(
model_output: jnp.ndarray,
timestep: int,
sample: jnp.ndarray,
- ) -> Union[FlaxPNDMSchedulerOutput, Tuple]:
+ ) -> FlaxPNDMSchedulerOutput | tuple:
"""
Step function propagating the sample with the linear multi-step method. This has one forward pass with multiple
times to approximate the solution.
diff --git a/src/diffusers/schedulers/scheduling_repaint.py b/src/diffusers/schedulers/scheduling_repaint.py
index bec4a1bdf652..43665f6c68e8 100644
--- a/src/diffusers/schedulers/scheduling_repaint.py
+++ b/src/diffusers/schedulers/scheduling_repaint.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -133,7 +133,7 @@ def __init__(
beta_end: float = 0.02,
beta_schedule: str = "linear",
eta: float = 0.0,
- trained_betas: Optional[np.ndarray] = None,
+ trained_betas: np.ndarray | None = None,
clip_sample: bool = True,
):
if trained_betas is not None:
@@ -168,7 +168,7 @@ def __init__(
self.eta = eta
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -190,7 +190,7 @@ def set_timesteps(
num_inference_steps: int,
jump_length: int = 10,
jump_n_sample: int = 10,
- device: Union[str, torch.device] = None,
+ device: str | torch.device = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -258,9 +258,9 @@ def step(
sample: torch.Tensor,
original_image: torch.Tensor,
mask: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[RePaintSchedulerOutput, Tuple]:
+ ) -> RePaintSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_sasolver.py b/src/diffusers/schedulers/scheduling_sasolver.py
index df8fbd7b1a47..28369e9e3384 100644
--- a/src/diffusers/schedulers/scheduling_sasolver.py
+++ b/src/diffusers/schedulers/scheduling_sasolver.py
@@ -16,7 +16,7 @@
# The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
import math
-from typing import Callable, List, Literal, Optional, Tuple, Union
+from typing import Callable, Literal
import numpy as np
import torch
@@ -162,23 +162,23 @@ def __init__(
beta_start: float = 0.0001,
beta_end: float = 0.02,
beta_schedule: str = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
predictor_order: int = 2,
corrector_order: int = 2,
prediction_type: str = "epsilon",
- tau_func: Optional[Callable] = None,
+ tau_func: Callable | None = None,
thresholding: bool = False,
dynamic_thresholding_ratio: float = 0.995,
sample_max_value: float = 1.0,
algorithm_type: str = "data_prediction",
lower_order_final: bool = True,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- use_flow_sigmas: Optional[bool] = False,
- flow_shift: Optional[float] = 1.0,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ use_flow_sigmas: bool = False,
+ flow_shift: float = 1.0,
lambda_min_clipped: float = -float("inf"),
- variance_type: Optional[str] = None,
+ variance_type: str | None = None,
timestep_spacing: str = "linspace",
steps_offset: int = 0,
):
@@ -266,7 +266,7 @@ def set_begin_index(self, begin_index: int = 0):
"""
self._begin_index = begin_index
- def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None):
+ def set_timesteps(self, num_inference_steps: int = None, device: str | torch.device = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -438,7 +438,7 @@ def _sigma_to_alpha_sigma_t(self, sigma):
The sigma value(s) to convert.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing (alpha_t, sigma_t) values.
"""
if self.config.use_flow_sigmas:
@@ -1123,8 +1123,8 @@ def stochastic_adams_moulton_update(
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
def index_for_timestep(
self,
- timestep: Union[int, torch.Tensor],
- schedule_timesteps: Optional[torch.Tensor] = None,
+ timestep: int | torch.Tensor,
+ schedule_timesteps: torch.Tensor | None = None,
) -> int:
"""
Find the index for a given timestep in the schedule.
@@ -1181,7 +1181,7 @@ def step(
sample: torch.Tensor,
generator=None,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the SA-Solver.
diff --git a/src/diffusers/schedulers/scheduling_scm.py b/src/diffusers/schedulers/scheduling_scm.py
index 7b01d886299c..fb3bbc40d726 100644
--- a/src/diffusers/schedulers/scheduling_scm.py
+++ b/src/diffusers/schedulers/scheduling_scm.py
@@ -16,7 +16,6 @@
# and https://github.com/hojonathanho/diffusion
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -46,7 +45,7 @@ class SCMSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
class SCMScheduler(SchedulerMixin, ConfigMixin):
@@ -118,7 +117,7 @@ def set_timesteps(
self,
num_inference_steps: int,
timesteps: torch.Tensor = None,
- device: Union[str, torch.device] = None,
+ device: str | torch.device = None,
max_timesteps: float = 1.57080,
intermediate_timesteps: float = 1.3,
):
@@ -173,7 +172,7 @@ def set_timesteps(
self._begin_index = None
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -190,7 +189,7 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -226,7 +225,7 @@ def step(
sample: torch.FloatTensor,
generator: torch.Generator = None,
return_dict: bool = True,
- ) -> Union[SCMSchedulerOutput, Tuple]:
+ ) -> SCMSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_sde_ve.py b/src/diffusers/schedulers/scheduling_sde_ve.py
index 1bfc08cce5e9..e4e37666c60d 100644
--- a/src/diffusers/schedulers/scheduling_sde_ve.py
+++ b/src/diffusers/schedulers/scheduling_sde_ve.py
@@ -16,7 +16,6 @@
import math
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import torch
@@ -86,7 +85,7 @@ def __init__(
self.set_sigmas(num_train_timesteps, sigma_min, sigma_max, sampling_eps)
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -103,9 +102,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
"""
return sample
- def set_timesteps(
- self, num_inference_steps: int, sampling_eps: float = None, device: Union[str, torch.device] = None
- ):
+ def set_timesteps(self, num_inference_steps: int, sampling_eps: float = None, device: str | torch.device = None):
"""
Sets the continuous timesteps used for the diffusion chain (to be run before inference).
@@ -162,9 +159,9 @@ def step_pred(
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[SdeVeOutput, Tuple]:
+ ) -> SdeVeOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -229,9 +226,9 @@ def step_correct(
self,
model_output: torch.Tensor,
sample: torch.Tensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Correct the predicted sample based on the `model_output` of the network. This is often run repeatedly after
making the prediction for the previous timestep.
diff --git a/src/diffusers/schedulers/scheduling_sde_ve_flax.py b/src/diffusers/schedulers/scheduling_sde_ve_flax.py
index f4fe6d8f6bbf..62a54c7dc948 100644
--- a/src/diffusers/schedulers/scheduling_sde_ve_flax.py
+++ b/src/diffusers/schedulers/scheduling_sde_ve_flax.py
@@ -15,7 +15,6 @@
# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import flax
import jax
@@ -37,9 +36,9 @@
@flax.struct.dataclass
class ScoreSdeVeSchedulerState:
# setable values
- timesteps: Optional[jnp.ndarray] = None
- discrete_sigmas: Optional[jnp.ndarray] = None
- sigmas: Optional[jnp.ndarray] = None
+ timesteps: jnp.ndarray | None = None
+ discrete_sigmas: jnp.ndarray | None = None
+ sigmas: jnp.ndarray | None = None
@classmethod
def create(cls):
@@ -62,7 +61,7 @@ class FlaxSdeVeOutput(FlaxSchedulerOutput):
state: ScoreSdeVeSchedulerState
prev_sample: jnp.ndarray
- prev_sample_mean: Optional[jnp.ndarray] = None
+ prev_sample_mean: jnp.ndarray | None = None
class FlaxScoreSdeVeScheduler(FlaxSchedulerMixin, ConfigMixin):
@@ -122,7 +121,7 @@ def set_timesteps(
self,
state: ScoreSdeVeSchedulerState,
num_inference_steps: int,
- shape: Tuple = (),
+ shape: tuple = (),
sampling_eps: float = None,
) -> ScoreSdeVeSchedulerState:
"""
@@ -187,7 +186,7 @@ def step_pred(
sample: jnp.ndarray,
key: jax.Array,
return_dict: bool = True,
- ) -> Union[FlaxSdeVeOutput, Tuple]:
+ ) -> FlaxSdeVeOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
@@ -246,7 +245,7 @@ def step_correct(
sample: jnp.ndarray,
key: jax.Array,
return_dict: bool = True,
- ) -> Union[FlaxSdeVeOutput, Tuple]:
+ ) -> FlaxSdeVeOutput | tuple:
"""
Correct the predicted sample based on the output model_output of the network. This is often run repeatedly
after making the prediction for the previous timestep.
diff --git a/src/diffusers/schedulers/scheduling_tcd.py b/src/diffusers/schedulers/scheduling_tcd.py
index a1303436cd8d..838d040c04a6 100644
--- a/src/diffusers/schedulers/scheduling_tcd.py
+++ b/src/diffusers/schedulers/scheduling_tcd.py
@@ -17,7 +17,7 @@
import math
from dataclasses import dataclass
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -45,7 +45,7 @@ class TCDSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_noised_sample: Optional[torch.Tensor] = None
+ pred_noised_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -208,7 +208,7 @@ def __init__(
beta_start: float = 0.00085,
beta_end: float = 0.012,
beta_schedule: str = "scaled_linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ trained_betas: np.ndarray | list[float] | None = None,
original_inference_steps: int = 50,
clip_sample: bool = False,
clip_sample_range: float = 1.0,
@@ -261,7 +261,7 @@ def __init__(
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
def index_for_timestep(
- self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None
+ self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None
) -> int:
"""
Find the index of a given timestep in the timestep schedule.
@@ -291,7 +291,7 @@ def index_for_timestep(
return indices[pos].item()
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
- def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
"""
Initialize the step index for the scheduler based on the given timestep.
@@ -328,7 +328,7 @@ def set_begin_index(self, begin_index: int = 0):
"""
self._begin_index = begin_index
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -420,10 +420,10 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device] = None,
- original_inference_steps: Optional[int] = None,
- timesteps: Optional[List[int]] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device = None,
+ original_inference_steps: int | None = None,
+ timesteps: list[int] | None = None,
strength: float = 1.0,
):
"""
@@ -440,7 +440,7 @@ def set_timesteps(
schedule (which is different from the standard `diffusers` implementation). We will then take
`num_inference_steps` timesteps from this schedule, evenly spaced in terms of indices, and use that as
our final timestep schedule. If not set, this will default to the `original_inference_steps` attribute.
- timesteps (`List[int]`, *optional*):
+ timesteps (`list[int]`, *optional*):
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
timestep spacing strategy of equal spacing between timesteps on the training/distillation timestep
schedule is used. If `timesteps` is passed, `num_inference_steps` must be `None`.
@@ -586,9 +586,9 @@ def step(
timestep: int,
sample: torch.Tensor,
eta: float = 0.3,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[TCDSchedulerOutput, Tuple]:
+ ) -> TCDSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_unclip.py b/src/diffusers/schedulers/scheduling_unclip.py
index 14b09277da04..26226b4eb35b 100644
--- a/src/diffusers/schedulers/scheduling_unclip.py
+++ b/src/diffusers/schedulers/scheduling_unclip.py
@@ -14,7 +14,7 @@
import math
from dataclasses import dataclass
-from typing import Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -41,7 +41,7 @@ class UnCLIPSchedulerOutput(BaseOutput):
"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
@@ -132,7 +132,7 @@ def __init__(
num_train_timesteps: int = 1000,
variance_type: str = "fixed_small_log",
clip_sample: bool = True,
- clip_sample_range: Optional[float] = 1.0,
+ clip_sample_range: float = 1.0,
prediction_type: str = "epsilon",
beta_schedule: str = "squaredcos_cap_v2",
):
@@ -154,7 +154,7 @@ def __init__(
self.variance_type = variance_type
- def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
"""
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
current timestep.
@@ -168,7 +168,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None
"""
return sample
- def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None):
"""
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -226,10 +226,10 @@ def step(
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
- prev_timestep: Optional[int] = None,
+ prev_timestep: int | None = None,
generator=None,
return_dict: bool = True,
- ) -> Union[UnCLIPSchedulerOutput, Tuple]:
+ ) -> UnCLIPSchedulerOutput | tuple:
"""
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
process from the learned model outputs (most often the predicted noise).
diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py
index de04e854d60c..71a5444491ed 100644
--- a/src/diffusers/schedulers/scheduling_unipc_multistep.py
+++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py
@@ -16,7 +16,7 @@
# The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
import math
-from typing import List, Literal, Optional, Tuple, Union
+from typing import Literal
import numpy as np
import torch
@@ -201,8 +201,8 @@ def __init__(
num_train_timesteps: int = 1000,
beta_start: float = 0.0001,
beta_end: float = 0.02,
- beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear",
- trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+ beta_schedule: str = "linear",
+ trained_betas: np.ndarray | list[float] | None = None,
solver_order: int = 2,
prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon",
thresholding: bool = False,
@@ -211,22 +211,22 @@ def __init__(
predict_x0: bool = True,
solver_type: Literal["bh1", "bh2"] = "bh2",
lower_order_final: bool = True,
- disable_corrector: List[int] = [],
- solver_p: Optional[SchedulerMixin] = None,
- use_karras_sigmas: Optional[bool] = False,
- use_exponential_sigmas: Optional[bool] = False,
- use_beta_sigmas: Optional[bool] = False,
- use_flow_sigmas: Optional[bool] = False,
- flow_shift: Optional[float] = 1.0,
+ disable_corrector: list[int] = [],
+ solver_p: SchedulerMixin = None,
+ use_karras_sigmas: bool = False,
+ use_exponential_sigmas: bool = False,
+ use_beta_sigmas: bool = False,
+ use_flow_sigmas: bool = False,
+ flow_shift: float = 1.0,
timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace",
steps_offset: int = 0,
- final_sigmas_type: Optional[Literal["zero", "sigma_min"]] = "zero",
+ final_sigmas_type: Literal["zero", "sigma_min"] = "zero",
rescale_betas_zero_snr: bool = False,
use_dynamic_shifting: bool = False,
time_shift_type: Literal["exponential"] = "exponential",
- sigma_min: Optional[float] = None,
- sigma_max: Optional[float] = None,
- shift_terminal: Optional[float] = None,
+ sigma_min: float | None = None,
+ sigma_max: bool | None = None,
+ shift_terminal: bool | None = None,
) -> None:
if self.config.use_beta_sigmas and not is_scipy_available():
raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
@@ -291,14 +291,14 @@ def __init__(
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
@property
- def step_index(self) -> Optional[int]:
+ def step_index(self) -> int:
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
"""
return self._step_index
@property
- def begin_index(self) -> Optional[int]:
+ def begin_index(self) -> int:
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
"""
@@ -317,10 +317,10 @@ def set_begin_index(self, begin_index: int = 0) -> None:
def set_timesteps(
self,
- num_inference_steps: Optional[int] = None,
- device: Union[str, torch.device] = None,
- sigmas: Optional[List[float]] = None,
- mu: Optional[float] = None,
+ num_inference_steps: int | None = None,
+ device: str | torch.device = None,
+ sigmas: list[float] | None = None,
+ mu: bool | None = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -482,6 +482,21 @@ def set_timesteps(
# Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.time_shift
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+ """
+ Apply time shifting to the sigmas.
+
+ Args:
+ mu (`float`):
+ The mu parameter for the time shift.
+ sigma (`float`):
+ The sigma parameter for the time shift.
+ t (`torch.Tensor`):
+ The input timesteps.
+
+ Returns:
+ `torch.Tensor`:
+ The time-shifted timesteps.
+ """
if self.config.time_shift_type == "exponential":
return self._time_shift_exponential(mu, sigma, t)
elif self.config.time_shift_type == "linear":
@@ -599,7 +614,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:
return t
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t
- def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Convert sigma values to alpha_t and sigma_t values.
@@ -608,7 +623,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to
The sigma value(s) to convert.
Returns:
- `Tuple[torch.Tensor, torch.Tensor]`:
+ `tuple[torch.Tensor, torch.Tensor]`:
A tuple containing (alpha_t, sigma_t) values.
"""
if self.config.use_flow_sigmas:
@@ -1084,8 +1099,8 @@ def multistep_uni_c_bh_update(
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
def index_for_timestep(
self,
- timestep: Union[int, torch.Tensor],
- schedule_timesteps: Optional[torch.Tensor] = None,
+ timestep: int | torch.Tensor,
+ schedule_timesteps: torch.Tensor | None = None,
) -> int:
"""
Find the index for a given timestep in the schedule.
@@ -1119,7 +1134,7 @@ def index_for_timestep(
return step_index
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
- def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
+ def _init_step_index(self, timestep: int | torch.Tensor) -> None:
"""
Initialize the step_index counter for the scheduler.
@@ -1138,10 +1153,10 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None:
def step(
self,
model_output: torch.Tensor,
- timestep: Union[int, torch.Tensor],
+ timestep: int | torch.Tensor,
sample: torch.Tensor,
return_dict: bool = True,
- ) -> Union[SchedulerOutput, Tuple]:
+ ) -> SchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
the multistep UniPC.
diff --git a/src/diffusers/schedulers/scheduling_utils.py b/src/diffusers/schedulers/scheduling_utils.py
index a355c7bb1a51..0d48d791ac42 100644
--- a/src/diffusers/schedulers/scheduling_utils.py
+++ b/src/diffusers/schedulers/scheduling_utils.py
@@ -15,7 +15,6 @@
import os
from dataclasses import dataclass
from enum import Enum
-from typing import Optional, Union
import torch
from huggingface_hub.utils import validate_hf_hub_args
@@ -83,7 +82,7 @@ class SchedulerMixin(PushToHubMixin):
the scheduler's `__init__` function, and the attributes can be accessed by `scheduler.config.num_train_timesteps`.
Class attributes:
- - **_compatibles** (`List[str]`) -- A list of scheduler classes that are compatible with the parent scheduler
+ - **_compatibles** (`list[str]`) -- A list of scheduler classes that are compatible with the parent scheduler
class. Use [`~ConfigMixin.from_config`] to load a different compatible scheduler class (should be overridden
by parent class).
"""
@@ -96,8 +95,8 @@ class SchedulerMixin(PushToHubMixin):
@validate_hf_hub_args
def from_pretrained(
cls,
- pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
- subfolder: Optional[str] = None,
+ pretrained_model_name_or_path: str | os.PathLike | None = None,
+ subfolder: str | None = None,
return_unused_kwargs=False,
**kwargs,
) -> Self:
@@ -116,14 +115,14 @@ def from_pretrained(
The subfolder location of a model file within a larger model repository on the Hub or locally.
return_unused_kwargs (`bool`, *optional*, defaults to `False`):
Whether kwargs that are not consumed by the Python class should be returned or not.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
@@ -153,7 +152,7 @@ def from_pretrained(
)
return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs)
- def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
+ def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
"""
Save a scheduler configuration object to a directory so that it can be reloaded using the
[`~SchedulerMixin.from_pretrained`] class method.
@@ -165,7 +164,7 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub:
Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
namespace).
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
"""
self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
@@ -176,7 +175,7 @@ def compatibles(self):
Returns all schedulers that are compatible with this scheduler
Returns:
- `List[SchedulerMixin]`: List of compatible schedulers
+ `list[SchedulerMixin]`: list of compatible schedulers
"""
return self._get_compatibles()
diff --git a/src/diffusers/schedulers/scheduling_utils_flax.py b/src/diffusers/schedulers/scheduling_utils_flax.py
index 0534e47d8a30..44de56a3980d 100644
--- a/src/diffusers/schedulers/scheduling_utils_flax.py
+++ b/src/diffusers/schedulers/scheduling_utils_flax.py
@@ -16,7 +16,6 @@
import os
from dataclasses import dataclass
from enum import Enum
-from typing import Optional, Tuple, Union
import flax
import jax.numpy as jnp
@@ -62,7 +61,7 @@ class FlaxSchedulerMixin(PushToHubMixin):
Mixin containing common functions for the schedulers.
Class attributes:
- - **_compatibles** (`List[str]`) -- A list of classes that are compatible with the parent class, so that
+ - **_compatibles** (`list[str]`) -- A list of classes that are compatible with the parent class, so that
`from_config` can be used from a class different than the one used to save the config (should be overridden
by parent class).
"""
@@ -76,8 +75,8 @@ class FlaxSchedulerMixin(PushToHubMixin):
@validate_hf_hub_args
def from_pretrained(
cls,
- pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
- subfolder: Optional[str] = None,
+ pretrained_model_name_or_path: str | os.PathLike | None = None,
+ subfolder: str | None = None,
return_unused_kwargs=False,
**kwargs,
):
@@ -98,14 +97,14 @@ def from_pretrained(
return_unused_kwargs (`bool`, *optional*, defaults to `False`):
Whether kwargs that are not consumed by the Python class should be returned or not.
- cache_dir (`Union[str, os.PathLike]`, *optional*):
+ cache_dir (`str | os.PathLike`, *optional*):
Path to a directory in which a downloaded pretrained model configuration should be cached if the
standard cache should not be used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
@@ -148,7 +147,7 @@ def from_pretrained(
return scheduler, state
- def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
+ def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
"""
Save a scheduler configuration object to the directory `save_directory`, so that it can be re-loaded using the
[`~FlaxSchedulerMixin.from_pretrained`] class method.
@@ -160,7 +159,7 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub:
Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
namespace).
- kwargs (`Dict[str, Any]`, *optional*):
+ kwargs (`dict[str, Any]`, *optional*):
Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
"""
self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
@@ -171,7 +170,7 @@ def compatibles(self):
Returns all schedulers that are compatible with this scheduler
Returns:
- `List[SchedulerMixin]`: List of compatible schedulers
+ `list[SchedulerMixin]`: list of compatible schedulers
"""
return self._get_compatibles()
@@ -185,7 +184,7 @@ def _get_compatibles(cls):
return compatible_classes
-def broadcast_to_shape_from_left(x: jnp.ndarray, shape: Tuple[int]) -> jnp.ndarray:
+def broadcast_to_shape_from_left(x: jnp.ndarray, shape: tuple[int]) -> jnp.ndarray:
assert len(shape) >= x.ndim
return jnp.broadcast_to(x.reshape(x.shape + (1,) * (len(shape) - x.ndim)), shape)
diff --git a/src/diffusers/schedulers/scheduling_vq_diffusion.py b/src/diffusers/schedulers/scheduling_vq_diffusion.py
index 57306301d023..7ba3a8e6f5c4 100644
--- a/src/diffusers/schedulers/scheduling_vq_diffusion.py
+++ b/src/diffusers/schedulers/scheduling_vq_diffusion.py
@@ -13,7 +13,6 @@
# limitations under the License.
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
import numpy as np
import torch
@@ -59,7 +58,7 @@ def index_to_log_onehot(x: torch.LongTensor, num_classes: int) -> torch.Tensor:
return log_x
-def gumbel_noised(logits: torch.Tensor, generator: Optional[torch.Generator]) -> torch.Tensor:
+def gumbel_noised(logits: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor:
"""
Apply gumbel noise to `logits`
"""
@@ -175,7 +174,7 @@ def __init__(
self.num_inference_steps = None
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy())
- def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -202,9 +201,9 @@ def step(
model_output: torch.Tensor,
timestep: torch.long,
sample: torch.LongTensor,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
return_dict: bool = True,
- ) -> Union[VQDiffusionSchedulerOutput, Tuple]:
+ ) -> VQDiffusionSchedulerOutput | tuple:
"""
Predict the sample from the previous timestep by the reverse transition distribution. See
[`~VQDiffusionScheduler.q_posterior`] for more details about how the distribution is computer.
diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py
index 3e9968d47fdd..6c07a30c2ccc 100644
--- a/src/diffusers/training_utils.py
+++ b/src/diffusers/training_utils.py
@@ -7,7 +7,7 @@
import warnings
from contextlib import contextmanager
from functools import partial
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Type, Union
+from typing import Any, Iterable
import numpy as np
import torch
@@ -162,7 +162,7 @@ def compute_dream_and_update_latents(
target: torch.Tensor,
encoder_hidden_states: torch.Tensor,
dream_detail_preservation: float = 1.0,
-) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+) -> tuple[torch.Tensor, torch.Tensor]:
"""
Implements "DREAM (Diffusion Rectification and Estimation-Adaptive Models)" from
https://huggingface.co/papers/2312.00210. DREAM helps align training with sampling to help training be more
@@ -207,7 +207,7 @@ def compute_dream_and_update_latents(
return _noisy_latents, _target
-def unet_lora_state_dict(unet: UNet2DConditionModel) -> Dict[str, torch.Tensor]:
+def unet_lora_state_dict(unet: UNet2DConditionModel) -> dict[str, torch.Tensor]:
r"""
Returns:
A state dict containing just the LoRA parameters.
@@ -226,7 +226,7 @@ def unet_lora_state_dict(unet: UNet2DConditionModel) -> Dict[str, torch.Tensor]:
return lora_state_dict
-def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32):
+def cast_training_params(model: torch.nn.Module | list[torch.nn.Module], dtype=torch.float32):
"""
Casts the training parameters of the model to the specified data type.
@@ -244,7 +244,7 @@ def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], d
def _set_state_dict_into_text_encoder(
- lora_state_dict: Dict[str, torch.Tensor], prefix: str, text_encoder: torch.nn.Module
+ lora_state_dict: dict[str, torch.Tensor], prefix: str, text_encoder: torch.nn.Module
):
"""
Sets the `lora_state_dict` into `text_encoder` coming from `transformers`.
@@ -262,7 +262,7 @@ def _set_state_dict_into_text_encoder(
set_peft_model_state_dict(text_encoder, text_encoder_state_dict, adapter_name="default")
-def _collate_lora_metadata(modules_to_save: Dict[str, torch.nn.Module]) -> Dict[str, Any]:
+def _collate_lora_metadata(modules_to_save: dict[str, torch.nn.Module]) -> dict[str, Any]:
metadatas = {}
for module_name, module in modules_to_save.items():
if module is not None:
@@ -276,8 +276,8 @@ def compute_density_for_timestep_sampling(
logit_mean: float = None,
logit_std: float = None,
mode_scale: float = None,
- device: Union[torch.device, str] = "cpu",
- generator: Optional[torch.Generator] = None,
+ device: torch.device | str = "cpu",
+ generator: torch.Generator | None = None,
):
"""
Compute the density for sampling the timesteps when doing SD3 training.
@@ -332,9 +332,7 @@ def free_memory():
@contextmanager
-def offload_models(
- *modules: Union[torch.nn.Module, DiffusionPipeline], device: Union[str, torch.device], offload: bool = True
-):
+def offload_models(*modules: torch.nn.Module | DiffusionPipeline, device: str | torch.device, offload: bool = True):
"""
Context manager that, if offload=True, moves each module to `device` on enter, then moves it back to its original
device on exit.
@@ -434,12 +432,12 @@ def get_fsdp_kwargs_from_accelerator(accelerator) -> dict:
def wrap_with_fsdp(
model: torch.nn.Module,
- device: Union[str, torch.device],
+ device: str | torch.device,
offload: bool = True,
use_orig_params: bool = True,
limit_all_gathers: bool = True,
- fsdp_kwargs: Optional[Dict[str, Any]] = None,
- transformer_layer_cls: Optional[Set[Type[torch.nn.Module]]] = None,
+ fsdp_kwargs: dict[str, Any] | None = None,
+ transformer_layer_cls: set[type[torch.nn.Module]] | None = None,
) -> FSDP:
"""
Wrap a model with FSDP using common defaults and optional transformer auto-wrapping.
@@ -465,10 +463,7 @@ def wrap_with_fsdp(
logger.info(f"transformer_layer_cls is not provided, auto-inferred as {transformer_layer_cls.__name__}")
# Add auto-wrap policy if transformer layers specified
- auto_wrap_policy = partial(
- transformer_auto_wrap_policy,
- transformer_layer_cls={transformer_layer_cls},
- )
+ auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={transformer_layer_cls})
config = {
"device_id": device,
@@ -498,11 +493,11 @@ def __init__(
min_decay: float = 0.0,
update_after_step: int = 0,
use_ema_warmup: bool = False,
- inv_gamma: Union[float, int] = 1.0,
- power: Union[float, int] = 2 / 3,
+ inv_gamma: float | int = 1.0,
+ power: float | int = 2 / 3,
foreach: bool = False,
- model_cls: Optional[Any] = None,
- model_config: Dict[str, Any] = None,
+ model_cls: Any | None = None,
+ model_config: dict[str, Any] | None = None,
**kwargs,
):
"""
@@ -516,7 +511,7 @@ def __init__(
Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True.
power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True.
foreach (bool): Use torch._foreach functions for updating shadow parameters. Should be faster.
- device (Optional[Union[str, torch.device]]): The device to store the EMA weights on. If None, the EMA
+ device (str | torch.device | None): The device to store the EMA weights on. If None, the EMA
weights will be stored on CPU.
@crowsonkb's notes on EMA Warmup:
diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py
index 3f736e2ee39b..dd50405c74b2 100644
--- a/src/diffusers/utils/__init__.py
+++ b/src/diffusers/utils/__init__.py
@@ -85,8 +85,6 @@
is_hpu_available,
is_inflect_available,
is_invisible_watermark_available,
- is_k_diffusion_available,
- is_k_diffusion_version,
is_kernels_available,
is_kornia_available,
is_librosa_available,
@@ -131,6 +129,7 @@
from .logging import get_logger
from .outputs import BaseOutput
from .peft_utils import (
+ apply_lora_scale,
check_peft_version,
delete_adapter_layers,
get_adapter_name,
diff --git a/src/diffusers/utils/deprecation_utils.py b/src/diffusers/utils/deprecation_utils.py
index d76623541b9f..7ec612499e9a 100644
--- a/src/diffusers/utils/deprecation_utils.py
+++ b/src/diffusers/utils/deprecation_utils.py
@@ -1,6 +1,6 @@
import inspect
import warnings
-from typing import Any, Dict, Optional, Union
+from typing import Any
from packaging import version
@@ -26,7 +26,7 @@
}
-def _maybe_remap_transformers_class(class_name: str) -> Optional[str]:
+def _maybe_remap_transformers_class(class_name: str) -> str | None:
"""
Check if a Transformers class should be remapped to a newer version.
@@ -53,7 +53,7 @@ def _maybe_remap_transformers_class(class_name: str) -> Optional[str]:
return None
-def deprecate(*args, take_from: Optional[Union[Dict, Any]] = None, standard_warn=True, stacklevel=2):
+def deprecate(*args, take_from: dict | Any | None = None, standard_warn=True, stacklevel=2):
from .. import __version__
deprecated_kwargs = take_from
diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py
index 6c436161c5a7..4e402921aa5f 100644
--- a/src/diffusers/utils/dummy_pt_objects.py
+++ b/src/diffusers/utils/dummy_pt_objects.py
@@ -896,6 +896,21 @@ def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch"])
+class CosmosControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
class CosmosTransformer3DModel(metaclass=DummyObject):
_backends = ["torch"]
@@ -1890,6 +1905,21 @@ def attention_backend(*args, **kwargs):
requires_backends(attention_backend, ["torch"])
+class AutoPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
class ComponentsManager(metaclass=DummyObject):
_backends = ["torch"]
@@ -1920,6 +1950,66 @@ def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch"])
+class ConditionalPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ConfigSpec(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class InputParam(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LoopSequentialPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
class ModularPipeline(metaclass=DummyObject):
_backends = ["torch"]
@@ -1950,6 +2040,36 @@ def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch"])
+class OutputParam(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SequentialPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
def get_constant_schedule(*args, **kwargs):
requires_backends(get_constant_schedule, ["torch"])
diff --git a/src/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py
deleted file mode 100644
index 2ab00c54ce2d..000000000000
--- a/src/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py
+++ /dev/null
@@ -1,32 +0,0 @@
-# This file is autogenerated by the command `make fix-copies`, do not edit.
-from ..utils import DummyObject, requires_backends
-
-
-class StableDiffusionKDiffusionPipeline(metaclass=DummyObject):
- _backends = ["torch", "transformers", "k_diffusion"]
-
- def __init__(self, *args, **kwargs):
- requires_backends(self, ["torch", "transformers", "k_diffusion"])
-
- @classmethod
- def from_config(cls, *args, **kwargs):
- requires_backends(cls, ["torch", "transformers", "k_diffusion"])
-
- @classmethod
- def from_pretrained(cls, *args, **kwargs):
- requires_backends(cls, ["torch", "transformers", "k_diffusion"])
-
-
-class StableDiffusionXLKDiffusionPipeline(metaclass=DummyObject):
- _backends = ["torch", "transformers", "k_diffusion"]
-
- def __init__(self, *args, **kwargs):
- requires_backends(self, ["torch", "transformers", "k_diffusion"])
-
- @classmethod
- def from_config(cls, *args, **kwargs):
- requires_backends(cls, ["torch", "transformers", "k_diffusion"])
-
- @classmethod
- def from_pretrained(cls, *args, **kwargs):
- requires_backends(cls, ["torch", "transformers", "k_diffusion"])
diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py
index 17c40613f9da..8758c549ca77 100644
--- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py
+++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py
@@ -977,6 +977,21 @@ def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
+class Cosmos2_5_TransferPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
class Cosmos2TextToImagePipeline(metaclass=DummyObject):
_backends = ["torch", "transformers"]
diff --git a/src/diffusers/utils/dynamic_modules_utils.py b/src/diffusers/utils/dynamic_modules_utils.py
index 1c0734cf35bb..2a1cea10e14d 100644
--- a/src/diffusers/utils/dynamic_modules_utils.py
+++ b/src/diffusers/utils/dynamic_modules_utils.py
@@ -24,7 +24,6 @@
import threading
from pathlib import Path
from types import ModuleType
-from typing import Dict, Optional, Union
from urllib import request
from huggingface_hub import hf_hub_download, model_info
@@ -65,7 +64,7 @@ def init_hf_modules():
init_path.touch()
-def create_dynamic_module(name: Union[str, os.PathLike]):
+def create_dynamic_module(name: str | os.PathLike):
"""
Creates a dynamic module in the cache directory for modules.
"""
@@ -197,7 +196,7 @@ def get_class_in_module(class_name, module_path, force_reload=False):
if force_reload:
sys.modules.pop(name, None)
importlib.invalidate_caches()
- cached_module: Optional[ModuleType] = sys.modules.get(name)
+ cached_module: ModuleType | None = sys.modules.get(name)
module_spec = importlib.util.spec_from_file_location(name, location=module_file)
module: ModuleType
@@ -245,16 +244,16 @@ def find_pipeline_class(loaded_module):
@validate_hf_hub_args
def get_cached_module_file(
- pretrained_model_name_or_path: Union[str, os.PathLike],
+ pretrained_model_name_or_path: str | os.PathLike,
module_file: str,
- subfolder: Optional[str] = None,
- cache_dir: Optional[Union[str, os.PathLike]] = None,
+ subfolder: str | None = None,
+ cache_dir: str | os.PathLike | None = None,
force_download: bool = False,
- proxies: Optional[Dict[str, str]] = None,
- token: Optional[Union[bool, str]] = None,
- revision: Optional[str] = None,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
local_files_only: bool = False,
- local_dir: Optional[str] = None,
+ local_dir: str | None = None,
):
"""
Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached
@@ -278,7 +277,7 @@ def get_cached_module_file(
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force to (re-)download the configuration files and override the cached versions if they
exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
token (`str` or *bool*, *optional*):
@@ -426,17 +425,17 @@ def get_cached_module_file(
@validate_hf_hub_args
def get_class_from_dynamic_module(
- pretrained_model_name_or_path: Union[str, os.PathLike],
+ pretrained_model_name_or_path: str | os.PathLike,
module_file: str,
- subfolder: Optional[str] = None,
- class_name: Optional[str] = None,
- cache_dir: Optional[Union[str, os.PathLike]] = None,
+ subfolder: str | None = None,
+ class_name: str | None = None,
+ cache_dir: str | os.PathLike | None = None,
force_download: bool = False,
- proxies: Optional[Dict[str, str]] = None,
- token: Optional[Union[bool, str]] = None,
- revision: Optional[str] = None,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
local_files_only: bool = False,
- local_dir: Optional[str] = None,
+ local_dir: str | None = None,
):
"""
Extracts a class from a module file, present in the local folder or repository of a model.
@@ -464,7 +463,7 @@ def get_class_from_dynamic_module(
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force to (re-)download the configuration files and override the cached versions if they
exist.
- proxies (`Dict[str, str]`, *optional*):
+ proxies (`dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
token (`str` or `bool`, *optional*):
diff --git a/src/diffusers/utils/export_utils.py b/src/diffusers/utils/export_utils.py
index 07cf46928a44..d9b4e173a862 100644
--- a/src/diffusers/utils/export_utils.py
+++ b/src/diffusers/utils/export_utils.py
@@ -3,7 +3,6 @@
import struct
import tempfile
from contextlib import contextmanager
-from typing import List, Optional, Union
import numpy as np
import PIL.Image
@@ -25,7 +24,7 @@ def buffered_writer(raw_f):
f.flush()
-def export_to_gif(image: List[PIL.Image.Image], output_gif_path: str = None, fps: int = 10) -> str:
+def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None, fps: int = 10) -> str:
if output_gif_path is None:
output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name
@@ -113,7 +112,7 @@ def export_to_obj(mesh, output_obj_path: str = None):
def _legacy_export_to_video(
- video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 10
+ video_frames: list[np.ndarray] | list[PIL.Image.Image], output_video_path: str = None, fps: int = 10
):
if is_opencv_available():
import cv2
@@ -139,12 +138,12 @@ def _legacy_export_to_video(
def export_to_video(
- video_frames: Union[List[np.ndarray], List[PIL.Image.Image]],
+ video_frames: list[np.ndarray] | list[PIL.Image.Image],
output_video_path: str = None,
fps: int = 10,
quality: float = 5.0,
- bitrate: Optional[int] = None,
- macro_block_size: Optional[int] = 16,
+ bitrate: int | None = None,
+ macro_block_size: int | None = 16,
) -> str:
"""
quality:
diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py
index 58695bae1e9d..ad1ce988870c 100644
--- a/src/diffusers/utils/hub_utils.py
+++ b/src/diffusers/utils/hub_utils.py
@@ -21,7 +21,6 @@
import tempfile
import warnings
from pathlib import Path
-from typing import Dict, List, Optional, Union
from uuid import uuid4
from huggingface_hub import (
@@ -72,7 +71,7 @@
SESSION_ID = uuid4().hex
-def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str:
+def http_user_agent(user_agent: dict | str | None = None) -> str:
"""
Formats a user-agent string with basic info about a request.
"""
@@ -98,15 +97,15 @@ def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str:
def load_or_create_model_card(
repo_id_or_path: str = None,
- token: Optional[str] = None,
+ token: str | None = None,
is_pipeline: bool = False,
from_training: bool = False,
- model_description: Optional[str] = None,
+ model_description: str | None = None,
base_model: str = None,
- prompt: Optional[str] = None,
- license: Optional[str] = None,
- widget: Optional[List[dict]] = None,
- inference: Optional[bool] = None,
+ prompt: str | None = None,
+ license: str | None = None,
+ widget: list[dict] | None = None,
+ inference: bool | None = None,
is_modular: bool = False,
) -> ModelCard:
"""
@@ -129,7 +128,7 @@ def load_or_create_model_card(
prompt (`str`, *optional*): Prompt used for training. Useful for DreamBooth-like training.
license: (`str`, *optional*): License of the output artifact. Helpful when using
`load_or_create_model_card` from a training script.
- widget (`List[dict]`, *optional*): Widget to accompany a gallery template.
+ widget (`list[dict]`, *optional*): Widget to accompany a gallery template.
inference: (`bool`, optional): Whether to turn on inference widget. Helpful when using
`load_or_create_model_card` from a training script.
is_modular: (`bool`, optional): Boolean flag to denote if the model card is for a modular pipeline.
@@ -174,7 +173,7 @@ def load_or_create_model_card(
return model_card
-def populate_model_card(model_card: ModelCard, tags: Union[str, List[str]] = None) -> ModelCard:
+def populate_model_card(model_card: ModelCard, tags: str | list[str] | None = None) -> ModelCard:
"""Populates the `model_card` with library name and optional tags."""
if model_card.data.library_name is None:
model_card.data.library_name = "diffusers"
@@ -190,7 +189,7 @@ def populate_model_card(model_card: ModelCard, tags: Union[str, List[str]] = Non
return model_card
-def extract_commit_hash(resolved_file: Optional[str], commit_hash: Optional[str] = None):
+def extract_commit_hash(resolved_file: str | None, commit_hash: str | None = None):
"""
Extracts the commit hash from a resolved filename toward a cache file.
"""
@@ -204,7 +203,7 @@ def extract_commit_hash(resolved_file: Optional[str], commit_hash: Optional[str]
return commit_hash if REGEX_COMMIT_HASH.match(commit_hash) else None
-def _add_variant(weights_name: str, variant: Optional[str] = None) -> str:
+def _add_variant(weights_name: str, variant: str | None = None) -> str:
if variant is not None:
splits = weights_name.split(".")
splits = splits[:-1] + [variant] + splits[-1:]
@@ -215,19 +214,19 @@ def _add_variant(weights_name: str, variant: Optional[str] = None) -> str:
@validate_hf_hub_args
def _get_model_file(
- pretrained_model_name_or_path: Union[str, Path],
+ pretrained_model_name_or_path: str | Path,
*,
weights_name: str,
- subfolder: Optional[str] = None,
- cache_dir: Optional[str] = None,
+ subfolder: str | None = None,
+ cache_dir: str | None = None,
force_download: bool = False,
- proxies: Optional[Dict] = None,
+ proxies: dict | None = None,
local_files_only: bool = False,
- token: Optional[str] = None,
- user_agent: Optional[Union[Dict, str]] = None,
- revision: Optional[str] = None,
- commit_hash: Optional[str] = None,
- dduf_entries: Optional[Dict[str, DDUFEntry]] = None,
+ token: str | None = None,
+ user_agent: dict | str | None = None,
+ revision: str | None = None,
+ commit_hash: str | None = None,
+ dduf_entries: dict[str, DDUFEntry] | None = None,
):
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
@@ -355,7 +354,7 @@ def _get_checkpoint_shard_files(
user_agent=None,
revision=None,
subfolder="",
- dduf_entries: Optional[Dict[str, DDUFEntry]] = None,
+ dduf_entries: dict[str, DDUFEntry] | None = None,
):
"""
For a given model:
@@ -456,7 +455,7 @@ def _get_checkpoint_shard_files(
return cached_filenames, sharded_metadata
-def _check_legacy_sharding_variant_format(folder: str = None, filenames: List[str] = None, variant: str = None):
+def _check_legacy_sharding_variant_format(folder: str = None, filenames: list[str] = None, variant: str = None):
if filenames and folder:
raise ValueError("Both `filenames` and `folder` cannot be provided.")
if not filenames:
@@ -476,12 +475,12 @@ class PushToHubMixin:
def _upload_folder(
self,
- working_dir: Union[str, os.PathLike],
+ working_dir: str | os.PathLike,
repo_id: str,
- token: Optional[str] = None,
- commit_message: Optional[str] = None,
+ token: str | None = None,
+ commit_message: str | None = None,
create_pr: bool = False,
- subfolder: Optional[str] = None,
+ subfolder: str | None = None,
):
"""
Uploads all files in `working_dir` to `repo_id`.
@@ -507,13 +506,13 @@ def _upload_folder(
def push_to_hub(
self,
repo_id: str,
- commit_message: Optional[str] = None,
- private: Optional[bool] = None,
- token: Optional[str] = None,
+ commit_message: str | None = None,
+ private: bool | None = None,
+ token: str | None = None,
create_pr: bool = False,
safe_serialization: bool = True,
- variant: Optional[str] = None,
- subfolder: Optional[str] = None,
+ variant: str | None = None,
+ subfolder: str | None = None,
) -> str:
"""
Upload model, scheduler, or pipeline files to the 🤗 Hugging Face Hub.
diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py
index e35fc9697076..8fb481946ebf 100644
--- a/src/diffusers/utils/import_utils.py
+++ b/src/diffusers/utils/import_utils.py
@@ -24,7 +24,7 @@
from functools import lru_cache as cache
from itertools import chain
from types import ModuleType
-from typing import Any, Tuple, Union
+from typing import Any
from huggingface_hub.utils import is_jinja_available # noqa: F401
from packaging.version import Version, parse
@@ -59,7 +59,7 @@
_is_google_colab = "google.colab" in sys.modules or any(k.startswith("COLAB_") for k in os.environ)
-def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[bool, str]:
+def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[bool, str]:
global _package_map
pkg_exists = importlib.util.find_spec(pkg_name) is not None
pkg_version = "N/A"
@@ -198,7 +198,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b
_kernels_available, _kernels_version = _is_package_available("kernels")
_inflect_available, _inflect_version = _is_package_available("inflect")
_unidecode_available, _unidecode_version = _is_package_available("unidecode")
-_k_diffusion_available, _k_diffusion_version = _is_package_available("k_diffusion")
+
_note_seq_available, _note_seq_version = _is_package_available("note_seq")
_wandb_available, _wandb_version = _is_package_available("wandb")
_tensorboard_available, _tensorboard_version = _is_package_available("tensorboard")
@@ -293,10 +293,6 @@ def is_kernels_available():
return _kernels_available
-def is_k_diffusion_available():
- return _k_diffusion_available
-
-
def is_note_seq_available():
return _note_seq_available
@@ -479,12 +475,6 @@ def is_av_available():
Unidecode`
"""
-# docstyle-ignore
-K_DIFFUSION_IMPORT_ERROR = """
-{0} requires the k-diffusion library but it was not found in your environment. You can install it with pip: `pip
-install k-diffusion`
-"""
-
# docstyle-ignore
NOTE_SEQ_IMPORT_ERROR = """
{0} requires the note-seq library but it was not found in your environment. You can install it with pip: `pip
@@ -601,7 +591,6 @@ def is_av_available():
("transformers", (is_transformers_available, TRANSFORMERS_IMPORT_ERROR)),
("unidecode", (is_unidecode_available, UNIDECODE_IMPORT_ERROR)),
("librosa", (is_librosa_available, LIBROSA_IMPORT_ERROR)),
- ("k_diffusion", (is_k_diffusion_available, K_DIFFUSION_IMPORT_ERROR)),
("note_seq", (is_note_seq_available, NOTE_SEQ_IMPORT_ERROR)),
("wandb", (is_wandb_available, WANDB_IMPORT_ERROR)),
("tensorboard", (is_tensorboard_available, TENSORBOARD_IMPORT_ERROR)),
@@ -668,7 +657,7 @@ def __getattr__(cls, key):
# This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L319
-def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+def compare_versions(library_or_version: str | Version, operation: str, requirement_version: str):
"""
Compares a library version to some requirement using a given operation.
@@ -830,22 +819,6 @@ def is_torchao_version(operation: str, version: str):
return compare_versions(parse(_torchao_version), operation, version)
-@cache
-def is_k_diffusion_version(operation: str, version: str):
- """
- Compares the current k-diffusion version to a given reference with an operation.
-
- Args:
- operation (`str`):
- A string representation of an operator, such as `">"` or `"<="`
- version (`str`):
- A version string
- """
- if not _k_diffusion_available:
- return False
- return compare_versions(parse(_k_diffusion_version), operation, version)
-
-
@cache
def is_optimum_quanto_version(operation: str, version: str):
"""
diff --git a/src/diffusers/utils/loading_utils.py b/src/diffusers/utils/loading_utils.py
index dd23ae73c861..c4fee0cfdd83 100644
--- a/src/diffusers/utils/loading_utils.py
+++ b/src/diffusers/utils/loading_utils.py
@@ -1,6 +1,6 @@
import os
import tempfile
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable
from urllib.parse import unquote, urlparse
import PIL.Image
@@ -12,7 +12,7 @@
def load_image(
- image: Union[str, PIL.Image.Image], convert_method: Optional[Callable[[PIL.Image.Image], PIL.Image.Image]] = None
+ image: str | PIL.Image.Image, convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None
) -> PIL.Image.Image:
"""
Loads `image` to a PIL Image.
@@ -56,20 +56,20 @@ def load_image(
def load_video(
video: str,
- convert_method: Optional[Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]]] = None,
-) -> List[PIL.Image.Image]:
+ convert_method: Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] | None = None,
+) -> list[PIL.Image.Image]:
"""
Loads `video` to a list of PIL Image.
Args:
video (`str`):
A URL or Path to a video to convert to a list of PIL Image format.
- convert_method (Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]], *optional*):
+ convert_method (Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]], *optional*):
A conversion method to apply to the video after loading it. When set to `None` the images will be converted
to "RGB".
Returns:
- `List[PIL.Image.Image]`:
+ `list[PIL.Image.Image]`:
The video as a list of PIL images.
"""
is_url = video.startswith("http://") or video.startswith("https://")
@@ -139,7 +139,7 @@ def load_video(
# Taken from `transformers`.
-def get_module_from_name(module, tensor_name: str) -> Tuple[Any, str]:
+def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]:
if "." in tensor_name:
splits = tensor_name.split(".")
for split in splits[:-1]:
diff --git a/src/diffusers/utils/logging.py b/src/diffusers/utils/logging.py
index 80e108e4a6ff..21c4bba8a830 100644
--- a/src/diffusers/utils/logging.py
+++ b/src/diffusers/utils/logging.py
@@ -28,7 +28,6 @@
WARN, # NOQA
WARNING, # NOQA
)
-from typing import Dict, Optional
from tqdm import auto as tqdm_lib
@@ -36,7 +35,7 @@
_lock = threading.Lock()
-_default_handler: Optional[logging.Handler] = None
+_default_handler: logging.Handler | None = None
log_levels = {
"debug": logging.DEBUG,
@@ -125,11 +124,11 @@ def _reset_library_root_logger() -> None:
_default_handler = None
-def get_log_levels_dict() -> Dict[str, int]:
+def get_log_levels_dict() -> dict[str, int]:
return log_levels
-def get_logger(name: Optional[str] = None) -> logging.Logger:
+def get_logger(name: str | None = None) -> logging.Logger:
"""
Return a logger with the specified name.
diff --git a/src/diffusers/utils/outputs.py b/src/diffusers/utils/outputs.py
index 2b20f6120ce3..b7d9a29ccfce 100644
--- a/src/diffusers/utils/outputs.py
+++ b/src/diffusers/utils/outputs.py
@@ -17,7 +17,7 @@
from collections import OrderedDict
from dataclasses import fields, is_dataclass
-from typing import Any, Tuple
+from typing import Any
import numpy as np
@@ -127,7 +127,7 @@ def __reduce__(self):
args = tuple(getattr(self, field.name) for field in fields(self))
return callable, args, *remaining
- def to_tuple(self) -> Tuple[Any, ...]:
+ def to_tuple(self) -> tuple[Any, ...]:
"""
Convert self to a tuple containing all the attributes/keys that are not `None`.
"""
diff --git a/src/diffusers/utils/peft_utils.py b/src/diffusers/utils/peft_utils.py
index 12066ee3f89b..65bcfe631e97 100644
--- a/src/diffusers/utils/peft_utils.py
+++ b/src/diffusers/utils/peft_utils.py
@@ -16,8 +16,8 @@
"""
import collections
+import functools
import importlib
-from typing import Optional
from packaging import version
@@ -123,7 +123,7 @@ def scale_lora_layers(model, weight):
module.scale_layer(weight)
-def unscale_lora_layers(model, weight: Optional[float] = None):
+def unscale_lora_layers(model, weight: float | None = None):
"""
Removes the previously passed weight given to the LoRA layers of the model.
@@ -275,6 +275,55 @@ def get_module_weight(weight_for_adapter, module_name):
module.set_scale(adapter_name, get_module_weight(weight, module_name))
+def apply_lora_scale(kwargs_name: str = "joint_attention_kwargs"):
+ """
+ Decorator to automatically handle LoRA layer scaling/unscaling in forward methods.
+
+ This decorator extracts the `lora_scale` from the specified kwargs parameter, applies scaling before the forward
+ pass, and ensures unscaling happens after, even if an exception occurs.
+
+ Args:
+ kwargs_name (`str`, defaults to `"joint_attention_kwargs"`):
+ The name of the keyword argument that contains the LoRA scale. Common values include
+ "joint_attention_kwargs", "attention_kwargs", "cross_attention_kwargs", etc.
+ """
+
+ def decorator(forward_fn):
+ @functools.wraps(forward_fn)
+ def wrapper(self, *args, **kwargs):
+ from . import USE_PEFT_BACKEND
+
+ lora_scale = 1.0
+ attention_kwargs = kwargs.get(kwargs_name)
+
+ if attention_kwargs is not None:
+ attention_kwargs = attention_kwargs.copy()
+ kwargs[kwargs_name] = attention_kwargs
+ lora_scale = attention_kwargs.pop("scale", 1.0)
+
+ if not USE_PEFT_BACKEND and lora_scale != 1.0:
+ logger.warning(
+ f"Passing `scale` via `{kwargs_name}` when not using the PEFT backend is ineffective."
+ )
+
+ # Apply LoRA scaling if using PEFT backend
+ if USE_PEFT_BACKEND:
+ scale_lora_layers(self, lora_scale)
+
+ try:
+ # Execute the forward pass
+ result = forward_fn(self, *args, **kwargs)
+ return result
+ finally:
+ # Always unscale, even if forward pass raises an exception
+ if USE_PEFT_BACKEND:
+ unscale_lora_layers(self, lora_scale)
+
+ return wrapper
+
+ return decorator
+
+
def check_peft_version(min_version: str) -> None:
r"""
Checks if the version of PEFT is compatible.
diff --git a/src/diffusers/utils/pil_utils.py b/src/diffusers/utils/pil_utils.py
index 76678070b697..72d4704fa945 100644
--- a/src/diffusers/utils/pil_utils.py
+++ b/src/diffusers/utils/pil_utils.py
@@ -1,5 +1,3 @@
-from typing import List
-
import PIL.Image
import PIL.ImageOps
from packaging import version
@@ -50,7 +48,7 @@ def numpy_to_pil(images):
return pil_images
-def make_image_grid(images: List[PIL.Image.Image], rows: int, cols: int, resize: int = None) -> PIL.Image.Image:
+def make_image_grid(images: list[PIL.Image.Image], rows: int, cols: int, resize: int = None) -> PIL.Image.Image:
"""
Prepares a single grid of images. Useful for visualization purposes.
"""
diff --git a/src/diffusers/utils/remote_utils.py b/src/diffusers/utils/remote_utils.py
index 6494dc14171a..2580bbf5dda9 100644
--- a/src/diffusers/utils/remote_utils.py
+++ b/src/diffusers/utils/remote_utils.py
@@ -13,9 +13,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
import io
import json
-from typing import List, Literal, Optional, Union, cast
+from typing import Literal, cast
import requests
@@ -58,18 +60,18 @@ def detect_image_type(data: bytes) -> str:
def check_inputs_decode(
endpoint: str,
tensor: "torch.Tensor",
- processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None,
+ processor: "VaeImageProcessor" | "VideoProcessor" | None = None,
do_scaling: bool = True,
- scaling_factor: Optional[float] = None,
- shift_factor: Optional[float] = None,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
output_type: Literal["mp4", "pil", "pt"] = "pil",
return_type: Literal["mp4", "pil", "pt"] = "pil",
image_format: Literal["png", "jpg"] = "jpg",
partial_postprocess: bool = False,
input_tensor_type: Literal["binary"] = "binary",
output_tensor_type: Literal["binary"] = "binary",
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
):
if tensor.ndim == 3 and height is None and width is None:
raise ValueError("`height` and `width` required for packed latents.")
@@ -91,7 +93,7 @@ def check_inputs_decode(
def postprocess_decode(
response: requests.Response,
- processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None,
+ processor: "VaeImageProcessor" | "VideoProcessor" | None = None,
output_type: Literal["mp4", "pil", "pt"] = "pil",
return_type: Literal["mp4", "pil", "pt"] = "pil",
partial_postprocess: bool = False,
@@ -117,7 +119,7 @@ def postprocess_decode(
else:
if isinstance(processor, VideoProcessor):
output = cast(
- List[Image.Image],
+ list[Image.Image],
processor.postprocess_video(output_tensor, output_type="pil")[0],
)
else:
@@ -144,15 +146,15 @@ def postprocess_decode(
def prepare_decode(
tensor: "torch.Tensor",
- processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None,
+ processor: "VaeImageProcessor" | "VideoProcessor" | None = None,
do_scaling: bool = True,
- scaling_factor: Optional[float] = None,
- shift_factor: Optional[float] = None,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
output_type: Literal["mp4", "pil", "pt"] = "pil",
image_format: Literal["png", "jpg"] = "jpg",
partial_postprocess: bool = False,
- height: Optional[int] = None,
- width: Optional[int] = None,
+ height: int | None = None,
+ width: int | None = None,
):
headers = {}
parameters = {
@@ -188,19 +190,19 @@ def prepare_decode(
def remote_decode(
endpoint: str,
tensor: "torch.Tensor",
- processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None,
+ processor: "VaeImageProcessor" | "VideoProcessor" | None = None,
do_scaling: bool = True,
- scaling_factor: Optional[float] = None,
- shift_factor: Optional[float] = None,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
output_type: Literal["mp4", "pil", "pt"] = "pil",
return_type: Literal["mp4", "pil", "pt"] = "pil",
image_format: Literal["png", "jpg"] = "jpg",
partial_postprocess: bool = False,
input_tensor_type: Literal["binary"] = "binary",
output_tensor_type: Literal["binary"] = "binary",
- height: Optional[int] = None,
- width: Optional[int] = None,
-) -> Union[Image.Image, List[Image.Image], bytes, "torch.Tensor"]:
+ height: int | None = None,
+ width: int | None = None,
+) -> Image.Image | list[Image.Image] | bytes | "torch.Tensor":
"""
Hugging Face Hybrid Inference that allow running VAE decode remotely.
@@ -275,7 +277,7 @@ def remote_decode(
Required for `"packed"` latents.
Returns:
- output (`Image.Image` or `List[Image.Image]` or `bytes` or `torch.Tensor`).
+ output (`Image.Image` or `list[Image.Image]` or `bytes` or `torch.Tensor`).
"""
if input_tensor_type == "base64":
deprecate(
@@ -336,9 +338,9 @@ def remote_decode(
def check_inputs_encode(
endpoint: str,
- image: Union["torch.Tensor", Image.Image],
- scaling_factor: Optional[float] = None,
- shift_factor: Optional[float] = None,
+ image: "torch.Tensor" | Image.Image,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
):
pass
@@ -356,9 +358,9 @@ def postprocess_encode(
def prepare_encode(
- image: Union["torch.Tensor", Image.Image],
- scaling_factor: Optional[float] = None,
- shift_factor: Optional[float] = None,
+ image: "torch.Tensor" | Image.Image,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
):
headers = {}
parameters = {}
@@ -379,9 +381,9 @@ def prepare_encode(
def remote_encode(
endpoint: str,
- image: Union["torch.Tensor", Image.Image],
- scaling_factor: Optional[float] = None,
- shift_factor: Optional[float] = None,
+ image: "torch.Tensor" | Image.Image,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
) -> "torch.Tensor":
"""
Hugging Face Hybrid Inference that allow running VAE encode remotely.
diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py
index 3297bb5fdcd6..cb7bf942c648 100644
--- a/src/diffusers/utils/testing_utils.py
+++ b/src/diffusers/utils/testing_utils.py
@@ -19,7 +19,7 @@
from contextlib import contextmanager
from io import BytesIO, StringIO
from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Set
import numpy as np
import PIL.Image
@@ -676,7 +676,7 @@ def get_python_version():
return major, minor
-def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -> np.ndarray:
+def load_numpy(arry: str | np.ndarray, local_path: str | None = None) -> np.ndarray:
if isinstance(arry, str):
if local_path is not None:
# local_path can be passed to correct images of tests
@@ -702,14 +702,14 @@ def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -
return arry
-def load_pt(url: str, map_location: Optional[str] = None, weights_only: Optional[bool] = True):
+def load_pt(url: str, map_location: str | None = None, weights_only: bool | None = True):
response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT)
response.raise_for_status()
arry = torch.load(BytesIO(response.content), map_location=map_location, weights_only=weights_only)
return arry
-def load_image(image: Union[str, PIL.Image.Image]) -> PIL.Image.Image:
+def load_image(image: str | PIL.Image.Image) -> PIL.Image.Image:
"""
Loads `image` to a PIL Image.
@@ -750,7 +750,7 @@ def preprocess_image(image: PIL.Image, batch_size: int):
return 2.0 * image - 1.0
-def export_to_gif(image: List[PIL.Image.Image], output_gif_path: str = None) -> str:
+def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None) -> str:
if output_gif_path is None:
output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name
@@ -843,7 +843,7 @@ def export_to_obj(mesh, output_obj_path: str = None):
f.writelines("\n".join(combined_data))
-def export_to_video(video_frames: List[np.ndarray], output_video_path: str = None) -> str:
+def export_to_video(video_frames: list[np.ndarray], output_video_path: str = None) -> str:
if is_opencv_available():
import cv2
else:
@@ -1024,7 +1024,7 @@ def summary_failures_short(tr):
# Adapted from https://github.com/huggingface/transformers/blob/000e52aec8850d3fe2f360adc6fd256e5b47fe4c/src/transformers/testing_utils.py#L1905
-def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None):
+def is_flaky(max_attempts: int = 5, wait_before_retry: float | None = None, description: str | None = None):
"""
To decorate flaky tests (methods or entire classes). They will be retried on failures.
@@ -1281,7 +1281,7 @@ def _is_torch_fp64_available(device):
# This dispatches a defined function according to the accelerator from the function definitions.
-def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable], *args, **kwargs):
+def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs):
if device not in dispatch_table:
return dispatch_table["default"](*args, **kwargs)
@@ -1381,7 +1381,7 @@ def backend_supports_training(device: str):
# Guard for when Torch is not available
if is_torch_available():
# Update device function dict mapping
- def update_mapping_from_spec(device_fn_dict: Dict[str, Callable], attribute_name: str):
+ def update_mapping_from_spec(device_fn_dict: dict[str, Callable], attribute_name: str):
try:
# Try to import the function directly
spec_fn = getattr(device_spec_module, attribute_name)
@@ -1430,7 +1430,7 @@ def update_mapping_from_spec(device_fn_dict: Dict[str, Callable], attribute_name
# Modified from https://github.com/huggingface/transformers/blob/cdfb018d0300fef3b07d9220f3efe9c2a9974662/src/transformers/testing_utils.py#L3090
# Type definition of key used in `Expectations` class.
-DeviceProperties = Tuple[Union[str, None], Union[int, None]]
+DeviceProperties = tuple[str | None, int | None]
@functools.lru_cache
@@ -1477,7 +1477,7 @@ def _get_expected_safetensors_files(
module: torch.nn.Module,
offload_to_disk_path: str,
offload_type: str,
- num_blocks_per_group: Optional[int] = None,
+ num_blocks_per_group: int | None = None,
) -> Set[str]:
expected_files = set()
@@ -1545,7 +1545,7 @@ def _check_safetensors_serialization(
module: torch.nn.Module,
offload_to_disk_path: str,
offload_type: str,
- num_blocks_per_group: Optional[int] = None,
+ num_blocks_per_group: int | None = None,
) -> bool:
if not os.path.isdir(offload_to_disk_path):
return False, None, None
diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py
index 3b66fdadbef8..7f4cb3e12766 100644
--- a/src/diffusers/utils/torch_utils.py
+++ b/src/diffusers/utils/torch_utils.py
@@ -15,9 +15,10 @@
PyTorch utilities: Utilities related to PyTorch
"""
+from __future__ import annotations
+
import functools
import os
-from typing import Callable, Dict, List, Optional, Tuple, Union
from . import logging
from .import_utils import is_torch_available, is_torch_mlu_available, is_torch_npu_available, is_torch_version
@@ -88,7 +89,7 @@ def maybe_allow_in_graph(cls):
# This dispatches a defined function according to the accelerator from the function definitions.
-def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable], *args, **kwargs):
+def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, callable], *args, **kwargs):
if device not in dispatch_table:
return dispatch_table["default"](*args, **kwargs)
@@ -144,11 +145,11 @@ def backend_supports_training(device: str):
def randn_tensor(
- shape: Union[Tuple, List],
- generator: Optional[Union[List["torch.Generator"], "torch.Generator"]] = None,
- device: Optional[Union[str, "torch.device"]] = None,
- dtype: Optional["torch.dtype"] = None,
- layout: Optional["torch.layout"] = None,
+ shape: tuple | list,
+ generator: list["torch.Generator"] | "torch.Generator" | None = None,
+ device: str | "torch.device" | None = None,
+ dtype: "torch.dtype" | None = None,
+ layout: "torch.layout" | None = None,
):
"""A helper function to create random tensors on the desired `device` with the desired `dtype`. When
passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor
@@ -241,9 +242,9 @@ def fourier_filter(x_in: "torch.Tensor", threshold: int, scale: int) -> "torch.T
def apply_freeu(
resolution_idx: int, hidden_states: "torch.Tensor", res_hidden_states: "torch.Tensor", **freeu_kwargs
-) -> Tuple["torch.Tensor", "torch.Tensor"]:
- """Applies the FreeU mechanism as introduced in https://huggingface.co/papers/2309.11497. Adapted from the official
- code repository: https://github.com/ChenyangSi/FreeU.
+) -> tuple["torch.Tensor", "torch.Tensor"]:
+ """Applies the FreeU mechanism as introduced in https:
+ //arxiv.org/abs/2309.11497. Adapted from the official code repository: https://github.com/ChenyangSi/FreeU.
Args:
resolution_idx (`int`): Integer denoting the UNet block where FreeU is being applied.
@@ -292,7 +293,7 @@ def get_device():
return "cpu"
-def empty_device_cache(device_type: Optional[str] = None):
+def empty_device_cache(device_type: str | None = None):
if device_type is None:
device_type = get_device()
if device_type in ["cpu"]:
@@ -301,7 +302,7 @@ def empty_device_cache(device_type: Optional[str] = None):
device_mod.empty_cache()
-def device_synchronize(device_type: Optional[str] = None):
+def device_synchronize(device_type: str | None = None):
if device_type is None:
device_type = get_device()
device_mod = getattr(torch, device_type, torch.cuda)
diff --git a/src/diffusers/utils/typing_utils.py b/src/diffusers/utils/typing_utils.py
index 2b5b1a4f5ab5..e910932f33d9 100644
--- a/src/diffusers/utils/typing_utils.py
+++ b/src/diffusers/utils/typing_utils.py
@@ -18,7 +18,7 @@
from typing import Any, Dict, List, Set, Tuple, Type, Union, get_args, get_origin
-def _is_valid_type(obj: Any, class_or_tuple: Union[Type, Tuple[Type, ...]]) -> bool:
+def _is_valid_type(obj: Any, class_or_tuple: Type | tuple[Type, ...]) -> bool:
"""
Checks if an object is an instance of any of the provided types. For collections, it checks if every element is of
the correct type as well.
diff --git a/src/diffusers/utils/versions.py b/src/diffusers/utils/versions.py
index 945a3977ce62..ccede7b5ee57 100644
--- a/src/diffusers/utils/versions.py
+++ b/src/diffusers/utils/versions.py
@@ -19,7 +19,6 @@
import operator
import re
import sys
-from typing import Optional
from packaging import version
@@ -46,7 +45,7 @@ def _compare_versions(op, got_ver, want_ver, requirement, pkg, hint):
)
-def require_version(requirement: str, hint: Optional[str] = None) -> None:
+def require_version(requirement: str, hint: str | None = None) -> None:
"""
Perform a runtime check of the dependency versions, using the exact same syntax used by pip.
diff --git a/src/diffusers/video_processor.py b/src/diffusers/video_processor.py
index abeb30bca102..34427686394d 100644
--- a/src/diffusers/video_processor.py
+++ b/src/diffusers/video_processor.py
@@ -13,7 +13,6 @@
# limitations under the License.
import warnings
-from typing import List, Optional, Tuple, Union
import numpy as np
import PIL
@@ -26,20 +25,20 @@
class VideoProcessor(VaeImageProcessor):
r"""Simple video processor."""
- def preprocess_video(self, video, height: Optional[int] = None, width: Optional[int] = None) -> torch.Tensor:
+ def preprocess_video(self, video, height: int | None = None, width: int | None = None) -> torch.Tensor:
r"""
Preprocesses input video(s).
Args:
- video (`List[PIL.Image]`, `List[List[PIL.Image]]`, `torch.Tensor`, `np.array`, `List[torch.Tensor]`, `List[np.array]`):
+ video (`list[PIL.Image]`, `list[list[PIL.Image]]`, `torch.Tensor`, `np.array`, `list[torch.Tensor]`, `list[np.array]`):
The input video. It can be one of the following:
- * List of the PIL images.
- * List of list of PIL images.
+ * list of the PIL images.
+ * list of list of PIL images.
* 4D Torch tensors (expected shape for each tensor `(num_frames, num_channels, height, width)`).
* 4D NumPy arrays (expected shape for each array `(num_frames, height, width, num_channels)`).
- * List of 4D Torch tensors (expected shape for each tensor `(num_frames, num_channels, height,
+ * list of 4D Torch tensors (expected shape for each tensor `(num_frames, num_channels, height,
width)`).
- * List of 4D NumPy arrays (expected shape for each array `(num_frames, height, width, num_channels)`).
+ * list of 4D NumPy arrays (expected shape for each array `(num_frames, height, width, num_channels)`).
* 5D NumPy arrays: expected shape for each array `(batch_size, num_frames, height, width,
num_channels)`.
* 5D Torch tensors: expected shape for each array `(batch_size, num_frames, num_channels, height,
@@ -89,7 +88,7 @@ def preprocess_video(self, video, height: Optional[int] = None, width: Optional[
def postprocess_video(
self, video: torch.Tensor, output_type: str = "np"
- ) -> Union[np.ndarray, torch.Tensor, List[PIL.Image.Image]]:
+ ) -> np.ndarray | torch.Tensor | list[PIL.Image.Image]:
r"""
Converts a video tensor to a list of frames for export.
@@ -114,7 +113,7 @@ def postprocess_video(
return outputs
@staticmethod
- def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[int, int]:
+ def classify_height_width_bin(height: int, width: int, ratios: dict) -> tuple[int, int]:
r"""
Returns the binned height and width based on the aspect ratio.
@@ -124,7 +123,7 @@ def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[in
ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width).
Returns:
- `Tuple[int, int]`: The closest binned height and width.
+ `tuple[int, int]`: The closest binned height and width.
"""
ar = float(height / width)
closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
diff --git a/tests/fixtures/custom_pipeline/pipeline.py b/tests/fixtures/custom_pipeline/pipeline.py
index 25673e566549..2d569d0f8abe 100644
--- a/tests/fixtures/custom_pipeline/pipeline.py
+++ b/tests/fixtures/custom_pipeline/pipeline.py
@@ -15,8 +15,6 @@
# limitations under the License.
-from typing import Optional, Tuple, Union
-
import torch
from diffusers import DiffusionPipeline, ImagePipelineOutput, SchedulerMixin, UNet2DModel
@@ -42,12 +40,12 @@ def __init__(self, unet: UNet2DModel, scheduler: SchedulerMixin):
def __call__(
self,
batch_size: int = 1,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
num_inference_steps: int = 50,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
**kwargs,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
r"""
Args:
batch_size (`int`, *optional*, defaults to 1):
diff --git a/tests/fixtures/custom_pipeline/what_ever.py b/tests/fixtures/custom_pipeline/what_ever.py
index 7504940780e8..afb140f64d32 100644
--- a/tests/fixtures/custom_pipeline/what_ever.py
+++ b/tests/fixtures/custom_pipeline/what_ever.py
@@ -15,8 +15,6 @@
# limitations under the License.
-from typing import Optional, Tuple, Union
-
import torch
from diffusers import SchedulerMixin, UNet2DModel
@@ -43,12 +41,12 @@ def __init__(self, unet: UNet2DModel, scheduler: SchedulerMixin):
def __call__(
self,
batch_size: int = 1,
- generator: Optional[torch.Generator] = None,
+ generator: torch.Generator | None = None,
num_inference_steps: int = 50,
- output_type: Optional[str] = "pil",
+ output_type: str | None = "pil",
return_dict: bool = True,
**kwargs,
- ) -> Union[ImagePipelineOutput, Tuple]:
+ ) -> ImagePipelineOutput | tuple:
r"""
Args:
batch_size (`int`, *optional*, defaults to 1):
diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py
index 5ae16ab4b9da..1a5b47ce4e5e 100644
--- a/tests/lora/test_lora_layers_wan.py
+++ b/tests/lora/test_lora_layers_wan.py
@@ -18,18 +18,9 @@
import torch
from transformers import AutoTokenizer, T5EncoderModel
-from diffusers import (
- AutoencoderKLWan,
- FlowMatchEulerDiscreteScheduler,
- WanPipeline,
- WanTransformer3DModel,
-)
-
-from ..testing_utils import (
- floats_tensor,
- require_peft_backend,
- skip_mps,
-)
+from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanPipeline, WanTransformer3DModel
+
+from ..testing_utils import floats_tensor, require_peft_backend, skip_mps
sys.path.append(".")
diff --git a/tests/models/controlnets/__init__.py b/tests/models/controlnets/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/models/controlnets/test_models_controlnet_cosmos.py b/tests/models/controlnets/test_models_controlnet_cosmos.py
new file mode 100644
index 000000000000..5241a49261b2
--- /dev/null
+++ b/tests/models/controlnets/test_models_controlnet_cosmos.py
@@ -0,0 +1,275 @@
+# Copyright 2025 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+
+from diffusers import CosmosControlNetModel
+from diffusers.models.controlnets.controlnet_cosmos import CosmosControlNetOutput
+
+from ...testing_utils import enable_full_determinism, torch_device
+from ..test_modeling_common import ModelTesterMixin
+
+
+enable_full_determinism()
+
+
+class CosmosControlNetModelTests(ModelTesterMixin, unittest.TestCase):
+ model_class = CosmosControlNetModel
+ main_input_name = "controls_latents"
+ uses_custom_attn_processor = True
+
+ @property
+ def dummy_input(self):
+ batch_size = 1
+ num_channels = 16
+ num_frames = 1
+ height = 16
+ width = 16
+ text_embed_dim = 32
+ sequence_length = 12
+ img_context_dim_in = 32
+ img_context_num_tokens = 4
+
+ # Raw latents (not patchified) - the controlnet computes embeddings internally
+ controls_latents = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device)
+ latents = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device)
+ timestep = torch.tensor([0.5]).to(torch_device) # Diffusion timestep
+ condition_mask = torch.ones(batch_size, 1, num_frames, height, width).to(torch_device)
+ padding_mask = torch.zeros(batch_size, 1, height, width).to(torch_device)
+
+ # Text embeddings
+ text_context = torch.randn((batch_size, sequence_length, text_embed_dim)).to(torch_device)
+ # Image context for Cosmos 2.5
+ img_context = torch.randn((batch_size, img_context_num_tokens, img_context_dim_in)).to(torch_device)
+ encoder_hidden_states = (text_context, img_context)
+
+ return {
+ "controls_latents": controls_latents,
+ "latents": latents,
+ "timestep": timestep,
+ "encoder_hidden_states": encoder_hidden_states,
+ "condition_mask": condition_mask,
+ "conditioning_scale": 1.0,
+ "padding_mask": padding_mask,
+ }
+
+ @property
+ def input_shape(self):
+ return (16, 1, 16, 16)
+
+ @property
+ def output_shape(self):
+ # Output is tuple of n_controlnet_blocks tensors, each with shape (batch, num_patches, model_channels)
+ # After stacking by normalize_output: (n_blocks, batch, num_patches, model_channels)
+ # For test config: n_blocks=2, num_patches=64 (1*8*8), model_channels=32
+ # output_shape is used as (batch_size,) + output_shape, so: (2, 64, 32)
+ return (2, 64, 32)
+
+ def prepare_init_args_and_inputs_for_common(self):
+ init_dict = {
+ "n_controlnet_blocks": 2,
+ "in_channels": 16 + 1 + 1, # control_latent_channels + condition_mask + padding_mask
+ "latent_channels": 16 + 1 + 1, # base_latent_channels (16) + condition_mask (1) + padding_mask (1) = 18
+ "model_channels": 32,
+ "num_attention_heads": 2,
+ "attention_head_dim": 16,
+ "mlp_ratio": 2,
+ "text_embed_dim": 32,
+ "adaln_lora_dim": 4,
+ "patch_size": (1, 2, 2),
+ "max_size": (4, 32, 32),
+ "rope_scale": (2.0, 1.0, 1.0),
+ "extra_pos_embed_type": None,
+ "img_context_dim_in": 32,
+ "img_context_dim_out": 32,
+ "use_crossattn_projection": False, # Test doesn't need this projection
+ }
+ inputs_dict = self.dummy_input
+ return init_dict, inputs_dict
+
+ def test_output_format(self):
+ """Test that the model outputs CosmosControlNetOutput with correct structure."""
+ init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+ model = self.model_class(**init_dict)
+ model.to(torch_device)
+ model.eval()
+
+ with torch.no_grad():
+ output = model(**inputs_dict)
+
+ self.assertIsInstance(output, CosmosControlNetOutput)
+ self.assertIsInstance(output.control_block_samples, list)
+ self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"])
+ for tensor in output.control_block_samples:
+ self.assertIsInstance(tensor, torch.Tensor)
+
+ def test_output_list_format(self):
+ """Test that return_dict=False returns a tuple containing a list."""
+ init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+ model = self.model_class(**init_dict)
+ model.to(torch_device)
+ model.eval()
+
+ with torch.no_grad():
+ output = model(**inputs_dict, return_dict=False)
+
+ self.assertIsInstance(output, tuple)
+ self.assertEqual(len(output), 1)
+ self.assertIsInstance(output[0], list)
+ self.assertEqual(len(output[0]), init_dict["n_controlnet_blocks"])
+
+ def test_condition_mask_changes_output(self):
+ """Test that condition mask affects control outputs."""
+ init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+ model = self.model_class(**init_dict)
+ model.to(torch_device)
+ model.eval()
+
+ inputs_no_mask = dict(inputs_dict)
+ inputs_no_mask["condition_mask"] = torch.zeros_like(inputs_dict["condition_mask"])
+
+ with torch.no_grad():
+ output_no_mask = model(**inputs_no_mask)
+ output_with_mask = model(**inputs_dict)
+
+ self.assertEqual(len(output_no_mask.control_block_samples), len(output_with_mask.control_block_samples))
+ for no_mask_tensor, with_mask_tensor in zip(
+ output_no_mask.control_block_samples, output_with_mask.control_block_samples
+ ):
+ self.assertFalse(torch.allclose(no_mask_tensor, with_mask_tensor))
+
+ def test_conditioning_scale_single(self):
+ """Test that a single conditioning scale is broadcast to all blocks."""
+ init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+ model = self.model_class(**init_dict)
+ model.to(torch_device)
+ model.eval()
+
+ inputs_dict["conditioning_scale"] = 0.5
+
+ with torch.no_grad():
+ output = model(**inputs_dict)
+
+ self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"])
+
+ def test_conditioning_scale_list(self):
+ """Test that a list of conditioning scales is applied per block."""
+ init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+ model = self.model_class(**init_dict)
+ model.to(torch_device)
+ model.eval()
+
+ # Provide a scale for each block
+ inputs_dict["conditioning_scale"] = [0.5, 1.0]
+
+ with torch.no_grad():
+ output = model(**inputs_dict)
+
+ self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"])
+
+ def test_forward_with_none_img_context(self):
+ """Test forward pass when img_context is None."""
+ init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+ model = self.model_class(**init_dict)
+ model.to(torch_device)
+ model.eval()
+
+ # Set encoder_hidden_states to (text_context, None)
+ text_context = inputs_dict["encoder_hidden_states"][0]
+ inputs_dict["encoder_hidden_states"] = (text_context, None)
+
+ with torch.no_grad():
+ output = model(**inputs_dict)
+
+ self.assertIsInstance(output, CosmosControlNetOutput)
+ self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"])
+
+ def test_forward_without_img_context_proj(self):
+ """Test forward pass when img_context_proj is not configured."""
+ init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+ # Disable img_context_proj
+ init_dict["img_context_dim_in"] = None
+ model = self.model_class(**init_dict)
+ model.to(torch_device)
+ model.eval()
+
+ # When img_context is disabled, pass only text context (not a tuple)
+ text_context = inputs_dict["encoder_hidden_states"][0]
+ inputs_dict["encoder_hidden_states"] = text_context
+
+ with torch.no_grad():
+ output = model(**inputs_dict)
+
+ self.assertIsInstance(output, CosmosControlNetOutput)
+ self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"])
+
+ def test_gradient_checkpointing_is_applied(self):
+ expected_set = {"CosmosControlNetModel"}
+ super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+ # Note: test_set_attn_processor_for_determinism already handles uses_custom_attn_processor=True
+ # so no explicit skip needed for it
+ # Note: test_forward_signature and test_set_default_attn_processor don't exist in base class
+
+ # Skip tests that don't apply to this architecture
+ @unittest.skip("CosmosControlNetModel doesn't use norm groups.")
+ def test_forward_with_norm_groups(self):
+ pass
+
+ # Skip tests that expect .sample attribute - ControlNets don't have this
+ @unittest.skip("ControlNet output doesn't have .sample attribute")
+ def test_effective_gradient_checkpointing(self):
+ pass
+
+ # Skip tests that compute MSE loss against single tensor output
+ @unittest.skip("ControlNet outputs list of control blocks, not single tensor for MSE loss")
+ def test_ema_training(self):
+ pass
+
+ @unittest.skip("ControlNet outputs list of control blocks, not single tensor for MSE loss")
+ def test_training(self):
+ pass
+
+ # Skip tests where output shape comparison doesn't apply to ControlNets
+ @unittest.skip("ControlNet output shape doesn't match input shape by design")
+ def test_output(self):
+ pass
+
+ # Skip outputs_equivalence - dict/list comparison logic not compatible (recursive_check expects dict.values())
+ @unittest.skip("ControlNet output structure not compatible with recursive dict check")
+ def test_outputs_equivalence(self):
+ pass
+
+ # Skip model parallelism - base test uses torch.allclose(base_output[0], new_output[0]) which fails
+ # because output[0] is the list of control_block_samples, not a tensor
+ @unittest.skip("test_model_parallelism uses torch.allclose on output[0] which is a list, not a tensor")
+ def test_model_parallelism(self):
+ pass
+
+ # Skip layerwise casting tests - these have two issues:
+ # 1. _inference and _memory: dtype compatibility issues with learnable_pos_embed and float8/bfloat16
+ # 2. _training: same as test_training - mse_loss expects tensor, not list
+ @unittest.skip("Layerwise casting has dtype issues with learnable_pos_embed")
+ def test_layerwise_casting_inference(self):
+ pass
+
+ @unittest.skip("Layerwise casting has dtype issues with learnable_pos_embed")
+ def test_layerwise_casting_memory(self):
+ pass
+
+ @unittest.skip("test_layerwise_casting_training computes mse_loss on list output")
+ def test_layerwise_casting_training(self):
+ pass
diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py
index b9dfe932335c..1b1a51d1e26f 100644
--- a/tests/models/test_modeling_common.py
+++ b/tests/models/test_modeling_common.py
@@ -26,7 +26,7 @@
import unittest.mock as mock
import uuid
from collections import defaultdict
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Tuple
import numpy as np
import pytest
@@ -168,8 +168,8 @@ def named_persistent_module_tensors(
def compute_module_persistent_sizes(
model: nn.Module,
- dtype: Optional[Union[str, torch.device]] = None,
- special_dtypes: Optional[Dict[str, Union[str, torch.device]]] = None,
+ dtype: str | torch.device | None = None,
+ special_dtypes: dict[str, str | torch.device] | None = None,
):
"""
Compute the size of each submodule of a given model (parameters + persistent buffers).
diff --git a/tests/models/test_models_auto.py b/tests/models/test_models_auto.py
index a70754343f30..3506f8fa0d17 100644
--- a/tests/models/test_models_auto.py
+++ b/tests/models/test_models_auto.py
@@ -1,5 +1,5 @@
import unittest
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
from transformers import CLIPTextModel, LongformerModel
@@ -20,7 +20,9 @@ def test_load_from_config_diffusers_with_subfolder(self, mock_load_config):
side_effect=[EnvironmentError("File not found"), {"model_type": "clip_text_model"}],
)
def test_load_from_config_transformers_with_subfolder(self, mock_load_config):
- model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder")
+ model = AutoModel.from_pretrained(
+ "hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder", use_safetensors=False
+ )
assert isinstance(model, CLIPTextModel)
def test_load_from_config_without_subfolder(self):
@@ -28,5 +30,73 @@ def test_load_from_config_without_subfolder(self):
assert isinstance(model, LongformerModel)
def test_load_from_model_index(self):
- model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder")
+ model = AutoModel.from_pretrained(
+ "hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder", use_safetensors=False
+ )
assert isinstance(model, CLIPTextModel)
+
+
+class TestAutoModelFromConfig(unittest.TestCase):
+ @patch(
+ "diffusers.pipelines.pipeline_loading_utils.get_class_obj_and_candidates",
+ return_value=(MagicMock(), None),
+ )
+ def test_from_config_with_dict_diffusers_class(self, mock_get_class):
+ config = {"_class_name": "UNet2DConditionModel", "sample_size": 64}
+ mock_model = MagicMock()
+ mock_get_class.return_value[0].from_config.return_value = mock_model
+
+ result = AutoModel.from_config(config)
+
+ mock_get_class.assert_called_once_with(
+ library_name="diffusers",
+ class_name="UNet2DConditionModel",
+ importable_classes=unittest.mock.ANY,
+ pipelines=None,
+ is_pipeline_module=False,
+ )
+ mock_get_class.return_value[0].from_config.assert_called_once_with(config)
+ assert result is mock_model
+
+ @patch(
+ "diffusers.pipelines.pipeline_loading_utils.get_class_obj_and_candidates",
+ return_value=(MagicMock(), None),
+ )
+ @patch("diffusers.models.AutoModel.load_config", return_value={"_class_name": "UNet2DConditionModel"})
+ def test_from_config_with_string_path(self, mock_load_config, mock_get_class):
+ mock_model = MagicMock()
+ mock_get_class.return_value[0].from_config.return_value = mock_model
+
+ result = AutoModel.from_config("hf-internal-testing/tiny-stable-diffusion-torch", subfolder="unet")
+
+ mock_load_config.assert_called_once()
+ assert result is mock_model
+
+ def test_from_config_raises_on_missing_class_info(self):
+ config = {"some_key": "some_value"}
+ with self.assertRaises(ValueError, msg="Couldn't find a model class"):
+ AutoModel.from_config(config)
+
+ @patch(
+ "diffusers.pipelines.pipeline_loading_utils.get_class_obj_and_candidates",
+ return_value=(MagicMock(), None),
+ )
+ def test_from_config_with_model_type_routes_to_transformers(self, mock_get_class):
+ config = {"model_type": "clip_text_model"}
+ mock_model = MagicMock()
+ mock_get_class.return_value[0].from_config.return_value = mock_model
+
+ result = AutoModel.from_config(config)
+
+ mock_get_class.assert_called_once_with(
+ library_name="transformers",
+ class_name="AutoModel",
+ importable_classes=unittest.mock.ANY,
+ pipelines=None,
+ is_pipeline_module=False,
+ )
+ assert result is mock_model
+
+ def test_from_config_raises_on_none(self):
+ with self.assertRaises(ValueError, msg="Please provide a `pretrained_model_name_or_path_or_dict`"):
+ AutoModel.from_config(None)
diff --git a/tests/models/testing_utils/common.py b/tests/models/testing_utils/common.py
index 145a6fc27f35..7036bb16203d 100644
--- a/tests/models/testing_utils/common.py
+++ b/tests/models/testing_utils/common.py
@@ -446,16 +446,17 @@ def test_getattr_is_correct(self, caplog):
torch_device not in ["cuda", "xpu"],
reason="float16 and bfloat16 can only be used with an accelerator",
)
- def test_keep_in_fp32_modules(self):
+ def test_keep_in_fp32_modules(self, tmp_path):
model = self.model_class(**self.get_init_dict())
fp32_modules = model._keep_in_fp32_modules
if fp32_modules is None or len(fp32_modules) == 0:
pytest.skip("Model does not have _keep_in_fp32_modules defined.")
- # Test with float16
- model.to(torch_device)
- model.to(torch.float16)
+ # Save the model and reload with float16 dtype
+ # _keep_in_fp32_modules is only enforced during from_pretrained loading
+ model.save_pretrained(tmp_path)
+ model = self.model_class.from_pretrained(tmp_path, torch_dtype=torch.float16).to(torch_device)
for name, param in model.named_parameters():
if any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules):
@@ -470,7 +471,7 @@ def test_keep_in_fp32_modules(self):
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
@torch.no_grad()
- def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
+ def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype, atol=1e-4, rtol=0):
model = self.model_class(**self.get_init_dict())
model.to(torch_device)
fp32_modules = model._keep_in_fp32_modules or []
@@ -490,10 +491,6 @@ def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
output = model(**inputs, return_dict=False)[0]
output_loaded = model_loaded(**inputs, return_dict=False)[0]
- self._check_dtype_inference_output(output, output_loaded, dtype)
-
- def _check_dtype_inference_output(self, output, output_loaded, dtype, atol=1e-4, rtol=0):
- """Check dtype inference output with configurable tolerance."""
assert_tensors_close(
output, output_loaded, atol=atol, rtol=rtol, msg=f"Loaded model output differs for {dtype}"
)
diff --git a/tests/models/testing_utils/compile.py b/tests/models/testing_utils/compile.py
index 950d4d5d1fa5..998b88fb469e 100644
--- a/tests/models/testing_utils/compile.py
+++ b/tests/models/testing_utils/compile.py
@@ -81,7 +81,7 @@ def test_torch_compile_recompilation_and_graph_break(self):
_ = model(**inputs_dict)
@torch.no_grad()
- def test_torch_compile_repeated_blocks(self):
+ def test_torch_compile_repeated_blocks(self, recompile_limit=1):
if self.model_class._repeated_blocks is None:
pytest.skip("Skipping test as the model class doesn't have `_repeated_blocks` set.")
@@ -92,7 +92,6 @@ def test_torch_compile_repeated_blocks(self):
model.eval()
model.compile_repeated_blocks(fullgraph=True)
- recompile_limit = 1
if self.model_class.__name__ == "UNet2DConditionModel":
recompile_limit = 2
diff --git a/tests/models/testing_utils/lora.py b/tests/models/testing_utils/lora.py
index 994aaed55ca7..dfdc4835ee88 100644
--- a/tests/models/testing_utils/lora.py
+++ b/tests/models/testing_utils/lora.py
@@ -375,7 +375,7 @@ def _check_model_hotswap(
# additionally check if dynamic compilation works.
if different_shapes is not None:
for height, width in different_shapes:
- new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
+ new_inputs_dict = self.get_dummy_inputs(height=height, width=width)
_ = model(**new_inputs_dict)
else:
output0_after = model(**inputs_dict)["sample"]
@@ -390,7 +390,7 @@ def _check_model_hotswap(
with torch.inference_mode():
if different_shapes is not None:
for height, width in different_shapes:
- new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
+ new_inputs_dict = self.get_dummy_inputs(height=height, width=width)
_ = model(**new_inputs_dict)
else:
output1_after = model(**inputs_dict)["sample"]
diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py
index f27e912766e5..0f1fbde72485 100644
--- a/tests/models/testing_utils/quantization.py
+++ b/tests/models/testing_utils/quantization.py
@@ -176,15 +176,7 @@ def _test_quantization_inference(self, config_kwargs):
model_quantized = self._create_quantized_model(config_kwargs)
model_quantized.to(torch_device)
- # Get model dtype from first parameter
- model_dtype = next(model_quantized.parameters()).dtype
-
inputs = self.get_dummy_inputs()
- # Cast inputs to model dtype
- inputs = {
- k: v.to(model_dtype) if isinstance(v, torch.Tensor) and v.is_floating_point() else v
- for k, v in inputs.items()
- }
output = model_quantized(**inputs, return_dict=False)[0]
assert output is not None, "Model output is None"
@@ -229,6 +221,8 @@ def _test_quantization_lora_inference(self, config_kwargs):
init_lora_weights=False,
)
model.add_adapter(lora_config)
+ # Move LoRA adapter weights to device (they default to CPU)
+ model.to(torch_device)
inputs = self.get_dummy_inputs()
output = model(**inputs, return_dict=False)[0]
@@ -634,6 +628,21 @@ def test_bnb_training(self):
"""Test that quantized models can be used for training with adapters."""
self._test_quantization_training(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"])
+ @pytest.mark.parametrize(
+ "config_name",
+ list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()),
+ ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()),
+ )
+ def test_cpu_device_map(self, config_name):
+ config_kwargs = BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]
+ model_quantized = self._create_quantized_model(config_kwargs, device_map="cpu")
+
+ assert hasattr(model_quantized, "hf_device_map"), "Model should have hf_device_map attribute"
+ assert model_quantized.hf_device_map is not None, "hf_device_map should not be None"
+ assert model_quantized.device == torch.device("cpu"), (
+ f"Model should be on CPU, but is on {model_quantized.device}"
+ )
+
@is_quantization
@is_quanto
@@ -1021,9 +1030,6 @@ def test_gguf_dequantize(self):
"""Test that dequantize() works correctly."""
self._test_dequantize({"compute_dtype": torch.bfloat16})
- def test_gguf_quantized_layers(self):
- self._test_quantized_layers({"compute_dtype": torch.bfloat16})
-
@is_quantization
@is_modelopt
diff --git a/tests/models/transformers/test_models_transformer_wan.py b/tests/models/transformers/test_models_transformer_wan.py
index 9f248f990c8a..26b0ac946434 100644
--- a/tests/models/transformers/test_models_transformer_wan.py
+++ b/tests/models/transformers/test_models_transformer_wan.py
@@ -12,57 +12,57 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-import unittest
-
+import pytest
import torch
from diffusers import WanTransformer3DModel
-
-from ...testing_utils import (
- enable_full_determinism,
- torch_device,
+from diffusers.utils.torch_utils import randn_tensor
+
+from ...testing_utils import enable_full_determinism, torch_device
+from ..testing_utils import (
+ AttentionTesterMixin,
+ BaseModelTesterConfig,
+ BitsAndBytesTesterMixin,
+ GGUFCompileTesterMixin,
+ GGUFTesterMixin,
+ MemoryTesterMixin,
+ ModelTesterMixin,
+ TorchAoTesterMixin,
+ TorchCompileTesterMixin,
+ TrainingTesterMixin,
)
-from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin
enable_full_determinism()
-class WanTransformer3DTests(ModelTesterMixin, unittest.TestCase):
- model_class = WanTransformer3DModel
- main_input_name = "hidden_states"
- uses_custom_attn_processor = True
+class WanTransformer3DTesterConfig(BaseModelTesterConfig):
+ @property
+ def model_class(self):
+ return WanTransformer3DModel
@property
- def dummy_input(self):
- batch_size = 1
- num_channels = 4
- num_frames = 2
- height = 16
- width = 16
- text_encoder_embedding_dim = 16
- sequence_length = 12
+ def pretrained_model_name_or_path(self):
+ return "hf-internal-testing/tiny-wan22-transformer"
- hidden_states = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device)
- timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device)
- encoder_hidden_states = torch.randn((batch_size, sequence_length, text_encoder_embedding_dim)).to(torch_device)
+ @property
+ def output_shape(self) -> tuple[int, ...]:
+ return (4, 2, 16, 16)
- return {
- "hidden_states": hidden_states,
- "encoder_hidden_states": encoder_hidden_states,
- "timestep": timestep,
- }
+ @property
+ def input_shape(self) -> tuple[int, ...]:
+ return (4, 2, 16, 16)
@property
- def input_shape(self):
- return (4, 1, 16, 16)
+ def main_input_name(self) -> str:
+ return "hidden_states"
@property
- def output_shape(self):
- return (4, 1, 16, 16)
+ def generator(self):
+ return torch.Generator("cpu").manual_seed(0)
- def prepare_init_args_and_inputs_for_common(self):
- init_dict = {
+ def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool]:
+ return {
"patch_size": (1, 2, 2),
"num_attention_heads": 2,
"attention_head_dim": 12,
@@ -76,16 +76,160 @@ def prepare_init_args_and_inputs_for_common(self):
"qk_norm": "rms_norm_across_heads",
"rope_max_seq_len": 32,
}
- inputs_dict = self.dummy_input
- return init_dict, inputs_dict
+
+ def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
+ batch_size = 1
+ num_channels = 4
+ num_frames = 2
+ height = 16
+ width = 16
+ text_encoder_embedding_dim = 16
+ sequence_length = 12
+
+ return {
+ "hidden_states": randn_tensor(
+ (batch_size, num_channels, num_frames, height, width),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (batch_size, sequence_length, text_encoder_embedding_dim),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
+ }
+
+
+class TestWanTransformer3D(WanTransformer3DTesterConfig, ModelTesterMixin):
+ """Core model tests for Wan Transformer 3D."""
+
+ @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
+ def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
+ # Skip: fp16/bf16 require very high atol to pass, providing little signal.
+ # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules.
+ pytest.skip("Tolerance requirements too high for meaningful test")
+
+
+class TestWanTransformer3DMemory(WanTransformer3DTesterConfig, MemoryTesterMixin):
+ """Memory optimization tests for Wan Transformer 3D."""
+
+
+class TestWanTransformer3DTraining(WanTransformer3DTesterConfig, TrainingTesterMixin):
+ """Training tests for Wan Transformer 3D."""
def test_gradient_checkpointing_is_applied(self):
expected_set = {"WanTransformer3DModel"}
super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
-class WanTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
- model_class = WanTransformer3DModel
+class TestWanTransformer3DAttention(WanTransformer3DTesterConfig, AttentionTesterMixin):
+ """Attention processor tests for Wan Transformer 3D."""
+
+
+class TestWanTransformer3DCompile(WanTransformer3DTesterConfig, TorchCompileTesterMixin):
+ """Torch compile tests for Wan Transformer 3D."""
- def prepare_init_args_and_inputs_for_common(self):
- return WanTransformer3DTests().prepare_init_args_and_inputs_for_common()
+
+class TestWanTransformer3DBitsAndBytes(WanTransformer3DTesterConfig, BitsAndBytesTesterMixin):
+ """BitsAndBytes quantization tests for Wan Transformer 3D."""
+
+ @property
+ def torch_dtype(self):
+ return torch.float16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the tiny Wan model dimensions."""
+ return {
+ "hidden_states": randn_tensor(
+ (1, 36, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanTransformer3DTorchAo(WanTransformer3DTesterConfig, TorchAoTesterMixin):
+ """TorchAO quantization tests for Wan Transformer 3D."""
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the tiny Wan model dimensions."""
+ return {
+ "hidden_states": randn_tensor(
+ (1, 36, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanTransformer3DGGUF(WanTransformer3DTesterConfig, GGUFTesterMixin):
+ """GGUF quantization tests for Wan Transformer 3D."""
+
+ @property
+ def gguf_filename(self):
+ return "https://huggingface.co/QuantStack/Wan2.2-I2V-A14B-GGUF/blob/main/LowNoise/Wan2.2-I2V-A14B-LowNoise-Q2_K.gguf"
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def _create_quantized_model(self, config_kwargs=None, **extra_kwargs):
+ return super()._create_quantized_model(
+ config_kwargs, config="Wan-AI/Wan2.2-I2V-A14B-Diffusers", subfolder="transformer", **extra_kwargs
+ )
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the real Wan I2V model dimensions.
+
+ Wan 2.2 I2V: in_channels=36, text_dim=4096
+ """
+ return {
+ "hidden_states": randn_tensor(
+ (1, 36, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanTransformer3DGGUFCompile(WanTransformer3DTesterConfig, GGUFCompileTesterMixin):
+ """GGUF + compile tests for Wan Transformer 3D."""
+
+ @property
+ def gguf_filename(self):
+ return "https://huggingface.co/QuantStack/Wan2.2-I2V-A14B-GGUF/blob/main/LowNoise/Wan2.2-I2V-A14B-LowNoise-Q2_K.gguf"
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def _create_quantized_model(self, config_kwargs=None, **extra_kwargs):
+ return super()._create_quantized_model(
+ config_kwargs, config="Wan-AI/Wan2.2-I2V-A14B-Diffusers", subfolder="transformer", **extra_kwargs
+ )
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the real Wan I2V model dimensions.
+
+ Wan 2.2 I2V: in_channels=36, text_dim=4096
+ """
+ return {
+ "hidden_states": randn_tensor(
+ (1, 36, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
diff --git a/tests/models/transformers/test_models_transformer_wan_animate.py b/tests/models/transformers/test_models_transformer_wan_animate.py
index 5d571b8c2e7d..ac0ef0698c63 100644
--- a/tests/models/transformers/test_models_transformer_wan_animate.py
+++ b/tests/models/transformers/test_models_transformer_wan_animate.py
@@ -12,76 +12,62 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-import unittest
-
+import pytest
import torch
from diffusers import WanAnimateTransformer3DModel
+from diffusers.utils.torch_utils import randn_tensor
-from ...testing_utils import (
- enable_full_determinism,
- torch_device,
+from ...testing_utils import enable_full_determinism, torch_device
+from ..testing_utils import (
+ AttentionTesterMixin,
+ BaseModelTesterConfig,
+ BitsAndBytesTesterMixin,
+ GGUFCompileTesterMixin,
+ GGUFTesterMixin,
+ MemoryTesterMixin,
+ ModelTesterMixin,
+ TorchAoTesterMixin,
+ TorchCompileTesterMixin,
+ TrainingTesterMixin,
)
-from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin
enable_full_determinism()
-class WanAnimateTransformer3DTests(ModelTesterMixin, unittest.TestCase):
- model_class = WanAnimateTransformer3DModel
- main_input_name = "hidden_states"
- uses_custom_attn_processor = True
-
+class WanAnimateTransformer3DTesterConfig(BaseModelTesterConfig):
@property
- def dummy_input(self):
- batch_size = 1
- num_channels = 4
- num_frames = 20 # To make the shapes work out; for complicated reasons we want 21 to divide num_frames + 1
- height = 16
- width = 16
- text_encoder_embedding_dim = 16
- sequence_length = 12
-
- clip_seq_len = 12
- clip_dim = 16
+ def model_class(self):
+ return WanAnimateTransformer3DModel
- inference_segment_length = 77 # The inference segment length in the full Wan2.2-Animate-14B model
- face_height = 16 # Should be square and match `motion_encoder_size` below
- face_width = 16
+ @property
+ def pretrained_model_name_or_path(self):
+ return "hf-internal-testing/tiny-wan-animate-transformer"
- hidden_states = torch.randn((batch_size, 2 * num_channels + 4, num_frames + 1, height, width)).to(torch_device)
- timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device)
- encoder_hidden_states = torch.randn((batch_size, sequence_length, text_encoder_embedding_dim)).to(torch_device)
- clip_ref_features = torch.randn((batch_size, clip_seq_len, clip_dim)).to(torch_device)
- pose_latents = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device)
- face_pixel_values = torch.randn((batch_size, 3, inference_segment_length, face_height, face_width)).to(
- torch_device
- )
+ @property
+ def output_shape(self) -> tuple[int, ...]:
+ # Output has fewer channels than input (4 vs 12)
+ return (4, 21, 16, 16)
- return {
- "hidden_states": hidden_states,
- "timestep": timestep,
- "encoder_hidden_states": encoder_hidden_states,
- "encoder_hidden_states_image": clip_ref_features,
- "pose_hidden_states": pose_latents,
- "face_pixel_values": face_pixel_values,
- }
+ @property
+ def input_shape(self) -> tuple[int, ...]:
+ return (12, 21, 16, 16)
@property
- def input_shape(self):
- return (12, 1, 16, 16)
+ def main_input_name(self) -> str:
+ return "hidden_states"
@property
- def output_shape(self):
- return (4, 1, 16, 16)
+ def generator(self):
+ return torch.Generator("cpu").manual_seed(0)
- def prepare_init_args_and_inputs_for_common(self):
+ def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool | float | dict]:
# Use custom channel sizes since the default Wan Animate channel sizes will cause the motion encoder to
# contain the vast majority of the parameters in the test model
channel_sizes = {"4": 16, "8": 16, "16": 16}
- init_dict = {
+ return {
"patch_size": (1, 2, 2),
"num_attention_heads": 2,
"attention_head_dim": 12,
@@ -105,22 +91,219 @@ def prepare_init_args_and_inputs_for_common(self):
"face_encoder_num_heads": 2,
"inject_face_latents_blocks": 2,
}
- inputs_dict = self.dummy_input
- return init_dict, inputs_dict
- def test_gradient_checkpointing_is_applied(self):
- expected_set = {"WanAnimateTransformer3DModel"}
- super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+ def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
+ batch_size = 1
+ num_channels = 4
+ num_frames = 20 # To make the shapes work out; for complicated reasons we want 21 to divide num_frames + 1
+ height = 16
+ width = 16
+ text_encoder_embedding_dim = 16
+ sequence_length = 12
+
+ clip_seq_len = 12
+ clip_dim = 16
+
+ inference_segment_length = 77 # The inference segment length in the full Wan2.2-Animate-14B model
+ face_height = 16 # Should be square and match `motion_encoder_size`
+ face_width = 16
+
+ return {
+ "hidden_states": randn_tensor(
+ (batch_size, 2 * num_channels + 4, num_frames + 1, height, width),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
+ "encoder_hidden_states": randn_tensor(
+ (batch_size, sequence_length, text_encoder_embedding_dim),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "encoder_hidden_states_image": randn_tensor(
+ (batch_size, clip_seq_len, clip_dim),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "pose_hidden_states": randn_tensor(
+ (batch_size, num_channels, num_frames, height, width),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "face_pixel_values": randn_tensor(
+ (batch_size, 3, inference_segment_length, face_height, face_width),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ }
+
+
+class TestWanAnimateTransformer3D(WanAnimateTransformer3DTesterConfig, ModelTesterMixin):
+ """Core model tests for Wan Animate Transformer 3D."""
- # Override test_output because the transformer output is expected to have less channels than the main transformer
- # input.
def test_output(self):
+ # Override test_output because the transformer output is expected to have less channels
+ # than the main transformer input.
expected_output_shape = (1, 4, 21, 16, 16)
super().test_output(expected_output_shape=expected_output_shape)
+ @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
+ def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
+ # Skip: fp16/bf16 require very high atol (~1e-2) to pass, providing little signal.
+ # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules.
+ pytest.skip("Tolerance requirements too high for meaningful test")
+
+
+class TestWanAnimateTransformer3DMemory(WanAnimateTransformer3DTesterConfig, MemoryTesterMixin):
+ """Memory optimization tests for Wan Animate Transformer 3D."""
-class WanAnimateTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
- model_class = WanAnimateTransformer3DModel
- def prepare_init_args_and_inputs_for_common(self):
- return WanAnimateTransformer3DTests().prepare_init_args_and_inputs_for_common()
+class TestWanAnimateTransformer3DTraining(WanAnimateTransformer3DTesterConfig, TrainingTesterMixin):
+ """Training tests for Wan Animate Transformer 3D."""
+
+ def test_gradient_checkpointing_is_applied(self):
+ expected_set = {"WanAnimateTransformer3DModel"}
+ super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+
+class TestWanAnimateTransformer3DAttention(WanAnimateTransformer3DTesterConfig, AttentionTesterMixin):
+ """Attention processor tests for Wan Animate Transformer 3D."""
+
+
+class TestWanAnimateTransformer3DCompile(WanAnimateTransformer3DTesterConfig, TorchCompileTesterMixin):
+ """Torch compile tests for Wan Animate Transformer 3D."""
+
+ def test_torch_compile_recompilation_and_graph_break(self):
+ # Skip: F.pad with mode="replicate" in WanAnimateFaceEncoder triggers importlib.import_module
+ # internally, which dynamo doesn't support tracing through.
+ pytest.skip("F.pad with replicate mode triggers unsupported import in torch.compile")
+
+
+class TestWanAnimateTransformer3DBitsAndBytes(WanAnimateTransformer3DTesterConfig, BitsAndBytesTesterMixin):
+ """BitsAndBytes quantization tests for Wan Animate Transformer 3D."""
+
+ @property
+ def torch_dtype(self):
+ return torch.float16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the tiny Wan Animate model dimensions."""
+ return {
+ "hidden_states": randn_tensor(
+ (1, 36, 21, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states_image": randn_tensor(
+ (1, 257, 1280), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "pose_hidden_states": randn_tensor(
+ (1, 16, 20, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "face_pixel_values": randn_tensor(
+ (1, 3, 77, 512, 512), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanAnimateTransformer3DTorchAo(WanAnimateTransformer3DTesterConfig, TorchAoTesterMixin):
+ """TorchAO quantization tests for Wan Animate Transformer 3D."""
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the tiny Wan Animate model dimensions."""
+ return {
+ "hidden_states": randn_tensor(
+ (1, 36, 21, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states_image": randn_tensor(
+ (1, 257, 1280), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "pose_hidden_states": randn_tensor(
+ (1, 16, 20, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "face_pixel_values": randn_tensor(
+ (1, 3, 77, 512, 512), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanAnimateTransformer3DGGUF(WanAnimateTransformer3DTesterConfig, GGUFTesterMixin):
+ """GGUF quantization tests for Wan Animate Transformer 3D."""
+
+ @property
+ def gguf_filename(self):
+ return "https://huggingface.co/QuantStack/Wan2.2-Animate-14B-GGUF/blob/main/Wan2.2-Animate-14B-Q2_K.gguf"
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the real Wan Animate model dimensions.
+
+ Wan 2.2 Animate: in_channels=36 (2*16+4), text_dim=4096, image_dim=1280
+ """
+ return {
+ "hidden_states": randn_tensor(
+ (1, 36, 21, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states_image": randn_tensor(
+ (1, 257, 1280), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "pose_hidden_states": randn_tensor(
+ (1, 16, 20, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "face_pixel_values": randn_tensor(
+ (1, 3, 77, 512, 512), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanAnimateTransformer3DGGUFCompile(WanAnimateTransformer3DTesterConfig, GGUFCompileTesterMixin):
+ """GGUF + compile tests for Wan Animate Transformer 3D."""
+
+ @property
+ def gguf_filename(self):
+ return "https://huggingface.co/QuantStack/Wan2.2-Animate-14B-GGUF/blob/main/Wan2.2-Animate-14B-Q2_K.gguf"
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the real Wan Animate model dimensions.
+
+ Wan 2.2 Animate: in_channels=36 (2*16+4), text_dim=4096, image_dim=1280
+ """
+ return {
+ "hidden_states": randn_tensor(
+ (1, 36, 21, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states_image": randn_tensor(
+ (1, 257, 1280), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "pose_hidden_states": randn_tensor(
+ (1, 16, 20, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "face_pixel_values": randn_tensor(
+ (1, 3, 77, 512, 512), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
diff --git a/tests/models/transformers/test_models_transformer_wan_vace.py b/tests/models/transformers/test_models_transformer_wan_vace.py
new file mode 100644
index 000000000000..5ab51bbb9003
--- /dev/null
+++ b/tests/models/transformers/test_models_transformer_wan_vace.py
@@ -0,0 +1,256 @@
+# Copyright 2025 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from diffusers import WanVACETransformer3DModel
+from diffusers.utils.torch_utils import randn_tensor
+
+from ...testing_utils import enable_full_determinism, torch_device
+from ..testing_utils import (
+ AttentionTesterMixin,
+ BaseModelTesterConfig,
+ BitsAndBytesTesterMixin,
+ GGUFCompileTesterMixin,
+ GGUFTesterMixin,
+ MemoryTesterMixin,
+ ModelTesterMixin,
+ TorchAoTesterMixin,
+ TorchCompileTesterMixin,
+ TrainingTesterMixin,
+)
+
+
+enable_full_determinism()
+
+
+class WanVACETransformer3DTesterConfig(BaseModelTesterConfig):
+ @property
+ def model_class(self):
+ return WanVACETransformer3DModel
+
+ @property
+ def pretrained_model_name_or_path(self):
+ return "hf-internal-testing/tiny-wan-vace-transformer"
+
+ @property
+ def output_shape(self) -> tuple[int, ...]:
+ return (16, 2, 16, 16)
+
+ @property
+ def input_shape(self) -> tuple[int, ...]:
+ return (16, 2, 16, 16)
+
+ @property
+ def main_input_name(self) -> str:
+ return "hidden_states"
+
+ @property
+ def generator(self):
+ return torch.Generator("cpu").manual_seed(0)
+
+ def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool | None]:
+ return {
+ "patch_size": (1, 2, 2),
+ "num_attention_heads": 2,
+ "attention_head_dim": 12,
+ "in_channels": 16,
+ "out_channels": 16,
+ "text_dim": 32,
+ "freq_dim": 256,
+ "ffn_dim": 32,
+ "num_layers": 4,
+ "cross_attn_norm": True,
+ "qk_norm": "rms_norm_across_heads",
+ "rope_max_seq_len": 32,
+ "vace_layers": [0, 2],
+ "vace_in_channels": 48, # 3 * in_channels = 3 * 16 = 48
+ }
+
+ def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
+ batch_size = 1
+ num_channels = 16
+ num_frames = 2
+ height = 16
+ width = 16
+ text_encoder_embedding_dim = 32
+ sequence_length = 12
+
+ # VACE requires control_hidden_states with vace_in_channels (3 * in_channels)
+ vace_in_channels = 48
+
+ return {
+ "hidden_states": randn_tensor(
+ (batch_size, num_channels, num_frames, height, width),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (batch_size, sequence_length, text_encoder_embedding_dim),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "control_hidden_states": randn_tensor(
+ (batch_size, vace_in_channels, num_frames, height, width),
+ generator=self.generator,
+ device=torch_device,
+ ),
+ "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
+ }
+
+
+class TestWanVACETransformer3D(WanVACETransformer3DTesterConfig, ModelTesterMixin):
+ """Core model tests for Wan VACE Transformer 3D."""
+
+ @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
+ def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
+ # Skip: fp16/bf16 require very high atol to pass, providing little signal.
+ # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules.
+ pytest.skip("Tolerance requirements too high for meaningful test")
+
+ def test_model_parallelism(self, tmp_path):
+ # Skip: Device mismatch between cuda:0 and cuda:1 in VACE control flow
+ pytest.skip("Model parallelism not yet supported for WanVACE")
+
+
+class TestWanVACETransformer3DMemory(WanVACETransformer3DTesterConfig, MemoryTesterMixin):
+ """Memory optimization tests for Wan VACE Transformer 3D."""
+
+
+class TestWanVACETransformer3DTraining(WanVACETransformer3DTesterConfig, TrainingTesterMixin):
+ """Training tests for Wan VACE Transformer 3D."""
+
+ def test_gradient_checkpointing_is_applied(self):
+ expected_set = {"WanVACETransformer3DModel"}
+ super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+
+class TestWanVACETransformer3DAttention(WanVACETransformer3DTesterConfig, AttentionTesterMixin):
+ """Attention processor tests for Wan VACE Transformer 3D."""
+
+
+class TestWanVACETransformer3DCompile(WanVACETransformer3DTesterConfig, TorchCompileTesterMixin):
+ """Torch compile tests for Wan VACE Transformer 3D."""
+
+ def test_torch_compile_repeated_blocks(self):
+ # WanVACE has two block types (WanTransformerBlock and WanVACETransformerBlock),
+ # so we need recompile_limit=2 instead of the default 1.
+ super().test_torch_compile_repeated_blocks(recompile_limit=2)
+
+
+class TestWanVACETransformer3DBitsAndBytes(WanVACETransformer3DTesterConfig, BitsAndBytesTesterMixin):
+ """BitsAndBytes quantization tests for Wan VACE Transformer 3D."""
+
+ @property
+ def torch_dtype(self):
+ return torch.float16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the tiny Wan VACE model dimensions."""
+ return {
+ "hidden_states": randn_tensor(
+ (1, 16, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "control_hidden_states": randn_tensor(
+ (1, 96, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanVACETransformer3DTorchAo(WanVACETransformer3DTesterConfig, TorchAoTesterMixin):
+ """TorchAO quantization tests for Wan VACE Transformer 3D."""
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the tiny Wan VACE model dimensions."""
+ return {
+ "hidden_states": randn_tensor(
+ (1, 16, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "control_hidden_states": randn_tensor(
+ (1, 96, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanVACETransformer3DGGUF(WanVACETransformer3DTesterConfig, GGUFTesterMixin):
+ """GGUF quantization tests for Wan VACE Transformer 3D."""
+
+ @property
+ def gguf_filename(self):
+ return "https://huggingface.co/QuantStack/Wan2.1_14B_VACE-GGUF/blob/main/Wan2.1_14B_VACE-Q3_K_S.gguf"
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the real Wan VACE model dimensions.
+
+ Wan 2.1 VACE: in_channels=16, text_dim=4096, vace_in_channels=96
+ """
+ return {
+ "hidden_states": randn_tensor(
+ (1, 16, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "control_hidden_states": randn_tensor(
+ (1, 96, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
+
+
+class TestWanVACETransformer3DGGUFCompile(WanVACETransformer3DTesterConfig, GGUFCompileTesterMixin):
+ """GGUF + compile tests for Wan VACE Transformer 3D."""
+
+ @property
+ def gguf_filename(self):
+ return "https://huggingface.co/QuantStack/Wan2.1_14B_VACE-GGUF/blob/main/Wan2.1_14B_VACE-Q3_K_S.gguf"
+
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
+ def get_dummy_inputs(self):
+ """Override to provide inputs matching the real Wan VACE model dimensions.
+
+ Wan 2.1 VACE: in_channels=16, text_dim=4096, vace_in_channels=96
+ """
+ return {
+ "hidden_states": randn_tensor(
+ (1, 16, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "encoder_hidden_states": randn_tensor(
+ (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "control_hidden_states": randn_tensor(
+ (1, 96, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+ ),
+ "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype),
+ }
diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py
index 854b5218c617..9a6b4b9b6fb4 100644
--- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py
+++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py
@@ -33,6 +33,19 @@
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
+FLUX_TEXT2IMAGE_WORKFLOWS = {
+ "text2image": [
+ ("text_encoder", "FluxTextEncoderStep"),
+ ("denoise.input", "FluxTextInputStep"),
+ ("denoise.before_denoise.prepare_latents", "FluxPrepareLatentsStep"),
+ ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"),
+ ("denoise.denoise", "FluxDenoiseStep"),
+ ("decode", "FluxDecodeStep"),
+ ]
+}
+
+
class TestFluxModularPipelineFast(ModularPipelineTesterMixin):
pipeline_class = FluxModularPipeline
pipeline_blocks_class = FluxAutoBlocks
@@ -40,6 +53,7 @@ class TestFluxModularPipelineFast(ModularPipelineTesterMixin):
params = frozenset(["prompt", "height", "width", "guidance_scale"])
batch_params = frozenset(["prompt"])
+ expected_workflow_blocks = FLUX_TEXT2IMAGE_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
@@ -59,6 +73,23 @@ def test_float16_inference(self):
super().test_float16_inference(9e-2)
+FLUX_IMAGE2IMAGE_WORKFLOWS = {
+ "image2image": [
+ ("text_encoder", "FluxTextEncoderStep"),
+ ("vae_encoder.preprocess", "FluxProcessImagesInputStep"),
+ ("vae_encoder.encode", "FluxVaeEncoderStep"),
+ ("denoise.input.text_inputs", "FluxTextInputStep"),
+ ("denoise.input.additional_inputs", "FluxAdditionalInputsStep"),
+ ("denoise.before_denoise.prepare_latents", "FluxPrepareLatentsStep"),
+ ("denoise.before_denoise.set_timesteps", "FluxImg2ImgSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_img2img_latents", "FluxImg2ImgPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"),
+ ("denoise.denoise", "FluxDenoiseStep"),
+ ("decode", "FluxDecodeStep"),
+ ]
+}
+
+
class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
pipeline_class = FluxModularPipeline
pipeline_blocks_class = FluxAutoBlocks
@@ -66,6 +97,7 @@ class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
params = frozenset(["prompt", "height", "width", "guidance_scale", "image"])
batch_params = frozenset(["prompt", "image"])
+ expected_workflow_blocks = FLUX_IMAGE2IMAGE_WORKFLOWS
def get_pipeline(self, components_manager=None, torch_dtype=torch.float32):
pipeline = super().get_pipeline(components_manager, torch_dtype)
@@ -125,6 +157,32 @@ def test_float16_inference(self):
super().test_float16_inference(8e-2)
+FLUX_KONTEXT_WORKFLOWS = {
+ "text2image": [
+ ("text_encoder", "FluxTextEncoderStep"),
+ ("denoise.input", "FluxTextInputStep"),
+ ("denoise.before_denoise.prepare_latents", "FluxPrepareLatentsStep"),
+ ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"),
+ ("denoise.denoise", "FluxKontextDenoiseStep"),
+ ("decode", "FluxDecodeStep"),
+ ],
+ "image_conditioned": [
+ ("text_encoder", "FluxTextEncoderStep"),
+ ("vae_encoder.preprocess", "FluxKontextProcessImagesInputStep"),
+ ("vae_encoder.encode", "FluxVaeEncoderStep"),
+ ("denoise.input.set_resolution", "FluxKontextSetResolutionStep"),
+ ("denoise.input.text_inputs", "FluxTextInputStep"),
+ ("denoise.input.additional_inputs", "FluxKontextAdditionalInputsStep"),
+ ("denoise.before_denoise.prepare_latents", "FluxPrepareLatentsStep"),
+ ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_rope_inputs", "FluxKontextRoPEInputsStep"),
+ ("denoise.denoise", "FluxKontextDenoiseStep"),
+ ("decode", "FluxDecodeStep"),
+ ],
+}
+
+
class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin):
pipeline_class = FluxKontextModularPipeline
pipeline_blocks_class = FluxKontextAutoBlocks
@@ -132,6 +190,7 @@ class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin):
params = frozenset(["prompt", "height", "width", "guidance_scale", "image"])
batch_params = frozenset(["prompt", "image"])
+ expected_workflow_blocks = FLUX_KONTEXT_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py
index 8fd529e97e71..3045af636841 100644
--- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py
+++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py
@@ -28,6 +28,21 @@
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
+FLUX2_TEXT2IMAGE_WORKFLOWS = {
+ "text2image": [
+ ("text_encoder", "Flux2TextEncoderStep"),
+ ("denoise.input", "Flux2TextInputStep"),
+ ("denoise.prepare_latents", "Flux2PrepareLatentsStep"),
+ ("denoise.set_timesteps", "Flux2SetTimestepsStep"),
+ ("denoise.prepare_guidance", "Flux2PrepareGuidanceStep"),
+ ("denoise.prepare_rope_inputs", "Flux2RoPEInputsStep"),
+ ("denoise.denoise", "Flux2DenoiseStep"),
+ ("denoise.after_denoise", "Flux2UnpackLatentsStep"),
+ ("decode", "Flux2DecodeStep"),
+ ],
+}
+
+
class TestFlux2ModularPipelineFast(ModularPipelineTesterMixin):
pipeline_class = Flux2ModularPipeline
pipeline_blocks_class = Flux2AutoBlocks
@@ -35,6 +50,7 @@ class TestFlux2ModularPipelineFast(ModularPipelineTesterMixin):
params = frozenset(["prompt", "height", "width", "guidance_scale"])
batch_params = frozenset(["prompt"])
+ expected_workflow_blocks = FLUX2_TEXT2IMAGE_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
@@ -56,6 +72,24 @@ def test_float16_inference(self):
super().test_float16_inference(9e-2)
+FLUX2_IMAGE_CONDITIONED_WORKFLOWS = {
+ "image_conditioned": [
+ ("text_encoder", "Flux2TextEncoderStep"),
+ ("vae_encoder.preprocess", "Flux2ProcessImagesInputStep"),
+ ("vae_encoder.encode", "Flux2VaeEncoderStep"),
+ ("denoise.input", "Flux2TextInputStep"),
+ ("denoise.prepare_image_latents", "Flux2PrepareImageLatentsStep"),
+ ("denoise.prepare_latents", "Flux2PrepareLatentsStep"),
+ ("denoise.set_timesteps", "Flux2SetTimestepsStep"),
+ ("denoise.prepare_guidance", "Flux2PrepareGuidanceStep"),
+ ("denoise.prepare_rope_inputs", "Flux2RoPEInputsStep"),
+ ("denoise.denoise", "Flux2DenoiseStep"),
+ ("denoise.after_denoise", "Flux2UnpackLatentsStep"),
+ ("decode", "Flux2DecodeStep"),
+ ],
+}
+
+
class TestFlux2ImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
pipeline_class = Flux2ModularPipeline
pipeline_blocks_class = Flux2AutoBlocks
@@ -63,6 +97,7 @@ class TestFlux2ImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
params = frozenset(["prompt", "height", "width", "guidance_scale", "image"])
batch_params = frozenset(["prompt", "image"])
+ expected_workflow_blocks = FLUX2_IMAGE_CONDITIONED_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py
index 26653b20f8c4..ad295a961357 100644
--- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py
+++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py
@@ -28,13 +28,28 @@
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
-class TestFlux2ModularPipelineFast(ModularPipelineTesterMixin):
+FLUX2_KLEIN_WORKFLOWS = {
+ "text2image": [
+ ("text_encoder", "Flux2KleinTextEncoderStep"),
+ ("denoise.input", "Flux2TextInputStep"),
+ ("denoise.prepare_latents", "Flux2PrepareLatentsStep"),
+ ("denoise.set_timesteps", "Flux2SetTimestepsStep"),
+ ("denoise.prepare_rope_inputs", "Flux2RoPEInputsStep"),
+ ("denoise.denoise", "Flux2KleinDenoiseStep"),
+ ("denoise.after_denoise", "Flux2UnpackLatentsStep"),
+ ("decode", "Flux2DecodeStep"),
+ ],
+}
+
+
+class TestFlux2KleinModularPipelineFast(ModularPipelineTesterMixin):
pipeline_class = Flux2KleinModularPipeline
pipeline_blocks_class = Flux2KleinAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-modular"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
+ expected_workflow_blocks = FLUX2_KLEIN_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
@@ -55,13 +70,31 @@ def test_float16_inference(self):
super().test_float16_inference(9e-2)
-class TestFlux2ImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
+FLUX2_KLEIN_IMAGE_CONDITIONED_WORKFLOWS = {
+ "image_conditioned": [
+ ("text_encoder", "Flux2KleinTextEncoderStep"),
+ ("vae_encoder.preprocess", "Flux2ProcessImagesInputStep"),
+ ("vae_encoder.encode", "Flux2VaeEncoderStep"),
+ ("denoise.input", "Flux2TextInputStep"),
+ ("denoise.prepare_image_latents", "Flux2PrepareImageLatentsStep"),
+ ("denoise.prepare_latents", "Flux2PrepareLatentsStep"),
+ ("denoise.set_timesteps", "Flux2SetTimestepsStep"),
+ ("denoise.prepare_rope_inputs", "Flux2RoPEInputsStep"),
+ ("denoise.denoise", "Flux2KleinDenoiseStep"),
+ ("denoise.after_denoise", "Flux2UnpackLatentsStep"),
+ ("decode", "Flux2DecodeStep"),
+ ],
+}
+
+
+class TestFlux2KleinImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
pipeline_class = Flux2KleinModularPipeline
pipeline_blocks_class = Flux2KleinAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-modular"
params = frozenset(["prompt", "height", "width", "image"])
batch_params = frozenset(["prompt", "image"])
+ expected_workflow_blocks = FLUX2_KLEIN_IMAGE_CONDITIONED_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py
index 701dd0fed896..b3aa79040317 100644
--- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py
+++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py
@@ -21,20 +21,35 @@
from diffusers.modular_pipelines import (
Flux2KleinBaseAutoBlocks,
- Flux2KleinModularPipeline,
+ Flux2KleinBaseModularPipeline,
)
from ...testing_utils import floats_tensor, torch_device
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
-class TestFlux2ModularPipelineFast(ModularPipelineTesterMixin):
- pipeline_class = Flux2KleinModularPipeline
+FLUX2_KLEIN_BASE_WORKFLOWS = {
+ "text2image": [
+ ("text_encoder", "Flux2KleinBaseTextEncoderStep"),
+ ("denoise.input", "Flux2KleinBaseTextInputStep"),
+ ("denoise.prepare_latents", "Flux2PrepareLatentsStep"),
+ ("denoise.set_timesteps", "Flux2SetTimestepsStep"),
+ ("denoise.prepare_rope_inputs", "Flux2KleinBaseRoPEInputsStep"),
+ ("denoise.denoise", "Flux2KleinBaseDenoiseStep"),
+ ("denoise.after_denoise", "Flux2UnpackLatentsStep"),
+ ("decode", "Flux2DecodeStep"),
+ ],
+}
+
+
+class TestFlux2KleinBaseModularPipelineFast(ModularPipelineTesterMixin):
+ pipeline_class = Flux2KleinBaseModularPipeline
pipeline_blocks_class = Flux2KleinBaseAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-base-modular"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
+ expected_workflow_blocks = FLUX2_KLEIN_BASE_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
@@ -55,13 +70,31 @@ def test_float16_inference(self):
super().test_float16_inference(9e-2)
-class TestFlux2ImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
- pipeline_class = Flux2KleinModularPipeline
+FLUX2_KLEIN_BASE_IMAGE_CONDITIONED_WORKFLOWS = {
+ "image_conditioned": [
+ ("text_encoder", "Flux2KleinBaseTextEncoderStep"),
+ ("vae_encoder.preprocess", "Flux2ProcessImagesInputStep"),
+ ("vae_encoder.encode", "Flux2VaeEncoderStep"),
+ ("denoise.input", "Flux2KleinBaseTextInputStep"),
+ ("denoise.prepare_latents", "Flux2PrepareLatentsStep"),
+ ("denoise.prepare_image_latents", "Flux2PrepareImageLatentsStep"),
+ ("denoise.set_timesteps", "Flux2SetTimestepsStep"),
+ ("denoise.prepare_rope_inputs", "Flux2KleinBaseRoPEInputsStep"),
+ ("denoise.denoise", "Flux2KleinBaseDenoiseStep"),
+ ("denoise.after_denoise", "Flux2UnpackLatentsStep"),
+ ("decode", "Flux2DecodeStep"),
+ ],
+}
+
+
+class TestFlux2KleinBaseImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
+ pipeline_class = Flux2KleinBaseModularPipeline
pipeline_blocks_class = Flux2KleinBaseAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-base-modular"
params = frozenset(["prompt", "height", "width", "image"])
batch_params = frozenset(["prompt", "image"])
+ expected_workflow_blocks = FLUX2_KLEIN_BASE_IMAGE_CONDITIONED_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
diff --git a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py
index f4bd27b7ea47..92573c202e49 100644
--- a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py
+++ b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py
@@ -30,6 +30,103 @@
from ..test_modular_pipelines_common import ModularGuiderTesterMixin, ModularPipelineTesterMixin
+QWEN_IMAGE_TEXT2IMAGE_WORKFLOWS = {
+ "text2image": [
+ ("text_encoder", "QwenImageTextEncoderStep"),
+ ("denoise.input", "QwenImageTextInputsStep"),
+ ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "QwenImageSetTimestepsStep"),
+ ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"),
+ ("denoise.denoise", "QwenImageDenoiseStep"),
+ ("denoise.after_denoise", "QwenImageAfterDenoiseStep"),
+ ("decode.decode", "QwenImageDecoderStep"),
+ ("decode.postprocess", "QwenImageProcessImagesOutputStep"),
+ ],
+ "image2image": [
+ ("text_encoder", "QwenImageTextEncoderStep"),
+ ("vae_encoder.preprocess", "QwenImageProcessImagesInputStep"),
+ ("vae_encoder.encode", "QwenImageVaeEncoderStep"),
+ ("denoise.input.text_inputs", "QwenImageTextInputsStep"),
+ ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"),
+ ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"),
+ ("denoise.prepare_img2img_latents", "QwenImagePrepareLatentsWithStrengthStep"),
+ ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"),
+ ("denoise.denoise", "QwenImageDenoiseStep"),
+ ("denoise.after_denoise", "QwenImageAfterDenoiseStep"),
+ ("decode.decode", "QwenImageDecoderStep"),
+ ("decode.postprocess", "QwenImageProcessImagesOutputStep"),
+ ],
+ "inpainting": [
+ ("text_encoder", "QwenImageTextEncoderStep"),
+ ("vae_encoder.preprocess", "QwenImageInpaintProcessImagesInputStep"),
+ ("vae_encoder.encode", "QwenImageVaeEncoderStep"),
+ ("denoise.input.text_inputs", "QwenImageTextInputsStep"),
+ ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"),
+ ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"),
+ ("denoise.prepare_inpaint_latents.add_noise_to_latents", "QwenImagePrepareLatentsWithStrengthStep"),
+ ("denoise.prepare_inpaint_latents.create_mask_latents", "QwenImageCreateMaskLatentsStep"),
+ ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"),
+ ("denoise.denoise", "QwenImageInpaintDenoiseStep"),
+ ("denoise.after_denoise", "QwenImageAfterDenoiseStep"),
+ ("decode.decode", "QwenImageDecoderStep"),
+ ("decode.postprocess", "QwenImageInpaintProcessImagesOutputStep"),
+ ],
+ "controlnet_text2image": [
+ ("text_encoder", "QwenImageTextEncoderStep"),
+ ("controlnet_vae_encoder", "QwenImageControlNetVaeEncoderStep"),
+ ("denoise.input", "QwenImageTextInputsStep"),
+ ("denoise.controlnet_input", "QwenImageControlNetInputsStep"),
+ ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "QwenImageSetTimestepsStep"),
+ ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"),
+ ("denoise.controlnet_before_denoise", "QwenImageControlNetBeforeDenoiserStep"),
+ ("denoise.controlnet_denoise", "QwenImageControlNetDenoiseStep"),
+ ("denoise.after_denoise", "QwenImageAfterDenoiseStep"),
+ ("decode.decode", "QwenImageDecoderStep"),
+ ("decode.postprocess", "QwenImageProcessImagesOutputStep"),
+ ],
+ "controlnet_image2image": [
+ ("text_encoder", "QwenImageTextEncoderStep"),
+ ("vae_encoder.preprocess", "QwenImageProcessImagesInputStep"),
+ ("vae_encoder.encode", "QwenImageVaeEncoderStep"),
+ ("controlnet_vae_encoder", "QwenImageControlNetVaeEncoderStep"),
+ ("denoise.input.text_inputs", "QwenImageTextInputsStep"),
+ ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"),
+ ("denoise.controlnet_input", "QwenImageControlNetInputsStep"),
+ ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"),
+ ("denoise.prepare_img2img_latents", "QwenImagePrepareLatentsWithStrengthStep"),
+ ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"),
+ ("denoise.controlnet_before_denoise", "QwenImageControlNetBeforeDenoiserStep"),
+ ("denoise.controlnet_denoise", "QwenImageControlNetDenoiseStep"),
+ ("denoise.after_denoise", "QwenImageAfterDenoiseStep"),
+ ("decode.decode", "QwenImageDecoderStep"),
+ ("decode.postprocess", "QwenImageProcessImagesOutputStep"),
+ ],
+ "controlnet_inpainting": [
+ ("text_encoder", "QwenImageTextEncoderStep"),
+ ("vae_encoder.preprocess", "QwenImageInpaintProcessImagesInputStep"),
+ ("vae_encoder.encode", "QwenImageVaeEncoderStep"),
+ ("controlnet_vae_encoder", "QwenImageControlNetVaeEncoderStep"),
+ ("denoise.input.text_inputs", "QwenImageTextInputsStep"),
+ ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"),
+ ("denoise.controlnet_input", "QwenImageControlNetInputsStep"),
+ ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"),
+ ("denoise.prepare_inpaint_latents.add_noise_to_latents", "QwenImagePrepareLatentsWithStrengthStep"),
+ ("denoise.prepare_inpaint_latents.create_mask_latents", "QwenImageCreateMaskLatentsStep"),
+ ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"),
+ ("denoise.controlnet_before_denoise", "QwenImageControlNetBeforeDenoiserStep"),
+ ("denoise.controlnet_denoise", "QwenImageInpaintControlNetDenoiseStep"),
+ ("denoise.after_denoise", "QwenImageAfterDenoiseStep"),
+ ("decode.decode", "QwenImageDecoderStep"),
+ ("decode.postprocess", "QwenImageInpaintProcessImagesOutputStep"),
+ ],
+}
+
+
class TestQwenImageModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin):
pipeline_class = QwenImageModularPipeline
pipeline_blocks_class = QwenImageAutoBlocks
@@ -37,6 +134,7 @@ class TestQwenImageModularPipelineFast(ModularPipelineTesterMixin, ModularGuider
params = frozenset(["prompt", "height", "width", "negative_prompt", "attention_kwargs", "image", "mask_image"])
batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"])
+ expected_workflow_blocks = QWEN_IMAGE_TEXT2IMAGE_WORKFLOWS
def get_dummy_inputs(self):
generator = self.get_generator()
@@ -56,6 +154,44 @@ def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=5e-4)
+QWEN_IMAGE_EDIT_WORKFLOWS = {
+ "image_conditioned": [
+ ("text_encoder.resize", "QwenImageEditResizeStep"),
+ ("text_encoder.encode", "QwenImageEditTextEncoderStep"),
+ ("vae_encoder.resize", "QwenImageEditResizeStep"),
+ ("vae_encoder.preprocess", "QwenImageEditProcessImagesInputStep"),
+ ("vae_encoder.encode", "QwenImageVaeEncoderStep"),
+ ("denoise.input.text_inputs", "QwenImageTextInputsStep"),
+ ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"),
+ ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "QwenImageSetTimestepsStep"),
+ ("denoise.prepare_rope_inputs", "QwenImageEditRoPEInputsStep"),
+ ("denoise.denoise", "QwenImageEditDenoiseStep"),
+ ("denoise.after_denoise", "QwenImageAfterDenoiseStep"),
+ ("decode.decode", "QwenImageDecoderStep"),
+ ("decode.postprocess", "QwenImageProcessImagesOutputStep"),
+ ],
+ "image_conditioned_inpainting": [
+ ("text_encoder.resize", "QwenImageEditResizeStep"),
+ ("text_encoder.encode", "QwenImageEditTextEncoderStep"),
+ ("vae_encoder.resize", "QwenImageEditResizeStep"),
+ ("vae_encoder.preprocess", "QwenImageEditInpaintProcessImagesInputStep"),
+ ("vae_encoder.encode", "QwenImageVaeEncoderStep"),
+ ("denoise.input.text_inputs", "QwenImageTextInputsStep"),
+ ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"),
+ ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"),
+ ("denoise.prepare_inpaint_latents.add_noise_to_latents", "QwenImagePrepareLatentsWithStrengthStep"),
+ ("denoise.prepare_inpaint_latents.create_mask_latents", "QwenImageCreateMaskLatentsStep"),
+ ("denoise.prepare_rope_inputs", "QwenImageEditRoPEInputsStep"),
+ ("denoise.denoise", "QwenImageEditInpaintDenoiseStep"),
+ ("denoise.after_denoise", "QwenImageAfterDenoiseStep"),
+ ("decode.decode", "QwenImageDecoderStep"),
+ ("decode.postprocess", "QwenImageInpaintProcessImagesOutputStep"),
+ ],
+}
+
+
class TestQwenImageEditModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin):
pipeline_class = QwenImageEditModularPipeline
pipeline_blocks_class = QwenImageEditAutoBlocks
@@ -63,6 +199,7 @@ class TestQwenImageEditModularPipelineFast(ModularPipelineTesterMixin, ModularGu
params = frozenset(["prompt", "height", "width", "negative_prompt", "attention_kwargs", "image", "mask_image"])
batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"])
+ expected_workflow_blocks = QWEN_IMAGE_EDIT_WORKFLOWS
def get_dummy_inputs(self):
generator = self.get_generator()
diff --git a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py
index 7b55933e4caf..f640f0ec83f2 100644
--- a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py
+++ b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py
@@ -267,6 +267,60 @@ def test_controlnet_cfg(self):
assert max_diff > 1e-2, "Output with CFG must be different from normal inference"
+TEXT2IMAGE_WORKFLOWS = {
+ "text2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"),
+ ("denoise.denoise", "StableDiffusionXLDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+ "controlnet_text2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"),
+ ("denoise.controlnet_input", "StableDiffusionXLControlNetInputStep"),
+ ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+ "controlnet_union_text2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"),
+ ("denoise.controlnet_input", "StableDiffusionXLControlNetUnionInputStep"),
+ ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+ "ip_adapter_text2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("ip_adapter", "StableDiffusionXLIPAdapterStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"),
+ ("denoise.denoise", "StableDiffusionXLDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+ "ip_adapter_controlnet_text2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("ip_adapter", "StableDiffusionXLIPAdapterStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"),
+ ("denoise.controlnet_input", "StableDiffusionXLControlNetInputStep"),
+ ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+}
+
+
class TestSDXLModularPipelineFast(
SDXLModularTesterMixin,
SDXLModularIPAdapterTesterMixin,
@@ -291,6 +345,8 @@ class TestSDXLModularPipelineFast(
batch_params = frozenset(["prompt", "negative_prompt"])
expected_image_output_shape = (1, 3, 64, 64)
+ expected_workflow_blocks = TEXT2IMAGE_WORKFLOWS
+
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
inputs = {
@@ -314,6 +370,65 @@ def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=3e-3)
+IMAGE2IMAGE_WORKFLOWS = {
+ "image2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("vae_encoder", "StableDiffusionXLVaeEncoderStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("denoise.denoise", "StableDiffusionXLDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+ "controlnet_image2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("vae_encoder", "StableDiffusionXLVaeEncoderStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("denoise.controlnet_input", "StableDiffusionXLControlNetInputStep"),
+ ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+ "controlnet_union_image2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("vae_encoder", "StableDiffusionXLVaeEncoderStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("denoise.controlnet_input", "StableDiffusionXLControlNetUnionInputStep"),
+ ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+ "ip_adapter_image2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("ip_adapter", "StableDiffusionXLIPAdapterStep"),
+ ("vae_encoder", "StableDiffusionXLVaeEncoderStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("denoise.denoise", "StableDiffusionXLDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+ "ip_adapter_controlnet_image2image": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("ip_adapter", "StableDiffusionXLIPAdapterStep"),
+ ("vae_encoder", "StableDiffusionXLVaeEncoderStep"),
+ ("denoise.input", "StableDiffusionXLInputStep"),
+ ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"),
+ ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("denoise.controlnet_input", "StableDiffusionXLControlNetInputStep"),
+ ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLDecodeStep"),
+ ],
+}
+
+
class TestSDXLImg2ImgModularPipelineFast(
SDXLModularTesterMixin,
SDXLModularIPAdapterTesterMixin,
@@ -338,6 +453,7 @@ class TestSDXLImg2ImgModularPipelineFast(
)
batch_params = frozenset(["prompt", "negative_prompt", "image"])
expected_image_output_shape = (1, 3, 64, 64)
+ expected_workflow_blocks = IMAGE2IMAGE_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
@@ -367,6 +483,65 @@ def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=3e-3)
+INPAINTING_WORKFLOWS = {
+ "inpainting": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"),
+ ("input", "StableDiffusionXLInputStep"),
+ ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"),
+ ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("denoise", "StableDiffusionXLInpaintDenoiseStep"),
+ ("decode", "StableDiffusionXLInpaintDecodeStep"),
+ ],
+ "controlnet_inpainting": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"),
+ ("input", "StableDiffusionXLInputStep"),
+ ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"),
+ ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("controlnet_input", "StableDiffusionXLControlNetInputStep"),
+ ("denoise", "StableDiffusionXLInpaintControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLInpaintDecodeStep"),
+ ],
+ "controlnet_union_inpainting": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"),
+ ("input", "StableDiffusionXLInputStep"),
+ ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"),
+ ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("controlnet_input", "StableDiffusionXLControlNetUnionInputStep"),
+ ("denoise", "StableDiffusionXLInpaintControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLInpaintDecodeStep"),
+ ],
+ "ip_adapter_inpainting": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("ip_adapter", "StableDiffusionXLIPAdapterStep"),
+ ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"),
+ ("input", "StableDiffusionXLInputStep"),
+ ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"),
+ ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("denoise", "StableDiffusionXLInpaintDenoiseStep"),
+ ("decode", "StableDiffusionXLInpaintDecodeStep"),
+ ],
+ "ip_adapter_controlnet_inpainting": [
+ ("text_encoder", "StableDiffusionXLTextEncoderStep"),
+ ("ip_adapter", "StableDiffusionXLIPAdapterStep"),
+ ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"),
+ ("input", "StableDiffusionXLInputStep"),
+ ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"),
+ ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"),
+ ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"),
+ ("controlnet_input", "StableDiffusionXLControlNetInputStep"),
+ ("denoise", "StableDiffusionXLInpaintControlNetDenoiseStep"),
+ ("decode", "StableDiffusionXLInpaintDecodeStep"),
+ ],
+}
+
+
class SDXLInpaintingModularPipelineFastTests(
SDXLModularTesterMixin,
SDXLModularIPAdapterTesterMixin,
@@ -392,6 +567,7 @@ class SDXLInpaintingModularPipelineFastTests(
)
batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"])
expected_image_output_shape = (1, 3, 64, 64)
+ expected_workflow_blocks = INPAINTING_WORKFLOWS
def get_dummy_inputs(self, device, seed=0):
generator = self.get_generator(seed)
diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py
index 9ee5c6c2ac80..e97b543ff85d 100644
--- a/tests/modular_pipelines/test_modular_pipelines_common.py
+++ b/tests/modular_pipelines/test_modular_pipelines_common.py
@@ -1,12 +1,12 @@
import gc
import tempfile
-from typing import Callable, Union
+from typing import Callable
import pytest
import torch
import diffusers
-from diffusers import ComponentsManager, ModularPipeline, ModularPipelineBlocks
+from diffusers import AutoModel, ComponentsManager, ModularPipeline, ModularPipelineBlocks
from diffusers.guiders import ClassifierFreeGuidance
from diffusers.modular_pipelines.modular_pipeline_utils import (
ComponentSpec,
@@ -46,7 +46,7 @@ def get_generator(self, seed=0):
return generator
@property
- def pipeline_class(self) -> Union[Callable, ModularPipeline]:
+ def pipeline_class(self) -> Callable | ModularPipeline:
raise NotImplementedError(
"You need to set the attribute `pipeline_class = ClassNameOfPipeline` in the child test class. "
"See existing pipeline tests for reference."
@@ -59,7 +59,7 @@ def pretrained_model_name_or_path(self) -> str:
)
@property
- def pipeline_blocks_class(self) -> Union[Callable, ModularPipelineBlocks]:
+ def pipeline_blocks_class(self) -> Callable | ModularPipelineBlocks:
raise NotImplementedError(
"You need to set the attribute `pipeline_blocks_class = ClassNameOfPipelineBlocks` in the child test class. "
"See existing pipeline tests for reference."
@@ -100,6 +100,14 @@ def batch_params(self) -> frozenset:
"See existing pipeline tests for reference."
)
+ @property
+ def expected_workflow_blocks(self) -> dict:
+ raise NotImplementedError(
+ "You need to set the attribute `expected_workflow_blocks` in the child test class. "
+ "`expected_workflow_blocks` is a dictionary that maps workflow names to list of block names. "
+ "See existing pipeline tests for reference."
+ )
+
def setup_method(self):
# clean up the VRAM before each test
torch.compiler.reset()
@@ -341,6 +349,34 @@ def test_save_from_pretrained(self):
assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
+ def test_workflow_map(self):
+ blocks = self.pipeline_blocks_class()
+ if blocks._workflow_map is None:
+ pytest.skip("Skipping test as _workflow_map is not set")
+
+ assert hasattr(self, "expected_workflow_blocks") and self.expected_workflow_blocks, (
+ "expected_workflow_blocks must be defined in the test class"
+ )
+
+ for workflow_name, expected_blocks in self.expected_workflow_blocks.items():
+ workflow_blocks = blocks.get_workflow(workflow_name)
+ actual_blocks = list(workflow_blocks.sub_blocks.items())
+
+ # Check that the number of blocks matches
+ assert len(actual_blocks) == len(expected_blocks), (
+ f"Workflow '{workflow_name}' has {len(actual_blocks)} blocks, expected {len(expected_blocks)}"
+ )
+
+ # Check that each block name and type matches
+ for i, ((actual_name, actual_block), (expected_name, expected_class_name)) in enumerate(
+ zip(actual_blocks, expected_blocks)
+ ):
+ assert actual_name == expected_name
+ assert actual_block.__class__.__name__ == expected_class_name, (
+ f"Workflow '{workflow_name}': block '{actual_name}' has type "
+ f"{actual_block.__class__.__name__}, expected {expected_class_name}"
+ )
+
class ModularGuiderTesterMixin:
def test_guider_cfg(self, expected_max_diff=1e-2):
@@ -598,3 +634,68 @@ def test_model_description_includes_block_count(self):
content = generate_modular_model_card_content(blocks)
assert "5-block architecture" in content["model_description"]
+
+
+class TestAutoModelLoadIdTagging:
+ def test_automodel_tags_load_id(self):
+ model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet")
+
+ assert hasattr(model, "_diffusers_load_id"), "Model should have _diffusers_load_id attribute"
+ assert model._diffusers_load_id != "null", "_diffusers_load_id should not be 'null'"
+
+ # Verify load_id contains the expected fields
+ load_id = model._diffusers_load_id
+ assert "hf-internal-testing/tiny-stable-diffusion-xl-pipe" in load_id
+ assert "unet" in load_id
+
+ def test_automodel_update_components(self):
+ pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+ pipe.load_components(torch_dtype=torch.float32)
+
+ auto_model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet")
+
+ pipe.update_components(unet=auto_model)
+
+ assert pipe.unet is auto_model
+
+ assert "unet" in pipe._component_specs
+ spec = pipe._component_specs["unet"]
+ assert spec.pretrained_model_name_or_path == "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
+ assert spec.subfolder == "unet"
+
+
+class TestLoadComponentsSkipBehavior:
+ def test_load_components_skips_already_loaded(self):
+ pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+ pipe.load_components(torch_dtype=torch.float32)
+
+ original_unet = pipe.unet
+
+ pipe.load_components()
+
+ # Verify that the unet is the same object (not reloaded)
+ assert pipe.unet is original_unet, "load_components should skip already loaded components"
+
+ def test_load_components_selective_loading(self):
+ pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+
+ pipe.load_components(names="unet", torch_dtype=torch.float32)
+
+ # Verify only requested component was loaded.
+ assert hasattr(pipe, "unet")
+ assert pipe.unet is not None
+ assert getattr(pipe, "vae", None) is None
+
+ def test_load_components_skips_invalid_pretrained_path(self):
+ pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+
+ pipe._component_specs["test_component"] = ComponentSpec(
+ name="test_component",
+ type_hint=torch.nn.Module,
+ pretrained_model_name_or_path=None,
+ default_creation_method="from_pretrained",
+ )
+ pipe.load_components(torch_dtype=torch.float32)
+
+ # Verify test_component was not loaded
+ assert not hasattr(pipe, "test_component") or pipe.test_component is None
diff --git a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py
index 29da18fce61b..ab45def3ef30 100644
--- a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py
+++ b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py
@@ -19,6 +19,30 @@
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
+ZIMAGE_WORKFLOWS = {
+ "text2image": [
+ ("text_encoder", "ZImageTextEncoderStep"),
+ ("denoise.input", "ZImageTextInputStep"),
+ ("denoise.prepare_latents", "ZImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "ZImageSetTimestepsStep"),
+ ("denoise.denoise", "ZImageDenoiseStep"),
+ ("decode", "ZImageVaeDecoderStep"),
+ ],
+ "image2image": [
+ ("text_encoder", "ZImageTextEncoderStep"),
+ ("vae_encoder", "ZImageVaeImageEncoderStep"),
+ ("denoise.input", "ZImageTextInputStep"),
+ ("denoise.additional_inputs", "ZImageAdditionalInputsStep"),
+ ("denoise.prepare_latents", "ZImagePrepareLatentsStep"),
+ ("denoise.set_timesteps", "ZImageSetTimestepsStep"),
+ ("denoise.set_timesteps_with_strength", "ZImageSetTimestepsWithStrengthStep"),
+ ("denoise.prepare_latents_with_image", "ZImagePrepareLatentswithImageStep"),
+ ("denoise.denoise", "ZImageDenoiseStep"),
+ ("decode", "ZImageVaeDecoderStep"),
+ ],
+}
+
+
class TestZImageModularPipelineFast(ModularPipelineTesterMixin):
pipeline_class = ZImageModularPipeline
pipeline_blocks_class = ZImageAutoBlocks
@@ -26,6 +50,7 @@ class TestZImageModularPipelineFast(ModularPipelineTesterMixin):
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
+ expected_workflow_blocks = ZIMAGE_WORKFLOWS
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
diff --git a/tests/others/test_check_copies.py b/tests/others/test_check_copies.py
index 4b6fa28eb9ac..134faf5ec6f0 100644
--- a/tests/others/test_check_copies.py
+++ b/tests/others/test_check_copies.py
@@ -41,7 +41,7 @@
\"""
prev_sample: torch.Tensor
- pred_original_sample: Optional[torch.Tensor] = None
+ pred_original_sample: torch.Tensor | None = None
"""
diff --git a/tests/others/test_outputs.py b/tests/others/test_outputs.py
index c8069e6916ed..90b8bfe9464e 100644
--- a/tests/others/test_outputs.py
+++ b/tests/others/test_outputs.py
@@ -1,7 +1,6 @@
import pickle as pkl
import unittest
from dataclasses import dataclass
-from typing import List, Union
import numpy as np
import PIL.Image
@@ -13,7 +12,7 @@
@dataclass
class CustomOutput(BaseOutput):
- images: Union[List[PIL.Image.Image], np.ndarray]
+ images: list[PIL.Image.Image] | np.ndarray
class ConfigTester(unittest.TestCase):
diff --git a/tests/pipelines/allegro/test_allegro.py b/tests/pipelines/allegro/test_allegro.py
index b2e588de0647..c126a94ce10e 100644
--- a/tests/pipelines/allegro/test_allegro.py
+++ b/tests/pipelines/allegro/test_allegro.py
@@ -158,6 +158,10 @@ def test_save_load_local(self):
def test_save_load_optional_components(self):
pass
+ @unittest.skip("Decoding without tiling is not yet implemented")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
+
def test_inference(self):
device = "cpu"
diff --git a/tests/pipelines/audioldm2/test_audioldm2.py b/tests/pipelines/audioldm2/test_audioldm2.py
index 5ccba1dabbfe..33f41bce8f16 100644
--- a/tests/pipelines/audioldm2/test_audioldm2.py
+++ b/tests/pipelines/audioldm2/test_audioldm2.py
@@ -282,6 +282,8 @@ def test_audioldm2_prompt_embeds(self):
text_inputs = text_inputs["input_ids"].to(torch_device)
clap_prompt_embeds = audioldm_pipe.text_encoder.get_text_features(text_inputs)
+ if hasattr(clap_prompt_embeds, "pooler_output"):
+ clap_prompt_embeds = clap_prompt_embeds.pooler_output
clap_prompt_embeds = clap_prompt_embeds[:, None, :]
text_inputs = audioldm_pipe.tokenizer_2(
@@ -341,6 +343,8 @@ def test_audioldm2_negative_prompt_embeds(self):
text_inputs = text_inputs["input_ids"].to(torch_device)
clap_prompt_embeds = audioldm_pipe.text_encoder.get_text_features(text_inputs)
+ if hasattr(clap_prompt_embeds, "pooler_output"):
+ clap_prompt_embeds = clap_prompt_embeds.pooler_output
clap_prompt_embeds = clap_prompt_embeds[:, None, :]
text_inputs = audioldm_pipe.tokenizer_2(
diff --git a/tests/pipelines/bria/test_pipeline_bria.py b/tests/pipelines/bria/test_pipeline_bria.py
index 844488e76f2e..dac9c428cfc9 100644
--- a/tests/pipelines/bria/test_pipeline_bria.py
+++ b/tests/pipelines/bria/test_pipeline_bria.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
from huggingface_hub import hf_hub_download
-from transformers import T5EncoderModel, T5TokenizerFast
+from transformers import AutoConfig, T5EncoderModel, T5TokenizerFast
from diffusers import (
AutoencoderKL,
@@ -89,7 +89,8 @@ def get_dummy_components(self):
scheduler = FlowMatchEulerDiscreteScheduler()
torch.manual_seed(0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = T5TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/chroma/test_pipeline_chroma.py b/tests/pipelines/chroma/test_pipeline_chroma.py
index 3edd58b75f82..6b856128dff0 100644
--- a/tests/pipelines/chroma/test_pipeline_chroma.py
+++ b/tests/pipelines/chroma/test_pipeline_chroma.py
@@ -2,7 +2,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKL, ChromaPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler
@@ -41,7 +41,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1):
)
torch.manual_seed(0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py
index 4ed1393037b9..8d991c42c749 100644
--- a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py
+++ b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py
@@ -3,7 +3,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKL, ChromaImg2ImgPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler
@@ -42,7 +42,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1):
)
torch.manual_seed(0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/chronoedit/test_chronoedit.py b/tests/pipelines/chronoedit/test_chronoedit.py
index 43e5b3159b1c..0b72f93eed3c 100644
--- a/tests/pipelines/chronoedit/test_chronoedit.py
+++ b/tests/pipelines/chronoedit/test_chronoedit.py
@@ -17,6 +17,7 @@
import torch
from PIL import Image
from transformers import (
+ AutoConfig,
AutoTokenizer,
CLIPImageProcessor,
CLIPVisionConfig,
@@ -71,7 +72,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
# TODO: impl FlowDPMSolverMultistepScheduler
scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
diff --git a/tests/pipelines/cogvideo/test_cogvideox.py b/tests/pipelines/cogvideo/test_cogvideox.py
index dca1725d8a74..e1797cfaac9c 100644
--- a/tests/pipelines/cogvideo/test_cogvideox.py
+++ b/tests/pipelines/cogvideo/test_cogvideox.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLCogVideoX, CogVideoXPipeline, CogVideoXTransformer3DModel, DDIMScheduler
@@ -117,7 +117,8 @@ def get_dummy_components(self, num_layers: int = 1):
torch.manual_seed(0)
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
@@ -235,6 +236,9 @@ def test_attention_slicing_forward_pass(
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
diff --git a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py
index 097e8df7b35f..30eab652c011 100644
--- a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py
+++ b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLCogVideoX, CogVideoXFunControlPipeline, CogVideoXTransformer3DModel, DDIMScheduler
@@ -104,7 +104,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
@@ -228,6 +229,9 @@ def test_attention_slicing_forward_pass(
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
diff --git a/tests/pipelines/cogvideo/test_cogvideox_image2video.py b/tests/pipelines/cogvideo/test_cogvideox_image2video.py
index 1dd5e2ae1405..c30cf56a933c 100644
--- a/tests/pipelines/cogvideo/test_cogvideox_image2video.py
+++ b/tests/pipelines/cogvideo/test_cogvideox_image2video.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel, DDIMScheduler
from diffusers.utils import load_image
@@ -113,7 +113,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
@@ -237,6 +238,9 @@ def test_attention_slicing_forward_pass(
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
diff --git a/tests/pipelines/cogvideo/test_cogvideox_video2video.py b/tests/pipelines/cogvideo/test_cogvideox_video2video.py
index 3a1da7c4e7f7..60424ad2a04e 100644
--- a/tests/pipelines/cogvideo/test_cogvideox_video2video.py
+++ b/tests/pipelines/cogvideo/test_cogvideox_video2video.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel, CogVideoXVideoToVideoPipeline, DDIMScheduler
@@ -99,7 +99,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/cogview3/test_cogview3plus.py b/tests/pipelines/cogview3/test_cogview3plus.py
index 819d4b952fc7..374cb6a2a295 100644
--- a/tests/pipelines/cogview3/test_cogview3plus.py
+++ b/tests/pipelines/cogview3/test_cogview3plus.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKL, CogVideoXDDIMScheduler, CogView3PlusPipeline, CogView3PlusTransformer2DModel
@@ -89,7 +89,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = CogVideoXDDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/cogview4/test_cogview4.py b/tests/pipelines/cogview4/test_cogview4.py
index a1f0fc7a715b..5f71b1b296d9 100644
--- a/tests/pipelines/cogview4/test_cogview4.py
+++ b/tests/pipelines/cogview4/test_cogview4.py
@@ -108,7 +108,7 @@ def get_dummy_inputs(self, device, seed=0):
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
"prompt": "dance monkey",
- "negative_prompt": "",
+ "negative_prompt": "bad",
"generator": generator,
"num_inference_steps": 2,
"guidance_scale": 6.0,
diff --git a/tests/pipelines/consisid/test_consisid.py b/tests/pipelines/consisid/test_consisid.py
index 4fd9e536cddc..b427eeea1d8c 100644
--- a/tests/pipelines/consisid/test_consisid.py
+++ b/tests/pipelines/consisid/test_consisid.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLCogVideoX, ConsisIDPipeline, ConsisIDTransformer3DModel, DDIMScheduler
from diffusers.utils import load_image
@@ -122,7 +122,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
@@ -248,6 +249,9 @@ def test_attention_slicing_forward_pass(
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux.py b/tests/pipelines/controlnet_flux/test_controlnet_flux.py
index 0895d9de3581..8607cd6944d9 100644
--- a/tests/pipelines/controlnet_flux/test_controlnet_flux.py
+++ b/tests/pipelines/controlnet_flux/test_controlnet_flux.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
from huggingface_hub import hf_hub_download
-from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from transformers import AutoConfig, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
from diffusers import (
AutoencoderKL,
@@ -97,7 +97,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = T5TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py
index 3d8378a5786d..a4749188dfd8 100644
--- a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py
+++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py
@@ -2,7 +2,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -13,9 +13,7 @@
)
from diffusers.utils.torch_utils import randn_tensor
-from ...testing_utils import (
- torch_device,
-)
+from ...testing_utils import torch_device
from ..test_pipelines_common import PipelineTesterMixin, check_qkv_fused_layers_exist
@@ -70,7 +68,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py
index 3ba475deb8a8..6eb560d90848 100644
--- a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py
+++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py
@@ -3,15 +3,7 @@
import numpy as np
import torch
-
-# torch_device, # {{ edit_1 }} Removed unused import
-from transformers import (
- AutoTokenizer,
- CLIPTextConfig,
- CLIPTextModel,
- CLIPTokenizer,
- T5EncoderModel,
-)
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -22,11 +14,7 @@
)
from diffusers.utils.torch_utils import randn_tensor
-from ...testing_utils import (
- enable_full_determinism,
- floats_tensor,
- torch_device,
-)
+from ...testing_utils import enable_full_determinism, floats_tensor, torch_device
from ..test_pipelines_common import PipelineTesterMixin
@@ -85,7 +73,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py b/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py
index bf31f2abcffb..034ef56b0fd3 100644
--- a/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py
+++ b/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, BertModel, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -96,7 +96,10 @@ def get_dummy_components(self):
scheduler = DDPMScheduler()
text_encoder = BertModel.from_pretrained("hf-internal-testing/tiny-random-BertModel")
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-BertModel")
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py
index 34c34b7a2ce7..072f9aa405d9 100644
--- a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py
+++ b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py
@@ -17,7 +17,14 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel
+from transformers import (
+ AutoConfig,
+ AutoTokenizer,
+ CLIPTextConfig,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ T5EncoderModel,
+)
from diffusers import (
AutoencoderKL,
@@ -28,10 +35,7 @@
from diffusers.models import SD3ControlNetModel
from diffusers.utils.torch_utils import randn_tensor
-from ...testing_utils import (
- enable_full_determinism,
- torch_device,
-)
+from ...testing_utils import enable_full_determinism, torch_device
from ..test_pipelines_common import PipelineTesterMixin
@@ -103,7 +107,8 @@ def get_dummy_components(self):
text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_3 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py
index 2b6cf8d1e8be..46b08cf1f00b 100644
--- a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py
+++ b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py
@@ -15,11 +15,17 @@
import gc
import unittest
-from typing import Optional
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel
+from transformers import (
+ AutoConfig,
+ AutoTokenizer,
+ CLIPTextConfig,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ T5EncoderModel,
+)
from diffusers import (
AutoencoderKL,
@@ -63,7 +69,7 @@ class StableDiffusion3ControlNetPipelineFastTests(unittest.TestCase, PipelineTes
test_group_offloading = True
def get_dummy_components(
- self, num_controlnet_layers: int = 3, qk_norm: Optional[str] = "rms_norm", use_dual_attention=False
+ self, num_controlnet_layers: int = 3, qk_norm: str | None = "rms_norm", use_dual_attention=False
):
torch.manual_seed(0)
transformer = SD3Transformer2DModel(
@@ -118,7 +124,8 @@ def get_dummy_components(
text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_3 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
diff --git a/tests/pipelines/cosmos/cosmos_guardrail.py b/tests/pipelines/cosmos/cosmos_guardrail.py
index c9ef597fdb36..e1d667608ce6 100644
--- a/tests/pipelines/cosmos/cosmos_guardrail.py
+++ b/tests/pipelines/cosmos/cosmos_guardrail.py
@@ -14,7 +14,6 @@
# ===== This file is an implementation of a dummy guardrail for the fast tests =====
-from typing import Union
import numpy as np
import torch
@@ -35,7 +34,7 @@ def check_text_safety(self, prompt: str) -> bool:
def check_video_safety(self, frames: np.ndarray) -> np.ndarray:
return frames
- def to(self, device: Union[str, torch.device] = None, dtype: torch.dtype = None):
+ def to(self, device: str | torch.device = None, dtype: torch.dtype = None):
module = super().to(device=device, dtype=dtype)
return module
diff --git a/tests/pipelines/cosmos/test_cosmos.py b/tests/pipelines/cosmos/test_cosmos.py
index 32eea9c98c2c..3aa92dca1677 100644
--- a/tests/pipelines/cosmos/test_cosmos.py
+++ b/tests/pipelines/cosmos/test_cosmos.py
@@ -20,7 +20,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLCosmos, CosmosTextToWorldPipeline, CosmosTransformer3DModel, EDMEulerScheduler
@@ -107,7 +107,8 @@ def get_dummy_components(self):
rho=7.0,
final_sigmas_type="sigma_min",
)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
@@ -232,6 +233,9 @@ def test_attention_slicing_forward_pass(
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
diff --git a/tests/pipelines/cosmos/test_cosmos2_5_transfer.py b/tests/pipelines/cosmos/test_cosmos2_5_transfer.py
new file mode 100644
index 000000000000..c098a46a8f8a
--- /dev/null
+++ b/tests/pipelines/cosmos/test_cosmos2_5_transfer.py
@@ -0,0 +1,440 @@
+# Copyright 2025 The HuggingFace Team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import json
+import os
+import tempfile
+import unittest
+
+import numpy as np
+import torch
+from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer
+
+from diffusers import (
+ AutoencoderKLWan,
+ Cosmos2_5_TransferPipeline,
+ CosmosControlNetModel,
+ CosmosTransformer3DModel,
+ UniPCMultistepScheduler,
+)
+
+from ...testing_utils import enable_full_determinism, torch_device
+from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
+from ..test_pipelines_common import PipelineTesterMixin, to_np
+from .cosmos_guardrail import DummyCosmosSafetyChecker
+
+
+enable_full_determinism()
+
+
+class Cosmos2_5_TransferWrapper(Cosmos2_5_TransferPipeline):
+ @staticmethod
+ def from_pretrained(*args, **kwargs):
+ if "safety_checker" not in kwargs or kwargs["safety_checker"] is None:
+ safety_checker = DummyCosmosSafetyChecker()
+ device_map = kwargs.get("device_map", "cpu")
+ torch_dtype = kwargs.get("torch_dtype")
+ if device_map is not None or torch_dtype is not None:
+ safety_checker = safety_checker.to(device_map, dtype=torch_dtype)
+ kwargs["safety_checker"] = safety_checker
+ return Cosmos2_5_TransferPipeline.from_pretrained(*args, **kwargs)
+
+
+class Cosmos2_5_TransferPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+ pipeline_class = Cosmos2_5_TransferWrapper
+ params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
+ batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"controls"})
+ image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
+ image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
+ required_optional_params = frozenset(
+ [
+ "num_inference_steps",
+ "generator",
+ "latents",
+ "return_dict",
+ "callback_on_step_end",
+ "callback_on_step_end_tensor_inputs",
+ ]
+ )
+ supports_dduf = False
+ test_xformers_attention = False
+ test_layerwise_casting = True
+ test_group_offloading = True
+
+ def get_dummy_components(self):
+ torch.manual_seed(0)
+ # Transformer with img_context support for Transfer2.5
+ transformer = CosmosTransformer3DModel(
+ in_channels=16 + 1,
+ out_channels=16,
+ num_attention_heads=2,
+ attention_head_dim=16,
+ num_layers=2,
+ mlp_ratio=2,
+ text_embed_dim=32,
+ adaln_lora_dim=4,
+ max_size=(4, 32, 32),
+ patch_size=(1, 2, 2),
+ rope_scale=(2.0, 1.0, 1.0),
+ concat_padding_mask=True,
+ extra_pos_embed_type="learnable",
+ controlnet_block_every_n=1,
+ img_context_dim_in=32,
+ img_context_num_tokens=4,
+ img_context_dim_out=32,
+ )
+
+ torch.manual_seed(0)
+ controlnet = CosmosControlNetModel(
+ n_controlnet_blocks=2,
+ in_channels=16 + 1 + 1, # control latent channels + condition_mask + padding_mask
+ latent_channels=16 + 1 + 1, # base latent channels (16) + condition_mask (1) + padding_mask (1) = 18
+ model_channels=32,
+ num_attention_heads=2,
+ attention_head_dim=16,
+ mlp_ratio=2,
+ text_embed_dim=32,
+ adaln_lora_dim=4,
+ patch_size=(1, 2, 2),
+ max_size=(4, 32, 32),
+ rope_scale=(2.0, 1.0, 1.0),
+ extra_pos_embed_type="learnable", # Match transformer's config
+ img_context_dim_in=32,
+ img_context_dim_out=32,
+ use_crossattn_projection=False, # Test doesn't need this projection
+ )
+
+ torch.manual_seed(0)
+ vae = AutoencoderKLWan(
+ base_dim=3,
+ z_dim=16,
+ dim_mult=[1, 1, 1, 1],
+ num_res_blocks=1,
+ temperal_downsample=[False, True, True],
+ )
+
+ torch.manual_seed(0)
+ scheduler = UniPCMultistepScheduler()
+
+ torch.manual_seed(0)
+ config = Qwen2_5_VLConfig(
+ text_config={
+ "hidden_size": 16,
+ "intermediate_size": 16,
+ "num_hidden_layers": 2,
+ "num_attention_heads": 2,
+ "num_key_value_heads": 2,
+ "rope_scaling": {
+ "mrope_section": [1, 1, 2],
+ "rope_type": "default",
+ "type": "default",
+ },
+ "rope_theta": 1000000.0,
+ },
+ vision_config={
+ "depth": 2,
+ "hidden_size": 16,
+ "intermediate_size": 16,
+ "num_heads": 2,
+ "out_hidden_size": 16,
+ },
+ hidden_size=16,
+ vocab_size=152064,
+ vision_end_token_id=151653,
+ vision_start_token_id=151652,
+ vision_token_id=151654,
+ )
+ text_encoder = Qwen2_5_VLForConditionalGeneration(config)
+ tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration")
+
+ components = {
+ "transformer": transformer,
+ "controlnet": controlnet,
+ "vae": vae,
+ "scheduler": scheduler,
+ "text_encoder": text_encoder,
+ "tokenizer": tokenizer,
+ "safety_checker": DummyCosmosSafetyChecker(),
+ }
+ return components
+
+ def get_dummy_inputs(self, device, seed=0):
+ if str(device).startswith("mps"):
+ generator = torch.manual_seed(seed)
+ else:
+ generator = torch.Generator(device=device).manual_seed(seed)
+
+ controls_generator = torch.Generator(device="cpu").manual_seed(seed)
+
+ inputs = {
+ "prompt": "dance monkey",
+ "negative_prompt": "bad quality",
+ "controls": [torch.randn(3, 32, 32, generator=controls_generator) for _ in range(5)],
+ "generator": generator,
+ "num_inference_steps": 2,
+ "guidance_scale": 3.0,
+ "height": 32,
+ "width": 32,
+ "num_frames": 3,
+ "num_frames_per_chunk": 16,
+ "max_sequence_length": 16,
+ "output_type": "pt",
+ }
+
+ return inputs
+
+ def test_components_function(self):
+ init_components = self.get_dummy_components()
+ init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float))}
+ pipe = self.pipeline_class(**init_components)
+ self.assertTrue(hasattr(pipe, "components"))
+ self.assertTrue(set(pipe.components.keys()) == set(init_components.keys()))
+
+ def test_inference(self):
+ device = "cpu"
+
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ pipe.to(device)
+ pipe.set_progress_bar_config(disable=None)
+
+ inputs = self.get_dummy_inputs(device)
+ video = pipe(**inputs).frames
+ generated_video = video[0]
+ self.assertEqual(generated_video.shape, (3, 3, 32, 32))
+ self.assertTrue(torch.isfinite(generated_video).all())
+
+ def test_inference_autoregressive_multi_chunk(self):
+ device = "cpu"
+
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ pipe.to(device)
+ pipe.set_progress_bar_config(disable=None)
+
+ inputs = self.get_dummy_inputs(device)
+ inputs["num_frames"] = 5
+ inputs["num_frames_per_chunk"] = 3
+ inputs["num_ar_conditional_frames"] = 1
+
+ video = pipe(**inputs).frames
+ generated_video = video[0]
+ self.assertEqual(generated_video.shape, (5, 3, 32, 32))
+ self.assertTrue(torch.isfinite(generated_video).all())
+
+ def test_inference_autoregressive_multi_chunk_no_condition_frames(self):
+ device = "cpu"
+
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ pipe.to(device)
+ pipe.set_progress_bar_config(disable=None)
+
+ inputs = self.get_dummy_inputs(device)
+ inputs["num_frames"] = 5
+ inputs["num_frames_per_chunk"] = 3
+ inputs["num_ar_conditional_frames"] = 0
+
+ video = pipe(**inputs).frames
+ generated_video = video[0]
+ self.assertEqual(generated_video.shape, (5, 3, 32, 32))
+ self.assertTrue(torch.isfinite(generated_video).all())
+
+ def test_num_frames_per_chunk_above_rope_raises(self):
+ device = "cpu"
+
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ pipe.to(device)
+ pipe.set_progress_bar_config(disable=None)
+
+ inputs = self.get_dummy_inputs(device)
+ inputs["num_frames_per_chunk"] = 17
+
+ with self.assertRaisesRegex(ValueError, "too large for RoPE setting"):
+ pipe(**inputs)
+
+ def test_inference_with_controls(self):
+ """Test inference with control inputs (ControlNet)."""
+ device = "cpu"
+
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ pipe.to(device)
+ pipe.set_progress_bar_config(disable=None)
+
+ inputs = self.get_dummy_inputs(device)
+ inputs["controls"] = [torch.randn(3, 32, 32) for _ in range(5)] # list of 5 frames (C, H, W)
+ inputs["controls_conditioning_scale"] = 1.0
+ inputs["num_frames"] = None
+
+ video = pipe(**inputs).frames
+ generated_video = video[0]
+ self.assertEqual(generated_video.shape, (5, 3, 32, 32))
+ self.assertTrue(torch.isfinite(generated_video).all())
+
+ def test_callback_inputs(self):
+ sig = inspect.signature(self.pipeline_class.__call__)
+ has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
+ has_callback_step_end = "callback_on_step_end" in sig.parameters
+
+ if not (has_callback_tensor_inputs and has_callback_step_end):
+ return
+
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ pipe = pipe.to(torch_device)
+ pipe.set_progress_bar_config(disable=None)
+ self.assertTrue(
+ hasattr(pipe, "_callback_tensor_inputs"),
+ f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
+ )
+
+ def callback_inputs_subset(pipe, i, t, callback_kwargs):
+ for tensor_name in callback_kwargs.keys():
+ assert tensor_name in pipe._callback_tensor_inputs
+ return callback_kwargs
+
+ def callback_inputs_all(pipe, i, t, callback_kwargs):
+ for tensor_name in pipe._callback_tensor_inputs:
+ assert tensor_name in callback_kwargs
+ for tensor_name in callback_kwargs.keys():
+ assert tensor_name in pipe._callback_tensor_inputs
+ return callback_kwargs
+
+ inputs = self.get_dummy_inputs(torch_device)
+
+ inputs["callback_on_step_end"] = callback_inputs_subset
+ inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
+ _ = pipe(**inputs)[0]
+
+ inputs["callback_on_step_end"] = callback_inputs_all
+ inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
+ _ = pipe(**inputs)[0]
+
+ def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
+ is_last = i == (pipe.num_timesteps - 1)
+ if is_last:
+ callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
+ return callback_kwargs
+
+ inputs["callback_on_step_end"] = callback_inputs_change_tensor
+ inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
+ output = pipe(**inputs)[0]
+ assert output.abs().sum() < 1e10
+
+ def test_inference_batch_single_identical(self):
+ self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=1e-2)
+
+ def test_attention_slicing_forward_pass(
+ self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
+ ):
+ if not getattr(self, "test_attention_slicing", True):
+ return
+
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ for component in pipe.components.values():
+ if hasattr(component, "set_default_attn_processor"):
+ component.set_default_attn_processor()
+ pipe.to(torch_device)
+ pipe.set_progress_bar_config(disable=None)
+
+ generator_device = "cpu"
+ inputs = self.get_dummy_inputs(generator_device)
+ output_without_slicing = pipe(**inputs)[0]
+
+ pipe.enable_attention_slicing(slice_size=1)
+ inputs = self.get_dummy_inputs(generator_device)
+ output_with_slicing1 = pipe(**inputs)[0]
+
+ pipe.enable_attention_slicing(slice_size=2)
+ inputs = self.get_dummy_inputs(generator_device)
+ output_with_slicing2 = pipe(**inputs)[0]
+
+ if test_max_difference:
+ max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
+ max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
+ self.assertLess(
+ max(max_diff1, max_diff2),
+ expected_max_diff,
+ "Attention slicing should not affect the inference results",
+ )
+
+ def test_serialization_with_variants(self):
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ model_components = [
+ component_name
+ for component_name, component in pipe.components.items()
+ if isinstance(component, torch.nn.Module)
+ ]
+ # Remove components that aren't saved as standard diffusers models
+ if "safety_checker" in model_components:
+ model_components.remove("safety_checker")
+ variant = "fp16"
+
+ with tempfile.TemporaryDirectory() as tmpdir:
+ pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False)
+
+ with open(f"{tmpdir}/model_index.json", "r") as f:
+ config = json.load(f)
+
+ for subfolder in os.listdir(tmpdir):
+ if not os.path.isfile(subfolder) and subfolder in model_components:
+ folder_path = os.path.join(tmpdir, subfolder)
+ is_folder = os.path.isdir(folder_path) and subfolder in config
+ assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path))
+
+ def test_torch_dtype_dict(self):
+ components = self.get_dummy_components()
+ if not components:
+ self.skipTest("No dummy components defined.")
+
+ pipe = self.pipeline_class(**components)
+
+ specified_key = next(iter(components.keys()))
+
+ with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
+ pipe.save_pretrained(tmpdirname, safe_serialization=False)
+ torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16}
+ loaded_pipe = self.pipeline_class.from_pretrained(
+ tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict
+ )
+
+ for name, component in loaded_pipe.components.items():
+ # Skip components that are not loaded from disk or have special handling
+ if name == "safety_checker":
+ continue
+ if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"):
+ expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32))
+ self.assertEqual(
+ component.dtype,
+ expected_dtype,
+ f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}",
+ )
+
+ def test_save_load_optional_components(self, expected_max_difference=1e-4):
+ self.pipeline_class._optional_components.remove("safety_checker")
+ super().test_save_load_optional_components(expected_max_difference=expected_max_difference)
+ self.pipeline_class._optional_components.append("safety_checker")
+
+ @unittest.skip(
+ "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in "
+ "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is "
+ "too large and slow to run on CI."
+ )
+ def test_encode_prompt_works_in_isolation(self):
+ pass
diff --git a/tests/pipelines/cosmos/test_cosmos2_text2image.py b/tests/pipelines/cosmos/test_cosmos2_text2image.py
index 8e3c5e4c29f4..71c61eff0054 100644
--- a/tests/pipelines/cosmos/test_cosmos2_text2image.py
+++ b/tests/pipelines/cosmos/test_cosmos2_text2image.py
@@ -20,7 +20,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKLWan,
@@ -95,7 +95,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/cosmos/test_cosmos2_video2world.py b/tests/pipelines/cosmos/test_cosmos2_video2world.py
index b0ca0e160d98..1b814257a30a 100644
--- a/tests/pipelines/cosmos/test_cosmos2_video2world.py
+++ b/tests/pipelines/cosmos/test_cosmos2_video2world.py
@@ -21,7 +21,7 @@
import numpy as np
import PIL.Image
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKLWan,
@@ -96,7 +96,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/cosmos/test_cosmos_video2world.py b/tests/pipelines/cosmos/test_cosmos_video2world.py
index 2633c2007ac2..925c3b448677 100644
--- a/tests/pipelines/cosmos/test_cosmos_video2world.py
+++ b/tests/pipelines/cosmos/test_cosmos_video2world.py
@@ -21,7 +21,7 @@
import numpy as np
import PIL.Image
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLCosmos, CosmosTransformer3DModel, CosmosVideoToWorldPipeline, EDMEulerScheduler
@@ -108,7 +108,8 @@ def get_dummy_components(self):
rho=7.0,
final_sigmas_type="sigma_min",
)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
@@ -245,6 +246,9 @@ def test_attention_slicing_forward_pass(
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
diff --git a/tests/pipelines/deepfloyd_if/__init__.py b/tests/pipelines/deepfloyd_if/__init__.py
index d47374b07e22..d8f72bd2ca5b 100644
--- a/tests/pipelines/deepfloyd_if/__init__.py
+++ b/tests/pipelines/deepfloyd_if/__init__.py
@@ -2,7 +2,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import DDPMScheduler, UNet2DConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
@@ -18,7 +18,8 @@
class IFPipelineTesterMixin:
def _get_dummy_components(self):
torch.manual_seed(0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
torch.manual_seed(0)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
@@ -75,7 +76,8 @@ def _get_dummy_components(self):
def _get_superresolution_dummy_components(self):
torch.manual_seed(0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
torch.manual_seed(0)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
@@ -250,6 +252,9 @@ def _test_save_load_optional_components(self):
# This should be handled in the base test and then this method can be removed.
def _test_save_load_local(self):
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
diff --git a/tests/pipelines/deepfloyd_if/test_if.py b/tests/pipelines/deepfloyd_if/test_if.py
index e1870ddcbae9..0fd1391decd0 100644
--- a/tests/pipelines/deepfloyd_if/test_if.py
+++ b/tests/pipelines/deepfloyd_if/test_if.py
@@ -18,9 +18,7 @@
import torch
-from diffusers import (
- IFPipeline,
-)
+from diffusers import IFPipeline
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
diff --git a/tests/pipelines/flux/test_pipeline_flux.py b/tests/pipelines/flux/test_pipeline_flux.py
index f7476a21de57..13336f0cde9b 100644
--- a/tests/pipelines/flux/test_pipeline_flux.py
+++ b/tests/pipelines/flux/test_pipeline_flux.py
@@ -4,7 +4,7 @@
import numpy as np
import torch
from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -93,7 +93,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/flux/test_pipeline_flux_control.py b/tests/pipelines/flux/test_pipeline_flux_control.py
index 7e966470a336..44efca9b9f0e 100644
--- a/tests/pipelines/flux/test_pipeline_flux_control.py
+++ b/tests/pipelines/flux/test_pipeline_flux_control.py
@@ -3,7 +3,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxControlPipeline, FluxTransformer2DModel
@@ -53,7 +53,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/flux/test_pipeline_flux_control_img2img.py b/tests/pipelines/flux/test_pipeline_flux_control_img2img.py
index e56136f2e91b..0f0bc0934115 100644
--- a/tests/pipelines/flux/test_pipeline_flux_control_img2img.py
+++ b/tests/pipelines/flux/test_pipeline_flux_control_img2img.py
@@ -3,7 +3,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -57,7 +57,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py
index e42c5fc2aab5..ae2b6b829e54 100644
--- a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py
+++ b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py
@@ -3,7 +3,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -58,7 +58,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/flux/test_pipeline_flux_fill.py b/tests/pipelines/flux/test_pipeline_flux_fill.py
index 25a4a3354820..42cd1efad495 100644
--- a/tests/pipelines/flux/test_pipeline_flux_fill.py
+++ b/tests/pipelines/flux/test_pipeline_flux_fill.py
@@ -3,7 +3,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxFillPipeline, FluxTransformer2DModel
@@ -58,7 +58,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/flux/test_pipeline_flux_img2img.py b/tests/pipelines/flux/test_pipeline_flux_img2img.py
index 6f435760aef5..00587905d337 100644
--- a/tests/pipelines/flux/test_pipeline_flux_img2img.py
+++ b/tests/pipelines/flux/test_pipeline_flux_img2img.py
@@ -3,7 +3,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxImg2ImgPipeline, FluxTransformer2DModel
@@ -55,7 +55,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/flux/test_pipeline_flux_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_inpaint.py
index 6324ff236e10..14edb9e441b5 100644
--- a/tests/pipelines/flux/test_pipeline_flux_inpaint.py
+++ b/tests/pipelines/flux/test_pipeline_flux_inpaint.py
@@ -3,7 +3,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxInpaintPipeline, FluxTransformer2DModel
@@ -55,7 +55,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext.py b/tests/pipelines/flux/test_pipeline_flux_kontext.py
index 5c78964ea54f..1c018f14b522 100644
--- a/tests/pipelines/flux/test_pipeline_flux_kontext.py
+++ b/tests/pipelines/flux/test_pipeline_flux_kontext.py
@@ -3,7 +3,7 @@
import numpy as np
import PIL.Image
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -79,7 +79,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py
index 9a2e32056dcb..b5f8570ebd1a 100644
--- a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py
+++ b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py
@@ -3,7 +3,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -79,7 +79,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/glm_image/test_glm_image.py b/tests/pipelines/glm_image/test_glm_image.py
index d907d082d275..f17776554b50 100644
--- a/tests/pipelines/glm_image/test_glm_image.py
+++ b/tests/pipelines/glm_image/test_glm_image.py
@@ -16,7 +16,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, GlmImagePipeline, GlmImageTransformer2DModel
from diffusers.utils import is_transformers_version
@@ -57,7 +57,8 @@ class GlmImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
def get_dummy_components(self):
torch.manual_seed(0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
glm_config = GlmImageConfig(
@@ -281,6 +282,86 @@ def test_batch_with_num_images_per_prompt(self):
# Should return 4 images (2 prompts × 2 images per prompt)
self.assertEqual(len(images), 4)
+ def test_prompt_with_prior_token_ids(self):
+ """Test that prompt and prior_token_ids can be provided together.
+
+ When both are given, the AR generation step is skipped (prior_token_ids is used
+ directly) and prompt is used to generate prompt_embeds via the glyph encoder.
+ """
+ device = "cpu"
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ pipe.to(device)
+ pipe.set_progress_bar_config(disable=None)
+
+ height, width = 32, 32
+
+ # Step 1: Run with prompt only to get prior_token_ids from AR model
+ generator = torch.Generator(device=device).manual_seed(0)
+ prior_token_ids, _, _ = pipe.generate_prior_tokens(
+ prompt="A photo of a cat",
+ height=height,
+ width=width,
+ device=torch.device(device),
+ generator=torch.Generator(device=device).manual_seed(0),
+ )
+
+ # Step 2: Run with both prompt and prior_token_ids — should not raise
+ generator = torch.Generator(device=device).manual_seed(0)
+ inputs_both = {
+ "prompt": "A photo of a cat",
+ "prior_token_ids": prior_token_ids,
+ "generator": generator,
+ "num_inference_steps": 2,
+ "guidance_scale": 1.5,
+ "height": height,
+ "width": width,
+ "max_sequence_length": 16,
+ "output_type": "pt",
+ }
+ images = pipe(**inputs_both).images
+ self.assertEqual(len(images), 1)
+ self.assertEqual(images[0].shape, (3, 32, 32))
+
+ def test_check_inputs_rejects_invalid_combinations(self):
+ """Test that check_inputs correctly rejects invalid input combinations."""
+ device = "cpu"
+ components = self.get_dummy_components()
+ pipe = self.pipeline_class(**components)
+ pipe.to(device)
+
+ height, width = 32, 32
+
+ # Neither prompt nor prior_token_ids → error
+ with self.assertRaises(ValueError):
+ pipe.check_inputs(
+ prompt=None,
+ height=height,
+ width=width,
+ callback_on_step_end_tensor_inputs=None,
+ prompt_embeds=torch.randn(1, 16, 32),
+ )
+
+ # prior_token_ids alone without prompt or prompt_embeds → error
+ with self.assertRaises(ValueError):
+ pipe.check_inputs(
+ prompt=None,
+ height=height,
+ width=width,
+ callback_on_step_end_tensor_inputs=None,
+ prior_token_ids=torch.randint(0, 100, (1, 64)),
+ )
+
+ # prompt + prompt_embeds together → error
+ with self.assertRaises(ValueError):
+ pipe.check_inputs(
+ prompt="A cat",
+ height=height,
+ width=width,
+ callback_on_step_end_tensor_inputs=None,
+ prompt_embeds=torch.randn(1, 16, 32),
+ )
+
@unittest.skip("Needs to be revisited.")
def test_encode_prompt_works_in_isolation(self):
pass
diff --git a/tests/pipelines/hidream_image/test_pipeline_hidream.py b/tests/pipelines/hidream_image/test_pipeline_hidream.py
index ddf39ba4c1e6..1d51872d5526 100644
--- a/tests/pipelines/hidream_image/test_pipeline_hidream.py
+++ b/tests/pipelines/hidream_image/test_pipeline_hidream.py
@@ -18,6 +18,7 @@
import numpy as np
import torch
from transformers import (
+ AutoConfig,
AutoTokenizer,
CLIPTextConfig,
CLIPTextModelWithProjection,
@@ -94,7 +95,8 @@ def get_dummy_components(self):
text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_3 = T5EncoderModel(config)
torch.manual_seed(0)
text_encoder_4 = LlamaForCausalLM.from_pretrained("hf-internal-testing/tiny-random-LlamaForCausalLM")
@@ -149,12 +151,12 @@ def test_inference(self):
self.assertEqual(generated_image.shape, (128, 128, 3))
# fmt: off
- expected_slice = np.array([0.4507, 0.5256, 0.4205, 0.5791, 0.4848, 0.4831, 0.4443, 0.5107, 0.6586, 0.3163, 0.7318, 0.5933, 0.6252, 0.5512, 0.5357, 0.5983])
+ expected_slice = np.array([0.4501, 0.5256, 0.4207, 0.5783, 0.4842, 0.4833, 0.4441, 0.5112, 0.6587, 0.3169, 0.7308, 0.5927, 0.6251, 0.5509, 0.5355, 0.5969])
# fmt: on
generated_slice = generated_image.flatten()
generated_slice = np.concatenate([generated_slice[:8], generated_slice[-8:]])
- self.assertTrue(np.allclose(generated_slice, expected_slice, atol=1e-3))
+ self.assertTrue(np.allclose(generated_slice, expected_slice, atol=5e-3))
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=3e-4)
diff --git a/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py b/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py
index e4b2c686b8b1..7cb6b0913bf6 100644
--- a/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py
+++ b/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py
@@ -34,11 +34,7 @@
)
from ...testing_utils import enable_full_determinism
-from ..test_pipelines_common import (
- FirstBlockCacheTesterMixin,
- PipelineTesterMixin,
- to_np,
-)
+from ..test_pipelines_common import FirstBlockCacheTesterMixin, PipelineTesterMixin, to_np
enable_full_determinism()
@@ -227,7 +223,7 @@ def test_inference_guider(self):
self.assertEqual(generated_image.shape, (3, 16, 16))
expected_slice_np = np.array(
- [0.61494756, 0.49616697, 0.60327923, 0.6115793, 0.49047345, 0.56977504, 0.53066164, 0.58880305, 0.5570612]
+ [0.6068114, 0.48716035, 0.5984431, 0.60241306, 0.48849544, 0.5624479, 0.53696984, 0.58964247, 0.54248774]
)
output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy()
diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py b/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py
index 27b5bde31050..2a28e5e42f7d 100644
--- a/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py
+++ b/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py
@@ -233,7 +233,7 @@ def test_inference(self):
self.assertEqual(generated_video.shape, (5, 3, 16, 16))
# fmt: off
- expected_slice = torch.tensor([0.444, 0.479, 0.4485, 0.5752, 0.3539, 0.1548, 0.2706, 0.3593, 0.5323, 0.6635, 0.6795, 0.5255, 0.5091, 0.345, 0.4276, 0.4128])
+ expected_slice = torch.tensor([0.4441, 0.4790, 0.4485, 0.5748, 0.3539, 0.1553, 0.2707, 0.3594, 0.5331, 0.6645, 0.6799, 0.5257, 0.5092, 0.3450, 0.4276, 0.4127])
# fmt: on
generated_slice = generated_video.flatten()
@@ -371,13 +371,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py b/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py
index 7ebe797febfa..0ac15d4c405a 100644
--- a/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py
+++ b/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py
@@ -333,13 +333,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video.py b/tests/pipelines/hunyuan_video/test_hunyuan_video.py
index 57a6daebad1f..719bb9e67fbf 100644
--- a/tests/pipelines/hunyuan_video/test_hunyuan_video.py
+++ b/tests/pipelines/hunyuan_video/test_hunyuan_video.py
@@ -346,13 +346,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py
index 51c258b15c38..f7f8908b6f8f 100644
--- a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py
+++ b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py
@@ -392,13 +392,13 @@ def test_sequential_offload_forward_pass_twice(self):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py b/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py
index 993c7ef6e4bb..de20148105bf 100644
--- a/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py
+++ b/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py
@@ -15,7 +15,14 @@
import unittest
import torch
-from transformers import ByT5Tokenizer, Qwen2_5_VLTextConfig, Qwen2_5_VLTextModel, Qwen2Tokenizer, T5EncoderModel
+from transformers import (
+ AutoConfig,
+ ByT5Tokenizer,
+ Qwen2_5_VLTextConfig,
+ Qwen2_5_VLTextModel,
+ Qwen2Tokenizer,
+ T5EncoderModel,
+)
from diffusers import (
AutoencoderKLHunyuanVideo15,
@@ -114,7 +121,8 @@ def get_dummy_components(self, num_layers: int = 1):
tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration")
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer_2 = ByT5Tokenizer()
guider = ClassifierFreeGuidance(guidance_scale=1.0)
diff --git a/tests/pipelines/hunyuandit/test_hunyuan_dit.py b/tests/pipelines/hunyuandit/test_hunyuan_dit.py
index 2a329f10bc80..ba57b6a3599a 100644
--- a/tests/pipelines/hunyuandit/test_hunyuan_dit.py
+++ b/tests/pipelines/hunyuandit/test_hunyuan_dit.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, BertModel, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel
from diffusers import AutoencoderKL, DDPMScheduler, HunyuanDiT2DModel, HunyuanDiTPipeline
@@ -74,7 +74,9 @@ def get_dummy_components(self):
scheduler = DDPMScheduler()
text_encoder = BertModel.from_pretrained("hf-internal-testing/tiny-random-BertModel")
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-BertModel")
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/kandinsky/test_kandinsky_combined.py b/tests/pipelines/kandinsky/test_kandinsky_combined.py
index eba897659700..fdb36b433a94 100644
--- a/tests/pipelines/kandinsky/test_kandinsky_combined.py
+++ b/tests/pipelines/kandinsky/test_kandinsky_combined.py
@@ -34,9 +34,7 @@
class KandinskyPipelineCombinedFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class = KandinskyCombinedPipeline
- params = [
- "prompt",
- ]
+ params = ["prompt"]
batch_params = ["prompt", "negative_prompt"]
required_optional_params = [
"generator",
@@ -148,6 +146,10 @@ def test_float16_inference(self):
def test_dict_tuple_outputs_equivalent(self):
super().test_dict_tuple_outputs_equivalent(expected_max_difference=5e-4)
+ @unittest.skip("Test not supported.")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
+
class KandinskyPipelineImg2ImgCombinedFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class = KandinskyImg2ImgCombinedPipeline
@@ -264,6 +266,10 @@ def test_dict_tuple_outputs_equivalent(self):
def test_save_load_optional_components(self):
super().test_save_load_optional_components(expected_max_difference=5e-4)
+ @unittest.skip("Test not supported.")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
+
class KandinskyPipelineInpaintCombinedFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class = KandinskyInpaintCombinedPipeline
@@ -384,3 +390,7 @@ def test_save_load_optional_components(self):
def test_save_load_local(self):
super().test_save_load_local(expected_max_difference=5e-3)
+
+ @unittest.skip("Test not supported.")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py b/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py
index 62f5853da9a5..30aa2896089b 100644
--- a/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py
+++ b/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py
@@ -36,9 +36,7 @@
class KandinskyV22PipelineCombinedFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class = KandinskyV22CombinedPipeline
- params = [
- "prompt",
- ]
+ params = ["prompt"]
batch_params = ["prompt", "negative_prompt"]
required_optional_params = [
"generator",
@@ -70,12 +68,7 @@ def get_dummy_components(self):
def get_dummy_inputs(self, device, seed=0):
prior_dummy = PriorDummies()
inputs = prior_dummy.get_dummy_inputs(device=device, seed=seed)
- inputs.update(
- {
- "height": 64,
- "width": 64,
- }
- )
+ inputs.update({"height": 64, "width": 64})
return inputs
def test_kandinsky(self):
@@ -155,12 +148,18 @@ def test_save_load_local(self):
def test_save_load_optional_components(self):
super().test_save_load_optional_components(expected_max_difference=5e-3)
+ @unittest.skip("Test not supported.")
def test_callback_inputs(self):
pass
+ @unittest.skip("Test not supported.")
def test_callback_cfg(self):
pass
+ @unittest.skip("Test not supported.")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
+
class KandinskyV22PipelineImg2ImgCombinedFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class = KandinskyV22Img2ImgCombinedPipeline
@@ -279,12 +278,18 @@ def test_save_load_optional_components(self):
def save_load_local(self):
super().test_save_load_local(expected_max_difference=5e-3)
+ @unittest.skip("Test not supported.")
def test_callback_inputs(self):
pass
+ @unittest.skip("Test not supported.")
def test_callback_cfg(self):
pass
+ @unittest.skip("Test not supported.")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
+
class KandinskyV22PipelineInpaintCombinedFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class = KandinskyV22InpaintCombinedPipeline
@@ -411,3 +416,7 @@ def test_callback_inputs(self):
def test_callback_cfg(self):
pass
+
+ @unittest.skip("`device_map` is not yet supported for connected pipelines.")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py b/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py
index df1dd2d9872c..5659699fc7aa 100644
--- a/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py
+++ b/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py
@@ -296,6 +296,9 @@ def callback_inputs_test(pipe, i, t, callback_kwargs):
output = pipe(**inputs)[0]
assert output.abs().sum() == 0
+ def test_pipeline_with_accelerator_device_map(self):
+ super().test_pipeline_with_accelerator_device_map(expected_max_difference=5e-3)
+
@slow
@require_torch_accelerator
diff --git a/tests/pipelines/kandinsky3/test_kandinsky3.py b/tests/pipelines/kandinsky3/test_kandinsky3.py
index 55500f729bbb..7deae194fb4b 100644
--- a/tests/pipelines/kandinsky3/test_kandinsky3.py
+++ b/tests/pipelines/kandinsky3/test_kandinsky3.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoPipelineForImage2Image,
@@ -108,7 +108,8 @@ def get_dummy_components(self, time_cond_proj_dim=None):
torch.manual_seed(0)
movq = self.dummy_movq
torch.manual_seed(0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config).eval()
torch.manual_seed(0)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
@@ -155,9 +156,9 @@ def test_kandinsky3(self):
assert image.shape == (1, 16, 16, 3)
- expected_slice = np.array([0.3768, 0.4373, 0.4865, 0.4890, 0.4299, 0.5122, 0.4921, 0.4924, 0.5599])
+ expected_slice = np.array([0.3944, 0.3680, 0.4842, 0.5333, 0.4412, 0.4812, 0.5089, 0.5381, 0.5578])
- assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, (
+ assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1, (
f" expected_slice {expected_slice}, but got {image_slice.flatten()}"
)
diff --git a/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py b/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py
index d3bfa4b3082c..469300e67788 100644
--- a/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py
+++ b/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py
@@ -20,7 +20,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoPipelineForImage2Image,
@@ -119,7 +119,8 @@ def get_dummy_components(self, time_cond_proj_dim=None):
torch.manual_seed(0)
movq = self.dummy_movq
torch.manual_seed(0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config).eval()
torch.manual_seed(0)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
@@ -155,10 +156,7 @@ def get_dummy_inputs(self, device, seed=0):
return inputs
def test_dict_tuple_outputs_equivalent(self):
- expected_slice = None
- if torch_device == "cpu":
- expected_slice = np.array([0.5762, 0.6112, 0.4150, 0.6018, 0.6167, 0.4626, 0.5426, 0.5641, 0.6536])
- super().test_dict_tuple_outputs_equivalent(expected_slice=expected_slice)
+ super().test_dict_tuple_outputs_equivalent()
def test_kandinsky3_img2img(self):
device = "cpu"
@@ -177,11 +175,9 @@ def test_kandinsky3_img2img(self):
assert image.shape == (1, 64, 64, 3)
- expected_slice = np.array(
- [0.576259, 0.6132097, 0.41703486, 0.603196, 0.62062526, 0.4655338, 0.5434324, 0.5660727, 0.65433365]
- )
+ expected_slice = np.array([0.5725, 0.6248, 0.4355, 0.5732, 0.6105, 0.5267, 0.5470, 0.5512, 0.6618])
- assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, (
+ assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1, (
f" expected_slice {expected_slice}, but got {image_slice.flatten()}"
)
@@ -194,6 +190,9 @@ def test_inference_batch_single_identical(self):
def test_save_load_dduf(self):
super().test_save_load_dduf(atol=1e-3, rtol=1e-3)
+ def test_pipeline_with_accelerator_device_map(self):
+ super().test_pipeline_with_accelerator_device_map(expected_max_difference=5e-3)
+
@slow
@require_torch_accelerator
diff --git a/tests/pipelines/latte/test_latte.py b/tests/pipelines/latte/test_latte.py
index a40d4bf8eede..873c06e11c5b 100644
--- a/tests/pipelines/latte/test_latte.py
+++ b/tests/pipelines/latte/test_latte.py
@@ -20,7 +20,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -109,7 +109,8 @@ def get_dummy_components(self, num_layers: int = 1):
vae = AutoencoderKL()
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/ltx/test_ltx.py b/tests/pipelines/ltx/test_ltx.py
index aaf4161b51fb..9836551d30a1 100644
--- a/tests/pipelines/ltx/test_ltx.py
+++ b/tests/pipelines/ltx/test_ltx.py
@@ -17,7 +17,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel
@@ -88,7 +88,8 @@ def get_dummy_components(self, num_layers: int = 1):
torch.manual_seed(0)
scheduler = FlowMatchEulerDiscreteScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/ltx/test_ltx_condition.py b/tests/pipelines/ltx/test_ltx_condition.py
index f5dfb0186209..b469662241fc 100644
--- a/tests/pipelines/ltx/test_ltx_condition.py
+++ b/tests/pipelines/ltx/test_ltx_condition.py
@@ -17,7 +17,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKLLTXVideo,
@@ -92,7 +92,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = FlowMatchEulerDiscreteScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/ltx/test_ltx_image2video.py b/tests/pipelines/ltx/test_ltx_image2video.py
index 2702993d4a59..7407c8bef5ea 100644
--- a/tests/pipelines/ltx/test_ltx_image2video.py
+++ b/tests/pipelines/ltx/test_ltx_image2video.py
@@ -17,7 +17,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKLLTXVideo,
@@ -91,7 +91,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = FlowMatchEulerDiscreteScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/mochi/test_mochi.py b/tests/pipelines/mochi/test_mochi.py
index 5615720a9343..df6f33f79cb2 100644
--- a/tests/pipelines/mochi/test_mochi.py
+++ b/tests/pipelines/mochi/test_mochi.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel
@@ -89,7 +89,8 @@ def get_dummy_components(self, num_layers: int = 2):
torch.manual_seed(0)
scheduler = FlowMatchEulerDiscreteScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
@@ -207,6 +208,9 @@ def test_attention_slicing_forward_pass(
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
diff --git a/tests/pipelines/pag/test_pag_hunyuan_dit.py b/tests/pipelines/pag/test_pag_hunyuan_dit.py
index f268a614f85c..38686ee448de 100644
--- a/tests/pipelines/pag/test_pag_hunyuan_dit.py
+++ b/tests/pipelines/pag/test_pag_hunyuan_dit.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, BertModel, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -67,7 +67,9 @@ def get_dummy_components(self):
scheduler = DDPMScheduler()
text_encoder = BertModel.from_pretrained("hf-internal-testing/tiny-random-BertModel")
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-BertModel")
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
diff --git a/tests/pipelines/pag/test_pag_pixart_sigma.py b/tests/pipelines/pag/test_pag_pixart_sigma.py
index c04ebad08fdc..9bc2f6eed395 100644
--- a/tests/pipelines/pag/test_pag_pixart_sigma.py
+++ b/tests/pipelines/pag/test_pag_pixart_sigma.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
import diffusers
from diffusers import (
@@ -80,7 +80,8 @@ def get_dummy_components(self):
vae = AutoencoderKL()
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/pag/test_pag_sana.py b/tests/pipelines/pag/test_pag_sana.py
index 5408595c729d..366267dc092b 100644
--- a/tests/pipelines/pag/test_pag_sana.py
+++ b/tests/pipelines/pag/test_pag_sana.py
@@ -325,13 +325,13 @@ def test_pag_applied_layers(self):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py
index 26e6ca099286..7f755ea8e170 100644
--- a/tests/pipelines/pag/test_pag_sd3.py
+++ b/tests/pipelines/pag/test_pag_sd3.py
@@ -3,7 +3,14 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel
+from transformers import (
+ AutoConfig,
+ AutoTokenizer,
+ CLIPTextConfig,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ T5EncoderModel,
+)
from diffusers import (
AutoencoderKL,
@@ -73,7 +80,9 @@ def get_dummy_components(self):
torch.manual_seed(0)
text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config)
- text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_3 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
diff --git a/tests/pipelines/pag/test_pag_sd3_img2img.py b/tests/pipelines/pag/test_pag_sd3_img2img.py
index 19a36e283de4..e4146b87803c 100644
--- a/tests/pipelines/pag/test_pag_sd3_img2img.py
+++ b/tests/pipelines/pag/test_pag_sd3_img2img.py
@@ -5,7 +5,14 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel
+from transformers import (
+ AutoConfig,
+ AutoTokenizer,
+ CLIPTextConfig,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ T5EncoderModel,
+)
from diffusers import (
AutoencoderKL,
@@ -84,7 +91,9 @@ def get_dummy_components(self):
torch.manual_seed(0)
text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config)
- text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_3 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
diff --git a/tests/pipelines/pixart_alpha/test_pixart.py b/tests/pipelines/pixart_alpha/test_pixart.py
index fd41c9887dcc..037a9f44f31e 100644
--- a/tests/pipelines/pixart_alpha/test_pixart.py
+++ b/tests/pipelines/pixart_alpha/test_pixart.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -77,7 +77,10 @@ def get_dummy_components(self):
vae = AutoencoderKL()
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/pixart_sigma/test_pixart.py b/tests/pipelines/pixart_sigma/test_pixart.py
index 6e8535062a79..51eebadd0ed0 100644
--- a/tests/pipelines/pixart_sigma/test_pixart.py
+++ b/tests/pipelines/pixart_sigma/test_pixart.py
@@ -19,7 +19,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
@@ -83,7 +83,10 @@ def get_dummy_components(self):
vae = AutoencoderKL()
scheduler = DDIMScheduler()
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
diff --git a/tests/pipelines/prx/test_pipeline_prx.py b/tests/pipelines/prx/test_pipeline_prx.py
index 46c6a5760e22..20c2727b3156 100644
--- a/tests/pipelines/prx/test_pipeline_prx.py
+++ b/tests/pipelines/prx/test_pipeline_prx.py
@@ -1,7 +1,6 @@
import unittest
import numpy as np
-import pytest
import torch
from transformers import AutoTokenizer
from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig, T5GemmaModuleConfig
@@ -11,17 +10,11 @@
from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel
from diffusers.pipelines.prx.pipeline_prx import PRXPipeline
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import is_transformers_version
from ..pipeline_params import TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
-@pytest.mark.xfail(
- condition=is_transformers_version(">", "4.57.1"),
- reason="See https://github.com/huggingface/diffusers/pull/12456#issuecomment-3424228544",
- strict=False,
-)
class PRXPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class = PRXPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
@@ -99,7 +92,7 @@ def get_dummy_components(self):
}
encoder_config = T5GemmaModuleConfig(**encoder_params)
text_encoder_config = T5GemmaConfig(encoder=encoder_config, is_encoder_decoder=False, **encoder_params)
- text_encoder = T5GemmaEncoder(text_encoder_config)
+ text_encoder = T5GemmaEncoder(text_encoder_config.encoder)
return {
"transformer": transformer,
@@ -263,3 +256,27 @@ def test_inference_with_autoencoder_dc(self):
expected_image = torch.zeros(3, 32, 32)
max_diff = np.abs(generated_image - expected_image).max()
self.assertLessEqual(max_diff, 1e10)
+
+ @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
+ def test_save_load_dduf(self):
+ pass
+
+ @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
+ def test_loading_with_variants(self):
+ pass
+
+ @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
+
+ @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
+ def test_save_load_local(self):
+ pass
+
+ @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
+ def test_save_load_optional_components(self):
+ pass
+
+ @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
+ def test_torch_dtype_dict(self):
+ pass
diff --git a/tests/pipelines/qwenimage/test_qwenimage.py b/tests/pipelines/qwenimage/test_qwenimage.py
index 8ebfe7d08bc1..ff53d1c234c7 100644
--- a/tests/pipelines/qwenimage/test_qwenimage.py
+++ b/tests/pipelines/qwenimage/test_qwenimage.py
@@ -113,7 +113,7 @@ def get_dummy_components(self):
vision_start_token_id=151652,
vision_token_id=151654,
)
- text_encoder = Qwen2_5_VLForConditionalGeneration(config)
+ text_encoder = Qwen2_5_VLForConditionalGeneration(config).eval()
tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration")
components = {
@@ -160,12 +160,12 @@ def test_inference(self):
self.assertEqual(generated_image.shape, (3, 32, 32))
# fmt: off
- expected_slice = torch.tensor([0.56331, 0.63677, 0.6015, 0.56369, 0.58166, 0.55277, 0.57176, 0.63261, 0.41466, 0.35561, 0.56229, 0.48334, 0.49714, 0.52622, 0.40872, 0.50208])
+ expected_slice = torch.tensor([0.5633, 0.6368, 0.6015, 0.5637, 0.5817, 0.5528, 0.5718, 0.6326, 0.4147, 0.3556, 0.5623, 0.4833, 0.4971, 0.5262, 0.4087, 0.5021])
# fmt: on
generated_slice = generated_image.flatten()
generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
- self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3))
+ self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3))
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1)
diff --git a/tests/pipelines/qwenimage/test_qwenimage_controlnet.py b/tests/pipelines/qwenimage/test_qwenimage_controlnet.py
index 188106b49b84..59a2dd497184 100644
--- a/tests/pipelines/qwenimage/test_qwenimage_controlnet.py
+++ b/tests/pipelines/qwenimage/test_qwenimage_controlnet.py
@@ -211,7 +211,7 @@ def test_qwen_controlnet(self):
generated_slice = generated_image.flatten()
generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
- self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3))
+ self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3))
def test_qwen_controlnet_multicondition(self):
device = "cpu"
@@ -255,7 +255,7 @@ def test_qwen_controlnet_multicondition(self):
generated_slice = generated_image.flatten()
generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
- self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3))
+ self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3))
def test_attention_slicing_forward_pass(
self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit.py b/tests/pipelines/qwenimage/test_qwenimage_edit.py
index 058548cf5f1b..8184c895f825 100644
--- a/tests/pipelines/qwenimage/test_qwenimage_edit.py
+++ b/tests/pipelines/qwenimage/test_qwenimage_edit.py
@@ -115,7 +115,7 @@ def get_dummy_components(self):
vision_start_token_id=151652,
vision_token_id=151654,
)
- text_encoder = Qwen2_5_VLForConditionalGeneration(config)
+ text_encoder = Qwen2_5_VLForConditionalGeneration(config).eval()
tokenizer = Qwen2Tokenizer.from_pretrained(tiny_ckpt_id)
components = {
@@ -163,12 +163,12 @@ def test_inference(self):
self.assertEqual(generated_image.shape, (3, 32, 32))
# fmt: off
- expected_slice = torch.tensor([[0.5637, 0.6341, 0.6001, 0.5620, 0.5794, 0.5498, 0.5757, 0.6389, 0.4174, 0.3597, 0.5649, 0.4894, 0.4969, 0.5255, 0.4083, 0.4986]])
+ expected_slice = torch.tensor([0.5637, 0.6341, 0.6001, 0.5620, 0.5794, 0.5498, 0.5757, 0.6389, 0.4174, 0.3597, 0.5649, 0.4894, 0.4969, 0.5255, 0.4083, 0.4986])
# fmt: on
generated_slice = generated_image.flatten()
generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
- self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3))
+ self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3))
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1)
diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py
index 6faf34728286..e8bc694ced84 100644
--- a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py
+++ b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py
@@ -164,7 +164,7 @@ def test_inference(self):
self.assertEqual(generated_image.shape, (3, 32, 32))
# fmt: off
- expected_slice = torch.tensor([[0.5637, 0.6341, 0.6001, 0.5620, 0.5794, 0.5498, 0.5757, 0.6389, 0.4174, 0.3597, 0.5649, 0.4894, 0.4969, 0.5255, 0.4083, 0.4986]])
+ expected_slice = torch.tensor([0.5640, 0.6339, 0.5997, 0.5607, 0.5799, 0.5496, 0.5760, 0.6393, 0.4172, 0.3595, 0.5655, 0.4896, 0.4971, 0.5255, 0.4088, 0.4987])
# fmt: on
generated_slice = generated_image.flatten()
diff --git a/tests/pipelines/sana/test_sana.py b/tests/pipelines/sana/test_sana.py
index f23303c966e5..fc533dd7bf5c 100644
--- a/tests/pipelines/sana/test_sana.py
+++ b/tests/pipelines/sana/test_sana.py
@@ -290,13 +290,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/sana/test_sana_controlnet.py b/tests/pipelines/sana/test_sana_controlnet.py
index df14d935edf5..0f432a998ddc 100644
--- a/tests/pipelines/sana/test_sana_controlnet.py
+++ b/tests/pipelines/sana/test_sana_controlnet.py
@@ -309,13 +309,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/sana/test_sana_sprint.py b/tests/pipelines/sana/test_sana_sprint.py
index 0d45205ea8c7..9f88d0f344ae 100644
--- a/tests/pipelines/sana/test_sana_sprint.py
+++ b/tests/pipelines/sana/test_sana_sprint.py
@@ -283,13 +283,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/sana/test_sana_sprint_img2img.py b/tests/pipelines/sana/test_sana_sprint_img2img.py
index 5de5c7f44606..312e2b2b8080 100644
--- a/tests/pipelines/sana/test_sana_sprint_img2img.py
+++ b/tests/pipelines/sana/test_sana_sprint_img2img.py
@@ -295,13 +295,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/sana_video/test_sana_video.py b/tests/pipelines/sana_video/test_sana_video.py
index 9f360a942a64..dd897a848be8 100644
--- a/tests/pipelines/sana_video/test_sana_video.py
+++ b/tests/pipelines/sana_video/test_sana_video.py
@@ -185,13 +185,13 @@ def test_save_load_local(self, expected_max_difference=5e-4):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/sana_video/test_sana_video_i2v.py b/tests/pipelines/sana_video/test_sana_video_i2v.py
index 36a646ca528f..2232968eb2ca 100644
--- a/tests/pipelines/sana_video/test_sana_video_i2v.py
+++ b/tests/pipelines/sana_video/test_sana_video_i2v.py
@@ -196,13 +196,13 @@ def test_save_load_local(self, expected_max_difference=5e-4):
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
- "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
+ "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
diff --git a/tests/pipelines/skyreels_v2/test_skyreels_v2.py b/tests/pipelines/skyreels_v2/test_skyreels_v2.py
index 1bcec877c30d..b6adb3cc1f2c 100644
--- a/tests/pipelines/skyreels_v2/test_skyreels_v2.py
+++ b/tests/pipelines/skyreels_v2/test_skyreels_v2.py
@@ -18,20 +18,11 @@
import torch
from transformers import AutoTokenizer, T5EncoderModel
-from diffusers import (
- AutoencoderKLWan,
- SkyReelsV2Pipeline,
- SkyReelsV2Transformer3DModel,
- UniPCMultistepScheduler,
-)
-
-from ...testing_utils import (
- enable_full_determinism,
-)
+from diffusers import AutoencoderKLWan, SkyReelsV2Pipeline, SkyReelsV2Transformer3DModel, UniPCMultistepScheduler
+
+from ...testing_utils import enable_full_determinism
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
-from ..test_pipelines_common import (
- PipelineTesterMixin,
-)
+from ..test_pipelines_common import PipelineTesterMixin
enable_full_determinism()
diff --git a/tests/pipelines/skyreels_v2/test_skyreels_v2_df.py b/tests/pipelines/skyreels_v2/test_skyreels_v2_df.py
index 74235d59efd6..213d085a6dca 100644
--- a/tests/pipelines/skyreels_v2/test_skyreels_v2_df.py
+++ b/tests/pipelines/skyreels_v2/test_skyreels_v2_df.py
@@ -25,13 +25,9 @@
UniPCMultistepScheduler,
)
-from ...testing_utils import (
- enable_full_determinism,
-)
+from ...testing_utils import enable_full_determinism
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
-from ..test_pipelines_common import (
- PipelineTesterMixin,
-)
+from ..test_pipelines_common import PipelineTesterMixin
enable_full_determinism()
diff --git a/tests/pipelines/skyreels_v2/test_skyreels_v2_df_image_to_video.py b/tests/pipelines/skyreels_v2/test_skyreels_v2_df_image_to_video.py
index f0cbc710df05..02daa61395c3 100644
--- a/tests/pipelines/skyreels_v2/test_skyreels_v2_df_image_to_video.py
+++ b/tests/pipelines/skyreels_v2/test_skyreels_v2_df_image_to_video.py
@@ -17,10 +17,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import (
- AutoTokenizer,
- T5EncoderModel,
-)
+from transformers import AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKLWan,
diff --git a/tests/pipelines/skyreels_v2/test_skyreels_v2_df_video_to_video.py b/tests/pipelines/skyreels_v2/test_skyreels_v2_df_video_to_video.py
index 1b0b23318e63..cbfb86f7696e 100644
--- a/tests/pipelines/skyreels_v2/test_skyreels_v2_df_video_to_video.py
+++ b/tests/pipelines/skyreels_v2/test_skyreels_v2_df_video_to_video.py
@@ -27,14 +27,9 @@
UniPCMultistepScheduler,
)
-from ...testing_utils import (
- enable_full_determinism,
- torch_device,
-)
+from ...testing_utils import enable_full_determinism, torch_device
from ..pipeline_params import TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
-from ..test_pipelines_common import (
- PipelineTesterMixin,
-)
+from ..test_pipelines_common import PipelineTesterMixin
enable_full_determinism()
diff --git a/tests/pipelines/stable_audio/test_stable_audio.py b/tests/pipelines/stable_audio/test_stable_audio.py
index dd03f4d07f07..492aa92252de 100644
--- a/tests/pipelines/stable_audio/test_stable_audio.py
+++ b/tests/pipelines/stable_audio/test_stable_audio.py
@@ -19,10 +19,7 @@
import numpy as np
import torch
-from transformers import (
- T5EncoderModel,
- T5Tokenizer,
-)
+from transformers import AutoConfig, T5EncoderModel, T5Tokenizer
from diffusers import (
AutoencoderOobleck,
@@ -111,7 +108,8 @@ def get_dummy_components(self):
)
torch.manual_seed(0)
t5_repo_id = "hf-internal-testing/tiny-random-T5ForConditionalGeneration"
- text_encoder = T5EncoderModel.from_pretrained(t5_repo_id)
+ config = AutoConfig.from_pretrained(t5_repo_id)
+ text_encoder = T5EncoderModel(config)
tokenizer = T5Tokenizer.from_pretrained(t5_repo_id, truncation=True, model_max_length=25)
torch.manual_seed(0)
diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py
index 3ccefe3de35d..200c832d0941 100644
--- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py
+++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py
@@ -3,7 +3,14 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel
+from transformers import (
+ AutoConfig,
+ AutoTokenizer,
+ CLIPTextConfig,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ T5EncoderModel,
+)
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Pipeline
@@ -72,7 +79,9 @@ def get_dummy_components(self):
torch.manual_seed(0)
text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config)
- text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_3 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py
index 9025b1060c9e..3f46b341a09e 100644
--- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py
+++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py
@@ -4,7 +4,14 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel
+from transformers import (
+ AutoConfig,
+ AutoTokenizer,
+ CLIPTextConfig,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ T5EncoderModel,
+)
from diffusers import (
AutoencoderKL,
@@ -73,7 +80,9 @@ def get_dummy_components(self):
torch.manual_seed(0)
text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config)
- text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_3 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py
index 628930340294..a90ca21a801b 100644
--- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py
+++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py
@@ -3,7 +3,14 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel
+from transformers import (
+ AutoConfig,
+ AutoTokenizer,
+ CLIPTextConfig,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ T5EncoderModel,
+)
from diffusers import (
AutoencoderKL,
@@ -73,7 +80,9 @@ def get_dummy_components(self):
torch.manual_seed(0)
text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config)
- text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ torch.manual_seed(0)
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_3 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py
index b3818e5fe4cc..af3573ce84cb 100644
--- a/tests/pipelines/test_pipelines_common.py
+++ b/tests/pipelines/test_pipelines_common.py
@@ -5,7 +5,7 @@
import tempfile
import unittest
import uuid
-from typing import Any, Callable, Dict, Union
+from typing import Any, Callable, Dict
import numpy as np
import PIL.Image
@@ -1072,7 +1072,7 @@ def get_generator(self, seed):
return generator
@property
- def pipeline_class(self) -> Union[Callable, DiffusionPipeline]:
+ def pipeline_class(self) -> Callable | DiffusionPipeline:
raise NotImplementedError(
"You need to set the attribute `pipeline_class = ClassNameOfPipeline` in the child test class. "
"See existing pipeline tests for reference."
@@ -1157,6 +1157,9 @@ def tearDown(self):
def test_save_load_local(self, expected_max_difference=5e-4):
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
@@ -1295,6 +1298,9 @@ def _test_inference_batch_single_identical(
additional_params_copy_to_batched_inputs=["num_inference_steps"],
):
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for components in pipe.components.values():
if hasattr(components, "set_default_attn_processor"):
@@ -1345,6 +1351,9 @@ def _test_inference_batch_single_identical(
def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4):
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
@@ -1477,6 +1486,9 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
if not self.pipeline_class._optional_components:
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
@@ -1557,6 +1569,9 @@ def _test_attention_slicing_forward_pass(
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
@@ -2065,7 +2080,16 @@ def is_nan(tensor):
for component_name in model_components_pipe:
pipe_component = model_components_pipe[component_name]
pipe_loaded_component = model_components_pipe_loaded[component_name]
- for p1, p2 in zip(pipe_component.parameters(), pipe_loaded_component.parameters()):
+
+ model_loaded_params = dict(pipe_loaded_component.named_parameters())
+ model_original_params = dict(pipe_component.named_parameters())
+
+ for name, p1 in model_original_params.items():
+ # Skip tied weights that aren't saved with variants (transformers v5 behavior)
+ if name not in model_loaded_params:
+ continue
+
+ p2 = model_loaded_params[name]
# nan check for luminanext (mps).
if not (is_nan(p1) and is_nan(p2)):
self.assertTrue(torch.equal(p1, p2))
@@ -2089,6 +2113,9 @@ def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
# We initialize the pipeline with only text encoders and tokenizers,
# mimicking a real-world scenario.
@@ -2220,6 +2247,9 @@ def test_save_load_dduf(self, atol=1e-4, rtol=1e-4):
from huggingface_hub import export_folder_as_dduf
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
@@ -2355,9 +2385,13 @@ def test_torch_dtype_dict(self):
f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}",
)
- @require_torch_accelerator
def test_pipeline_with_accelerator_device_map(self, expected_max_difference=1e-4):
components = self.get_dummy_components()
+ # Set text encoders to eval mode to match from_pretrained behavior
+ # This ensures deterministic outputs when models are loaded with device_map
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
@@ -2406,7 +2440,11 @@ def test_pipeline_level_group_offloading_sanity_checks(self):
if name not in [exclude_module_name] and isinstance(component, torch.nn.Module):
# `component.device` prints the `onload_device` type. We should probably override the
# `device` property in `ModelMixin`.
- component_device = next(component.parameters())[0].device
+ # Skip modules with no parameters (e.g., dummy safety checkers with only buffers)
+ params = list(component.parameters())
+ if not params:
+ continue
+ component_device = params[0].device
self.assertTrue(torch.device(component_device).type == torch.device(offload_device).type)
@require_torch_accelerator
@@ -2677,6 +2715,9 @@ def test_pyramid_attention_broadcast_inference(self, expected_atol: float = 0.2)
device = "cpu" # ensure determinism for the device-dependent torch.Generator
num_layers = 2
components = self.get_dummy_components(num_layers=num_layers)
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
diff --git a/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py b/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py
index 00ae0441fe99..bfbcd141d9b3 100644
--- a/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py
+++ b/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py
@@ -5,7 +5,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
import diffusers
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel, VisualClozePipeline
@@ -77,7 +77,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
@@ -215,7 +216,7 @@ def test_different_task_prompts(self, expected_min_diff=1e-1):
def test_callback_cfg(self):
pass
- def test_save_load_local(self, expected_max_difference=5e-4):
+ def test_save_load_local(self, expected_max_difference=1e-3):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
@@ -260,6 +261,9 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
if not hasattr(self.pipeline_class, "_optional_components"):
return
components = self.get_dummy_components()
+ for key in components:
+ if "text_encoder" in key and hasattr(components[key], "eval"):
+ components[key].eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
@@ -342,3 +346,7 @@ def test_save_load_float16(self, expected_max_diff=1e-2):
self.assertLess(
max_diff, expected_max_diff, "The output of the fp16 pipeline changed after saving and loading."
)
+
+ @unittest.skip("Test not supported.")
+ def test_pipeline_with_accelerator_device_map(self):
+ pass
diff --git a/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py b/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py
index ab6b3ca5c587..5640f2f634d3 100644
--- a/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py
+++ b/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py
@@ -5,7 +5,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
import diffusers
from diffusers import (
@@ -79,7 +79,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(clip_text_encoder_config)
torch.manual_seed(0)
- text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder_2 = T5EncoderModel(config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
@@ -310,3 +311,7 @@ def test_save_load_float16(self, expected_max_diff=1e-2):
@unittest.skip("Skipped due to missing layout_prompt. Needs further investigation.")
def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=0.0001, rtol=0.0001):
pass
+
+ @unittest.skip("Needs to be revisited later.")
+ def test_pipeline_with_accelerator_device_map(self, expected_max_difference=0.0001):
+ pass
diff --git a/tests/pipelines/wan/test_wan.py b/tests/pipelines/wan/test_wan.py
index 106a7b294646..958e1b8c8eaf 100644
--- a/tests/pipelines/wan/test_wan.py
+++ b/tests/pipelines/wan/test_wan.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanPipeline, WanTransformer3DModel
@@ -68,7 +68,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
# TODO: impl FlowDPMSolverMultistepScheduler
scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
diff --git a/tests/pipelines/wan/test_wan_22.py b/tests/pipelines/wan/test_wan_22.py
index 56ef5ceb97ed..fd17ca414af4 100644
--- a/tests/pipelines/wan/test_wan_22.py
+++ b/tests/pipelines/wan/test_wan_22.py
@@ -17,14 +17,11 @@
import numpy as np
import torch
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanPipeline, WanTransformer3DModel
-from ...testing_utils import (
- enable_full_determinism,
- torch_device,
-)
+from ...testing_utils import enable_full_determinism, torch_device
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
@@ -63,7 +60,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
@@ -235,7 +233,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
diff --git a/tests/pipelines/wan/test_wan_22_image_to_video.py b/tests/pipelines/wan/test_wan_22_image_to_video.py
index 6294d62044f3..4634047ebb73 100644
--- a/tests/pipelines/wan/test_wan_22_image_to_video.py
+++ b/tests/pipelines/wan/test_wan_22_image_to_video.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanImageToVideoPipeline, WanTransformer3DModel
@@ -64,7 +64,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
@@ -248,7 +249,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
diff --git a/tests/pipelines/wan/test_wan_animate.py b/tests/pipelines/wan/test_wan_animate.py
index d6d1b09f3620..5d634fb71849 100644
--- a/tests/pipelines/wan/test_wan_animate.py
+++ b/tests/pipelines/wan/test_wan_animate.py
@@ -19,6 +19,7 @@
import torch
from PIL import Image
from transformers import (
+ AutoConfig,
AutoTokenizer,
CLIPImageProcessor,
CLIPVisionConfig,
@@ -78,7 +79,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
diff --git a/tests/pipelines/wan/test_wan_image_to_video.py b/tests/pipelines/wan/test_wan_image_to_video.py
index 07a9142f2553..7ed263abdcb5 100644
--- a/tests/pipelines/wan/test_wan_image_to_video.py
+++ b/tests/pipelines/wan/test_wan_image_to_video.py
@@ -19,6 +19,7 @@
import torch
from PIL import Image
from transformers import (
+ AutoConfig,
AutoTokenizer,
CLIPImageProcessor,
CLIPVisionConfig,
@@ -68,7 +69,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
# TODO: impl FlowDPMSolverMultistepScheduler
scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
@@ -239,7 +241,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
# TODO: impl FlowDPMSolverMultistepScheduler
scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
diff --git a/tests/pipelines/wan/test_wan_vace.py b/tests/pipelines/wan/test_wan_vace.py
index fe078c0deb8a..53becce1685d 100644
--- a/tests/pipelines/wan/test_wan_vace.py
+++ b/tests/pipelines/wan/test_wan_vace.py
@@ -18,7 +18,7 @@
import numpy as np
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKLWan,
@@ -67,7 +67,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
diff --git a/tests/pipelines/wan/test_wan_video_to_video.py b/tests/pipelines/wan/test_wan_video_to_video.py
index 27ada121ca48..3804e972b97f 100644
--- a/tests/pipelines/wan/test_wan_video_to_video.py
+++ b/tests/pipelines/wan/test_wan_video_to_video.py
@@ -16,7 +16,7 @@
import torch
from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanTransformer3DModel, WanVideoToVideoPipeline
@@ -62,7 +62,8 @@ def get_dummy_components(self):
torch.manual_seed(0)
scheduler = UniPCMultistepScheduler(flow_shift=3.0)
- text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+ config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+ text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
torch.manual_seed(0)
diff --git a/tests/remote/test_remote_decode.py b/tests/remote/test_remote_decode.py
index 27170cba0835..e48ddcb17438 100644
--- a/tests/remote/test_remote_decode.py
+++ b/tests/remote/test_remote_decode.py
@@ -14,7 +14,6 @@
# limitations under the License.
import unittest
-from typing import Tuple, Union
import numpy as np
import PIL.Image
@@ -44,13 +43,13 @@
class RemoteAutoencoderKLMixin:
- shape: Tuple[int, ...] = None
- out_hw: Tuple[int, int] = None
+ shape: tuple[int, ...] = None
+ out_hw: tuple[int, int] = None
endpoint: str = None
dtype: torch.dtype = None
scaling_factor: float = None
shift_factor: float = None
- processor_cls: Union[VaeImageProcessor, VideoProcessor] = None
+ processor_cls: VaeImageProcessor | VideoProcessor = None
output_pil_slice: torch.Tensor = None
output_pt_slice: torch.Tensor = None
partial_postprocess_return_pt_slice: torch.Tensor = None
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
index 43f3253925df..53c1b8aa26ce 100644
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -18,7 +18,7 @@
from contextlib import contextmanager
from io import BytesIO, StringIO
from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Tuple, Union
import numpy as np
import PIL.Image
@@ -168,7 +168,7 @@ def assert_tensors_close(
max_diff = abs_diff.max().item()
flat_idx = abs_diff.argmax().item()
- max_idx = tuple(torch.unravel_index(torch.tensor(flat_idx), actual.shape).tolist())
+ max_idx = tuple(idx.item() for idx in torch.unravel_index(torch.tensor(flat_idx), actual.shape))
threshold = atol + rtol * expected.abs()
mismatched = (abs_diff > threshold).sum().item()
@@ -869,7 +869,7 @@ def get_python_version():
return major, minor
-def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -> np.ndarray:
+def load_numpy(arry: str | np.ndarray, local_path: str | None = None) -> np.ndarray:
if isinstance(arry, str):
if local_path is not None:
# local_path can be passed to correct images of tests
@@ -895,14 +895,14 @@ def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -
return arry
-def load_pt(url: str, map_location: Optional[str] = None, weights_only: Optional[bool] = True):
+def load_pt(url: str, map_location: str | None = None, weights_only: bool = True):
response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT)
response.raise_for_status()
arry = torch.load(BytesIO(response.content), map_location=map_location, weights_only=weights_only)
return arry
-def load_image(image: Union[str, PIL.Image.Image]) -> PIL.Image.Image:
+def load_image(image: str | PIL.Image.Image) -> PIL.Image.Image:
"""
Loads `image` to a PIL Image.
@@ -943,7 +943,7 @@ def preprocess_image(image: PIL.Image, batch_size: int):
return 2.0 * image - 1.0
-def export_to_gif(image: List[PIL.Image.Image], output_gif_path: str = None) -> str:
+def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None) -> str:
if output_gif_path is None:
output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name
@@ -1037,7 +1037,7 @@ def export_to_obj(mesh, output_obj_path: str = None):
f.writelines("\n".join(combined_data))
-def export_to_video(video_frames: List[np.ndarray], output_video_path: str = None) -> str:
+def export_to_video(video_frames: list[np.ndarray], output_video_path: str = None) -> str:
if is_opencv_available():
import cv2
else:
@@ -1218,7 +1218,7 @@ def summary_failures_short(tr):
# Adapted from https://github.com/huggingface/transformers/blob/000e52aec8850d3fe2f360adc6fd256e5b47fe4c/src/transformers..testing_utils.py#L1905
-def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None):
+def is_flaky(max_attempts: int = 5, wait_before_retry: float | None = None, description: str | None = None):
"""
To decorate flaky tests (methods or entire classes). They will be retried on failures.
@@ -1357,7 +1357,12 @@ def enable_full_determinism():
# variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set,
# depending on the CUDA version, so we set them both here
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
- os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
+ # Use larger workspace size for PyTorch 2.10+ to avoid CUBLAS_STATUS_NOT_INITIALIZED errors
+ # (catches 2.11 dev versions which report as >= 2.10)
+ if is_torch_version(">=", "2.10"):
+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+ else:
+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
torch.use_deterministic_algorithms(True)
# Enable CUDNN deterministic mode
@@ -1475,7 +1480,7 @@ def _is_torch_fp64_available(device):
# This dispatches a defined function according to the accelerator from the function definitions.
-def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable], *args, **kwargs):
+def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs):
if device not in dispatch_table:
return dispatch_table["default"](*args, **kwargs)
@@ -1533,7 +1538,7 @@ def backend_supports_training(device: str):
# Guard for when Torch is not available
if is_torch_available():
# Update device function dict mapping
- def update_mapping_from_spec(device_fn_dict: Dict[str, Callable], attribute_name: str):
+ def update_mapping_from_spec(device_fn_dict: dict[str, Callable], attribute_name: str):
try:
# Try to import the function directly
spec_fn = getattr(device_spec_module, attribute_name)
@@ -1629,10 +1634,10 @@ def _get_expected_safetensors_files(
module: torch.nn.Module,
offload_to_disk_path: str,
offload_type: str,
- num_blocks_per_group: Optional[int] = None,
- block_modules: Optional[List[str]] = None,
+ num_blocks_per_group: int | None = None,
+ block_modules: list[str] | None = None,
module_prefix: str = "",
- ) -> Set[str]:
+ ) -> set[str]:
expected_files = set()
def get_hashed_filename(group_id: str) -> str:
@@ -1712,8 +1717,8 @@ def _check_safetensors_serialization(
module: torch.nn.Module,
offload_to_disk_path: str,
offload_type: str,
- num_blocks_per_group: Optional[int] = None,
- block_modules: Optional[List[str]] = None,
+ num_blocks_per_group: int | None = None,
+ block_modules: list[str] | None = None,
) -> bool:
if not os.path.isdir(offload_to_disk_path):
return False, None, None
diff --git a/utils/custom_init_isort.py b/utils/custom_init_isort.py
index cc3bccb9bd63..2194b5a57063 100644
--- a/utils/custom_init_isort.py
+++ b/utils/custom_init_isort.py
@@ -63,7 +63,7 @@ def get_indent(line: str) -> str:
def split_code_in_indented_blocks(
- code: str, indent_level: str = "", start_prompt: Optional[str] = None, end_prompt: Optional[str] = None
+ code: str, indent_level: str = "", start_prompt: str | None = None, end_prompt: str | None = None
) -> List[str]:
"""
Split some code into its indented blocks, starting at a given level.
diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py
index abdc9fd409db..59ef5499cc2a 100644
--- a/utils/tests_fetcher.py
+++ b/utils/tests_fetcher.py
@@ -56,7 +56,7 @@
import re
from contextlib import contextmanager
from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Tuple, Union
from git import Repo
@@ -726,7 +726,7 @@ def init_test_examples_dependencies() -> Tuple[Dict[str, List[str]], List[str]]:
return test_example_deps, all_examples
-def create_reverse_dependency_map() -> Dict[str, List[str]]:
+def create_reverse_dependency_map() -> dict[str, List[str]]:
"""
Create the dependency map from module/test filename to the list of modules/tests that depend on it recursively.
@@ -778,7 +778,7 @@ def create_reverse_dependency_map() -> Dict[str, List[str]]:
return reverse_map
-def create_module_to_test_map(reverse_map: Dict[str, List[str]] = None) -> Dict[str, List[str]]:
+def create_module_to_test_map(reverse_map: Dict[str, List[str]] = None) -> dict[str, List[str]]:
"""
Extract the tests from the reverse_dependency_map and potentially filters the model tests.
@@ -864,7 +864,7 @@ def update_test_map_with_core_pipelines(json_output_file: str):
json.dump(test_map, fp, ensure_ascii=False)
-def create_json_map(test_files_to_run: List[str], json_output_file: Optional[str] = None):
+def create_json_map(test_files_to_run: List[str], json_output_file: str | None = None):
"""
Creates a map from a list of tests to run to easily split them by category, when running parallelism of slow tests.
@@ -909,7 +909,7 @@ def create_json_map(test_files_to_run: List[str], json_output_file: Optional[str
def infer_tests_to_run(
output_file: str,
diff_with_last_commit: bool = False,
- json_output_file: Optional[str] = None,
+ json_output_file: str | None = None,
):
"""
The main function called by the test fetcher. Determines the tests to run from the diff.
@@ -1023,7 +1023,7 @@ def filter_tests(output_file: str, filters: List[str]):
f.write(" ".join(test_files))
-def parse_commit_message(commit_message: str) -> Dict[str, bool]:
+def parse_commit_message(commit_message: str) -> dict[str, bool]:
"""
Parses the commit message to detect if a command is there to skip, force all or part of the CI.