This view is limited to 50 files because it contains too many changes.  See the raw diff here.
Files changed (50) hide show
  1. LICENSE +21 -0
  2. PowerPaint_v2 +1 -0
  3. README.md +97 -13
  4. app.py +449 -0
  5. gradio_PowerPaint.py +573 -0
  6. model/BrushNet_CA.py +960 -0
  7. model/__init__.py +0 -0
  8. model/__pycache__/BrushNet_CA.cpython-310.pyc +0 -0
  9. model/__pycache__/__init__.cpython-310.pyc +0 -0
  10. model/diffusers_c/__init__.py +789 -0
  11. model/diffusers_c/__pycache__/__init__.cpython-310.pyc +0 -0
  12. model/diffusers_c/__pycache__/configuration_utils.cpython-310.pyc +0 -0
  13. model/diffusers_c/__pycache__/dependency_versions_check.cpython-310.pyc +0 -0
  14. model/diffusers_c/__pycache__/dependency_versions_table.cpython-310.pyc +0 -0
  15. model/diffusers_c/__pycache__/image_processor.cpython-310.pyc +0 -0
  16. model/diffusers_c/commands/__init__.py +27 -0
  17. model/diffusers_c/commands/diffusers_cli.py +43 -0
  18. model/diffusers_c/commands/env.py +84 -0
  19. model/diffusers_c/commands/fp16_safetensors.py +132 -0
  20. model/diffusers_c/configuration_utils.py +703 -0
  21. model/diffusers_c/dependency_versions_check.py +34 -0
  22. model/diffusers_c/dependency_versions_table.py +45 -0
  23. model/diffusers_c/experimental/README.md +5 -0
  24. model/diffusers_c/experimental/__init__.py +1 -0
  25. model/diffusers_c/experimental/rl/__init__.py +1 -0
  26. model/diffusers_c/experimental/rl/value_guided_sampling.py +153 -0
  27. model/diffusers_c/image_processor.py +990 -0
  28. model/diffusers_c/loaders/__init__.py +88 -0
  29. model/diffusers_c/loaders/__pycache__/__init__.cpython-310.pyc +0 -0
  30. model/diffusers_c/loaders/__pycache__/__init__.cpython-39.pyc +0 -0
  31. model/diffusers_c/loaders/__pycache__/autoencoder.cpython-310.pyc +0 -0
  32. model/diffusers_c/loaders/__pycache__/autoencoder.cpython-39.pyc +0 -0
  33. model/diffusers_c/loaders/__pycache__/controlnet.cpython-39.pyc +0 -0
  34. model/diffusers_c/loaders/__pycache__/ip_adapter.cpython-39.pyc +0 -0
  35. model/diffusers_c/loaders/__pycache__/lora.cpython-39.pyc +0 -0
  36. model/diffusers_c/loaders/__pycache__/lora_conversion_utils.cpython-39.pyc +0 -0
  37. model/diffusers_c/loaders/__pycache__/peft.cpython-310.pyc +0 -0
  38. model/diffusers_c/loaders/__pycache__/peft.cpython-39.pyc +0 -0
  39. model/diffusers_c/loaders/__pycache__/single_file.cpython-39.pyc +0 -0
  40. model/diffusers_c/loaders/__pycache__/single_file_utils.cpython-310.pyc +0 -0
  41. model/diffusers_c/loaders/__pycache__/single_file_utils.cpython-39.pyc +0 -0
  42. model/diffusers_c/loaders/__pycache__/textual_inversion.cpython-39.pyc +0 -0
  43. model/diffusers_c/loaders/__pycache__/unet.cpython-310.pyc +0 -0
  44. model/diffusers_c/loaders/__pycache__/unet.cpython-39.pyc +0 -0
  45. model/diffusers_c/loaders/__pycache__/utils.cpython-310.pyc +0 -0
  46. model/diffusers_c/loaders/__pycache__/utils.cpython-39.pyc +0 -0
  47. model/diffusers_c/loaders/autoencoder.py +146 -0
  48. model/diffusers_c/loaders/controlnet.py +136 -0
  49. model/diffusers_c/loaders/ip_adapter.py +281 -0
  50. model/diffusers_c/loaders/lora.py +1349 -0
LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2024 OpenMMLab
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
PowerPaint_v2 ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit 8e039c8f98d8bbdf0c7258104275a4dcf1d1f5fb
README.md CHANGED
@@ -1,13 +1,97 @@
1
- ---
2
- title: PowerPaint
3
- emoji: 🐠
4
- colorFrom: green
5
- colorTo: gray
6
- sdk: gradio
7
- sdk_version: 4.33.0
8
- app_file: app.py
9
- pinned: false
10
- license: mit
11
- ---
12
-
13
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A Task is Worth One Word: Learning with Task Prompts for High-Quality Versatile Image Inpainting
2
+
3
+
4
+ ### [Project Page](https://powerpaint.github.io/) | [Paper](https://arxiv.org/abs/2312.03594) | [Online Demo(OpenXlab)](https://openxlab.org.cn/apps/detail/rangoliu/PowerPaint#basic-information)
5
+
6
+ This README provides a step-by-step guide to download the repository, set up the required virtual environment named "PowerPaint" using conda, and run PowerPaint with or without ControlNet.
7
+
8
+ **Feel free to try it and give it a star!**:star:
9
+
10
+ ## 🚀 News
11
+
12
+ **May 22, 2024**:fire:
13
+
14
+ - We have open-sourced the model weights for PowerPaint v2-1, rectifying some existing issues that were present during the training process of version 2. [![HuggingFace Model](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-blue)](https://huggingface.co/JunhaoZhuang/PowerPaint-v2-1)
15
+
16
+ **April 7, 2024**:fire:
17
+
18
+ - We open source the model weights and code for PowerPaint v2. [![Open in OpenXLab](https://cdn-static.openxlab.org.cn/header/openxlab_models.svg)](https://openxlab.org.cn/models/detail/zhuangjunhao/PowerPaint_v2) [![HuggingFace Model](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-blue)](https://huggingface.co/JunhaoZhuang/PowerPaint_v2)
19
+
20
+ **April 6, 2024**:
21
+
22
+ - We have retrained a new PowerPaint, taking inspiration from Brushnet. The [Online Demo](https://openxlab.org.cn/apps/detail/rangoliu/PowerPaint) has been updated accordingly. **We plan to release the model weights and code as open source in the next few days**.
23
+ - Tips: We preserve the cross-attention layer that was deleted by BrushNet for the task prompts input.
24
+
25
+ | | Object insertion | Object Removal|Shape-guided Object Insertion|Outpainting|
26
+ |-----------------|-----------------|-----------------|-----------------|-----------------|
27
+ | Original Image| ![cropinput](https://github.com/Sanster/IOPaint/assets/108931120/bf91a1e8-8eaf-4be6-b47d-b8e43c9d182a)|![cropinput](https://github.com/Sanster/IOPaint/assets/108931120/c7e56119-aa57-4761-b6aa-56f8a0b72456)|![image](https://github.com/Sanster/IOPaint/assets/108931120/cbbfe84e-2bf1-425b-8349-f7874f2e978c)|![cropinput](https://github.com/Sanster/IOPaint/assets/108931120/134bb707-0fe5-4d22-a0ca-d440fa521365)|
28
+ | Output| ![image](https://github.com/Sanster/IOPaint/assets/108931120/ee777506-d336-4275-94f6-31abf9521866)| ![image](https://github.com/Sanster/IOPaint/assets/108931120/e9d8cf6c-13b8-443c-b327-6f27da54cda6)|![image](https://github.com/Sanster/IOPaint/assets/108931120/cc3008c9-37dd-4d98-ad43-58f67be872dc)|![image](https://github.com/Sanster/IOPaint/assets/108931120/18d8ca23-e6d7-4680-977f-e66341312476)|
29
+
30
+ **December 22, 2023**:wrench:
31
+
32
+ - The logical error in loading ControlNet has been rectified. The `gradio_PowerPaint.py` file and [Online Demo](https://openxlab.org.cn/apps/detail/rangoliu/PowerPaint) have also been updated.
33
+
34
+ **December 18, 2023**
35
+
36
+ *Enhanced PowerPaint Model*
37
+
38
+ - We are delighted to announce the release of more stable model weights. These refined weights can now be accessed on [Hugging Face](https://huggingface.co/JunhaoZhuang/PowerPaint-v1/tree/main). The `gradio_PowerPaint.py` file and [Online Demo](https://openxlab.org.cn/apps/detail/rangoliu/PowerPaint) have also been updated as part of this release.
39
+
40
+
41
+
42
+ ________________
43
+ <img src='https://github.com/open-mmlab/mmagic/assets/12782558/acd01391-c73f-4997-aafd-0869aebcc915'/>
44
+
45
+ ## Getting Started
46
+
47
+ ```bash
48
+ # Clone the Repository
49
+ git clone https://github.com/zhuang2002/PowerPaint.git
50
+
51
+ # Navigate to the Repository
52
+ cd projects/powerpaint
53
+
54
+ # Create Virtual Environment with Conda
55
+ conda create --name PowerPaint python=3.9
56
+ conda activate PowerPaint
57
+
58
+ # Install Dependencies
59
+ pip install -r requirements.txt
60
+ ```
61
+ ## PowerPaint v2
62
+
63
+ ```bash
64
+ python gradio_PowerPaint_BrushNet.py
65
+ ```
66
+
67
+ ## PowerPaint v1
68
+
69
+ ```bash
70
+ # Create Models Folder
71
+ mkdir models
72
+
73
+ # Set up Git LFS
74
+ git lfs install
75
+
76
+ # Clone PowerPaint Model
77
+ git lfs clone https://huggingface.co/JunhaoZhuang/PowerPaint-v1/ ./models
78
+
79
+ python gradio_PowerPaint.py
80
+ ```
81
+
82
+ This command will launch the Gradio interface for PowerPaint.
83
+
84
+ Feel free to explore and edit images with PowerPaint!
85
+
86
+ ## BibTeX
87
+
88
+ ```
89
+ @misc{zhuang2023task,
90
+ title={A Task is Worth One Word: Learning with Task Prompts for High-Quality Versatile Image Inpainting},
91
+ author={Junhao Zhuang and Yanhong Zeng and Wenran Liu and Chun Yuan and Kai Chen},
92
+ year={2023},
93
+ eprint={2312.03594},
94
+ archivePrefix={arXiv},
95
+ primaryClass={cs.CV}
96
+ }
97
+ ```
app.py ADDED
@@ -0,0 +1,449 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import random
3
+
4
+ import gradio as gr
5
+ import numpy as np
6
+ import torch
7
+ from PIL import Image, ImageFilter
8
+ from transformers import CLIPTextModel
9
+
10
+ from diffusers import UniPCMultistepScheduler
11
+ from model.BrushNet_CA import BrushNetModel
12
+ from model.diffusers_c.models import UNet2DConditionModel
13
+ from pipeline.pipeline_PowerPaint_Brushnet_CA import StableDiffusionPowerPaintBrushNetPipeline
14
+ from utils.utils import TokenizerWrapper, add_tokens
15
+
16
+
17
+ base_path = "./PowerPaint_v2"
18
+ os.system("apt install git")
19
+ os.system("apt install git-lfs")
20
+ os.system(f"git lfs clone https://code.openxlab.org.cn/zhuangjunhao/PowerPaint_v2.git {base_path}")
21
+ os.system(f"cd {base_path} && git lfs pull")
22
+ os.system("cd ..")
23
+ torch.set_grad_enabled(False)
24
+ context_prompt = ""
25
+ context_negative_prompt = ""
26
+ base_model_path = "./PowerPaint_v2/realisticVisionV60B1_v51VAE/"
27
+ dtype = torch.float16
28
+ unet = UNet2DConditionModel.from_pretrained(
29
+ "runwayml/stable-diffusion-v1-5", subfolder="unet", revision=None, torch_dtype=dtype
30
+ )
31
+ text_encoder_brushnet = CLIPTextModel.from_pretrained(
32
+ "runwayml/stable-diffusion-v1-5", subfolder="text_encoder", revision=None, torch_dtype=dtype
33
+ )
34
+ brushnet = BrushNetModel.from_unet(unet)
35
+ global pipe
36
+ pipe = StableDiffusionPowerPaintBrushNetPipeline.from_pretrained(
37
+ base_model_path,
38
+ brushnet=brushnet,
39
+ text_encoder_brushnet=text_encoder_brushnet,
40
+ torch_dtype=dtype,
41
+ low_cpu_mem_usage=False,
42
+ safety_checker=None,
43
+ )
44
+ pipe.unet = UNet2DConditionModel.from_pretrained(base_model_path, subfolder="unet", revision=None, torch_dtype=dtype)
45
+ pipe.tokenizer = TokenizerWrapper(from_pretrained=base_model_path, subfolder="tokenizer", revision=None)
46
+ add_tokens(
47
+ tokenizer=pipe.tokenizer,
48
+ text_encoder=pipe.text_encoder_brushnet,
49
+ placeholder_tokens=["P_ctxt", "P_shape", "P_obj"],
50
+ initialize_tokens=["a", "a", "a"],
51
+ num_vectors_per_token=10,
52
+ )
53
+ from safetensors.torch import load_model
54
+
55
+
56
+ load_model(pipe.brushnet, "./PowerPaint_v2/PowerPaint_Brushnet/diffusion_pytorch_model.safetensors")
57
+
58
+ pipe.text_encoder_brushnet.load_state_dict(
59
+ torch.load("./PowerPaint_v2/PowerPaint_Brushnet/pytorch_model.bin"), strict=False
60
+ )
61
+
62
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
63
+
64
+ pipe.enable_model_cpu_offload()
65
+ global current_control
66
+ current_control = "canny"
67
+ # controlnet_conditioning_scale = 0.8
68
+
69
+
70
+ def set_seed(seed):
71
+ torch.manual_seed(seed)
72
+ torch.cuda.manual_seed(seed)
73
+ torch.cuda.manual_seed_all(seed)
74
+ np.random.seed(seed)
75
+ random.seed(seed)
76
+
77
+
78
+ def add_task(control_type):
79
+ # print(control_type)
80
+ if control_type == "object-removal":
81
+ promptA = "P_ctxt"
82
+ promptB = "P_ctxt"
83
+ negative_promptA = "P_obj"
84
+ negative_promptB = "P_obj"
85
+ elif control_type == "context-aware":
86
+ promptA = "P_ctxt"
87
+ promptB = "P_ctxt"
88
+ negative_promptA = ""
89
+ negative_promptB = ""
90
+ elif control_type == "shape-guided":
91
+ promptA = "P_shape"
92
+ promptB = "P_ctxt"
93
+ negative_promptA = "P_shape"
94
+ negative_promptB = "P_ctxt"
95
+ elif control_type == "image-outpainting":
96
+ promptA = "P_ctxt"
97
+ promptB = "P_ctxt"
98
+ negative_promptA = "P_obj"
99
+ negative_promptB = "P_obj"
100
+ else:
101
+ promptA = "P_obj"
102
+ promptB = "P_obj"
103
+ negative_promptA = "P_obj"
104
+ negative_promptB = "P_obj"
105
+
106
+ return promptA, promptB, negative_promptA, negative_promptB
107
+
108
+
109
+ def predict(
110
+ input_image,
111
+ prompt,
112
+ fitting_degree,
113
+ ddim_steps,
114
+ scale,
115
+ seed,
116
+ negative_prompt,
117
+ task,
118
+ vertical_expansion_ratio,
119
+ horizontal_expansion_ratio,
120
+ ):
121
+ size1, size2 = input_image["image"].convert("RGB").size
122
+
123
+ if task != "image-outpainting":
124
+ if size1 < size2:
125
+ input_image["image"] = input_image["image"].convert("RGB").resize((640, int(size2 / size1 * 640)))
126
+ else:
127
+ input_image["image"] = input_image["image"].convert("RGB").resize((int(size1 / size2 * 640), 640))
128
+ else:
129
+ if size1 < size2:
130
+ input_image["image"] = input_image["image"].convert("RGB").resize((512, int(size2 / size1 * 512)))
131
+ else:
132
+ input_image["image"] = input_image["image"].convert("RGB").resize((int(size1 / size2 * 512), 512))
133
+
134
+ if task == "image-outpainting" or task == "context-aware":
135
+ prompt = prompt + " empty scene"
136
+ if task == "object-removal":
137
+ prompt = prompt + " empty scene blur"
138
+
139
+ if vertical_expansion_ratio != None and horizontal_expansion_ratio != None:
140
+ o_W, o_H = input_image["image"].convert("RGB").size
141
+ c_W = int(horizontal_expansion_ratio * o_W)
142
+ c_H = int(vertical_expansion_ratio * o_H)
143
+
144
+ expand_img = np.ones((c_H, c_W, 3), dtype=np.uint8) * 127
145
+ original_img = np.array(input_image["image"])
146
+ expand_img[
147
+ int((c_H - o_H) / 2.0) : int((c_H - o_H) / 2.0) + o_H,
148
+ int((c_W - o_W) / 2.0) : int((c_W - o_W) / 2.0) + o_W,
149
+ :,
150
+ ] = original_img
151
+
152
+ blurry_gap = 10
153
+
154
+ expand_mask = np.ones((c_H, c_W, 3), dtype=np.uint8) * 255
155
+ if vertical_expansion_ratio == 1 and horizontal_expansion_ratio != 1:
156
+ expand_mask[
157
+ int((c_H - o_H) / 2.0) : int((c_H - o_H) / 2.0) + o_H,
158
+ int((c_W - o_W) / 2.0) + blurry_gap : int((c_W - o_W) / 2.0) + o_W - blurry_gap,
159
+ :,
160
+ ] = 0
161
+ elif vertical_expansion_ratio != 1 and horizontal_expansion_ratio != 1:
162
+ expand_mask[
163
+ int((c_H - o_H) / 2.0) + blurry_gap : int((c_H - o_H) / 2.0) + o_H - blurry_gap,
164
+ int((c_W - o_W) / 2.0) + blurry_gap : int((c_W - o_W) / 2.0) + o_W - blurry_gap,
165
+ :,
166
+ ] = 0
167
+ elif vertical_expansion_ratio != 1 and horizontal_expansion_ratio == 1:
168
+ expand_mask[
169
+ int((c_H - o_H) / 2.0) + blurry_gap : int((c_H - o_H) / 2.0) + o_H - blurry_gap,
170
+ int((c_W - o_W) / 2.0) : int((c_W - o_W) / 2.0) + o_W,
171
+ :,
172
+ ] = 0
173
+
174
+ input_image["image"] = Image.fromarray(expand_img)
175
+ input_image["mask"] = Image.fromarray(expand_mask)
176
+
177
+ promptA, promptB, negative_promptA, negative_promptB = add_task(task)
178
+ img = np.array(input_image["image"].convert("RGB"))
179
+
180
+ W = int(np.shape(img)[0] - np.shape(img)[0] % 8)
181
+ H = int(np.shape(img)[1] - np.shape(img)[1] % 8)
182
+ input_image["image"] = input_image["image"].resize((H, W))
183
+ input_image["mask"] = input_image["mask"].resize((H, W))
184
+
185
+ np_inpimg = np.array(input_image["image"])
186
+ np_inmask = np.array(input_image["mask"]) / 255.0
187
+
188
+ np_inpimg = np_inpimg * (1 - np_inmask)
189
+
190
+ input_image["image"] = Image.fromarray(np_inpimg.astype(np.uint8)).convert("RGB")
191
+
192
+ set_seed(seed)
193
+ global pipe
194
+ result = pipe(
195
+ promptA=promptA,
196
+ promptB=promptB,
197
+ promptU=prompt,
198
+ tradoff=fitting_degree,
199
+ tradoff_nag=fitting_degree,
200
+ image=input_image["image"].convert("RGB"),
201
+ mask=input_image["mask"].convert("RGB"),
202
+ num_inference_steps=ddim_steps,
203
+ generator=torch.Generator("cuda").manual_seed(seed),
204
+ brushnet_conditioning_scale=1.0,
205
+ negative_promptA=negative_promptA,
206
+ negative_promptB=negative_promptB,
207
+ negative_promptU=negative_prompt,
208
+ guidance_scale=scale,
209
+ width=H,
210
+ height=W,
211
+ ).images[0]
212
+ mask_np = np.array(input_image["mask"].convert("RGB"))
213
+ red = np.array(result).astype("float") * 1
214
+ red[:, :, 0] = 180.0
215
+ red[:, :, 2] = 0
216
+ red[:, :, 1] = 0
217
+ result_m = np.array(result)
218
+ result_m = Image.fromarray(
219
+ (
220
+ result_m.astype("float") * (1 - mask_np.astype("float") / 512.0) + mask_np.astype("float") / 512.0 * red
221
+ ).astype("uint8")
222
+ )
223
+ m_img = input_image["mask"].convert("RGB").filter(ImageFilter.GaussianBlur(radius=3))
224
+ m_img = np.asarray(m_img) / 255.0
225
+ img_np = np.asarray(input_image["image"].convert("RGB")) / 255.0
226
+ ours_np = np.asarray(result) / 255.0
227
+ ours_np = ours_np * m_img + (1 - m_img) * img_np
228
+ result_paste = Image.fromarray(np.uint8(ours_np * 255))
229
+
230
+ dict_res = [input_image["mask"].convert("RGB"), result_m]
231
+
232
+ dict_out = [result]
233
+
234
+ return dict_out, dict_res
235
+
236
+
237
+ def infer(
238
+ input_image,
239
+ text_guided_prompt,
240
+ text_guided_negative_prompt,
241
+ shape_guided_prompt,
242
+ shape_guided_negative_prompt,
243
+ fitting_degree,
244
+ ddim_steps,
245
+ scale,
246
+ seed,
247
+ task,
248
+ vertical_expansion_ratio,
249
+ horizontal_expansion_ratio,
250
+ outpaint_prompt,
251
+ outpaint_negative_prompt,
252
+ removal_prompt,
253
+ removal_negative_prompt,
254
+ context_prompt,
255
+ context_negative_prompt,
256
+ ):
257
+ if task == "text-guided":
258
+ prompt = text_guided_prompt
259
+ negative_prompt = text_guided_negative_prompt
260
+ elif task == "shape-guided":
261
+ prompt = shape_guided_prompt
262
+ negative_prompt = shape_guided_negative_prompt
263
+ elif task == "object-removal":
264
+ prompt = removal_prompt
265
+ negative_prompt = removal_negative_prompt
266
+ elif task == "context-aware":
267
+ prompt = context_prompt
268
+ negative_prompt = context_negative_prompt
269
+ elif task == "image-outpainting":
270
+ prompt = outpaint_prompt
271
+ negative_prompt = outpaint_negative_prompt
272
+ return predict(
273
+ input_image,
274
+ prompt,
275
+ fitting_degree,
276
+ ddim_steps,
277
+ scale,
278
+ seed,
279
+ negative_prompt,
280
+ task,
281
+ vertical_expansion_ratio,
282
+ horizontal_expansion_ratio,
283
+ )
284
+ else:
285
+ task = "text-guided"
286
+ prompt = text_guided_prompt
287
+ negative_prompt = text_guided_negative_prompt
288
+
289
+ return predict(input_image, prompt, fitting_degree, ddim_steps, scale, seed, negative_prompt, task, None, None)
290
+
291
+
292
+ def select_tab_text_guided():
293
+ return "text-guided"
294
+
295
+
296
+ def select_tab_object_removal():
297
+ return "object-removal"
298
+
299
+
300
+ def select_tab_context_aware():
301
+ return "context-aware"
302
+
303
+
304
+ def select_tab_image_outpainting():
305
+ return "image-outpainting"
306
+
307
+
308
+ def select_tab_shape_guided():
309
+ return "shape-guided"
310
+
311
+
312
+ with gr.Blocks(css="style.css") as demo:
313
+ with gr.Row():
314
+ gr.Markdown(
315
+ "<div align='center'><font size='18'>PowerPaint: High-Quality Versatile Image Inpainting</font></div>" # noqa
316
+ )
317
+ with gr.Row():
318
+ gr.Markdown(
319
+ "<div align='center'><font size='5'><a href='https://powerpaint.github.io/'>Project Page</a> &ensp;" # noqa
320
+ "<a href='https://arxiv.org/abs/2312.03594/'>Paper</a> &ensp;"
321
+ "<a href='https://github.com/zhuang2002/PowerPaint'>Code</a> </font></div>" # noqa
322
+ )
323
+ with gr.Row():
324
+ gr.Markdown(
325
+ "**Note:** Due to network-related factors, the page may experience occasional bugs! If the inpainting results deviate significantly from expectations, consider toggling between task options to refresh the content." # noqa
326
+ )
327
+ # Attention: Due to network-related factors, the page may experience occasional bugs. If the inpainting results deviate significantly from expectations, consider toggling between task options to refresh the content.
328
+ with gr.Row():
329
+ with gr.Column():
330
+ gr.Markdown("### Input image and draw mask")
331
+ input_image = gr.Image(source="upload", tool="sketch", type="pil")
332
+
333
+ task = gr.Radio(
334
+ ["text-guided", "object-removal", "shape-guided", "image-outpainting"], show_label=False, visible=False
335
+ )
336
+
337
+ # Text-guided object inpainting
338
+ with gr.Tab("Text-guided object inpainting") as tab_text_guided:
339
+ enable_text_guided = gr.Checkbox(
340
+ label="Enable text-guided object inpainting", value=True, interactive=False
341
+ )
342
+ text_guided_prompt = gr.Textbox(label="Prompt")
343
+ text_guided_negative_prompt = gr.Textbox(label="negative_prompt")
344
+ tab_text_guided.select(fn=select_tab_text_guided, inputs=None, outputs=task)
345
+
346
+ # Object removal inpainting
347
+ with gr.Tab("Object removal inpainting") as tab_object_removal:
348
+ enable_object_removal = gr.Checkbox(
349
+ label="Enable object removal inpainting",
350
+ value=True,
351
+ info="The recommended configuration for the Guidance Scale is 10 or higher. \
352
+ If undesired objects appear in the masked area, \
353
+ you can address this by specifically increasing the Guidance Scale.",
354
+ interactive=False,
355
+ )
356
+ removal_prompt = gr.Textbox(label="Prompt")
357
+ removal_negative_prompt = gr.Textbox(label="negative_prompt")
358
+ context_prompt = removal_prompt
359
+ context_negative_prompt = removal_negative_prompt
360
+ tab_object_removal.select(fn=select_tab_object_removal, inputs=None, outputs=task)
361
+
362
+ # Object image outpainting
363
+ with gr.Tab("Image outpainting") as tab_image_outpainting:
364
+ enable_object_removal = gr.Checkbox(
365
+ label="Enable image outpainting",
366
+ value=True,
367
+ info="The recommended configuration for the Guidance Scale is 15 or higher. \
368
+ If unwanted random objects appear in the extended image region, \
369
+ you can enhance the cleanliness of the extension area by increasing the Guidance Scale.",
370
+ interactive=False,
371
+ )
372
+ outpaint_prompt = gr.Textbox(label="Outpainting_prompt")
373
+ outpaint_negative_prompt = gr.Textbox(label="Outpainting_negative_prompt")
374
+ horizontal_expansion_ratio = gr.Slider(
375
+ label="horizontal expansion ratio",
376
+ minimum=1,
377
+ maximum=4,
378
+ step=0.05,
379
+ value=1,
380
+ )
381
+ vertical_expansion_ratio = gr.Slider(
382
+ label="vertical expansion ratio",
383
+ minimum=1,
384
+ maximum=4,
385
+ step=0.05,
386
+ value=1,
387
+ )
388
+ tab_image_outpainting.select(fn=select_tab_image_outpainting, inputs=None, outputs=task)
389
+
390
+ # Shape-guided object inpainting
391
+ with gr.Tab("Shape-guided object inpainting") as tab_shape_guided:
392
+ enable_shape_guided = gr.Checkbox(
393
+ label="Enable shape-guided object inpainting", value=True, interactive=False
394
+ )
395
+ shape_guided_prompt = gr.Textbox(label="shape_guided_prompt")
396
+ shape_guided_negative_prompt = gr.Textbox(label="shape_guided_negative_prompt")
397
+ fitting_degree = gr.Slider(
398
+ label="fitting degree",
399
+ minimum=0.3,
400
+ maximum=1,
401
+ step=0.05,
402
+ value=1,
403
+ )
404
+ tab_shape_guided.select(fn=select_tab_shape_guided, inputs=None, outputs=task)
405
+
406
+ run_button = gr.Button(label="Run")
407
+ with gr.Accordion("Advanced options", open=False):
408
+ ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=50, step=1)
409
+ scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=45.0, value=12, step=0.1)
410
+ seed = gr.Slider(
411
+ label="Seed",
412
+ minimum=0,
413
+ maximum=2147483647,
414
+ step=1,
415
+ randomize=True,
416
+ )
417
+ with gr.Column():
418
+ gr.Markdown("### Inpainting result")
419
+ inpaint_result = gr.Gallery(label="Generated images", show_label=False, columns=2)
420
+ gr.Markdown("### Mask")
421
+ gallery = gr.Gallery(label="Generated masks", show_label=False, columns=2)
422
+
423
+ run_button.click(
424
+ fn=infer,
425
+ inputs=[
426
+ input_image,
427
+ text_guided_prompt,
428
+ text_guided_negative_prompt,
429
+ shape_guided_prompt,
430
+ shape_guided_negative_prompt,
431
+ fitting_degree,
432
+ ddim_steps,
433
+ scale,
434
+ seed,
435
+ task,
436
+ vertical_expansion_ratio,
437
+ horizontal_expansion_ratio,
438
+ outpaint_prompt,
439
+ outpaint_negative_prompt,
440
+ removal_prompt,
441
+ removal_negative_prompt,
442
+ context_prompt,
443
+ context_negative_prompt,
444
+ ],
445
+ outputs=[inpaint_result, gallery],
446
+ )
447
+
448
+ demo.queue()
449
+ demo.launch(share=False, server_name="0.0.0.0", server_port=7860)
gradio_PowerPaint.py ADDED
@@ -0,0 +1,573 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import random
2
+
3
+ import cv2
4
+ import gradio as gr
5
+ import numpy as np
6
+ import torch
7
+ from controlnet_aux import HEDdetector, OpenposeDetector
8
+ from PIL import Image, ImageFilter
9
+ from transformers import DPTFeatureExtractor, DPTForDepthEstimation
10
+
11
+ from diffusers.pipelines.controlnet.pipeline_controlnet import ControlNetModel
12
+ from pipeline.pipeline_PowerPaint import StableDiffusionInpaintPipeline as Pipeline
13
+ from pipeline.pipeline_PowerPaint_ControlNet import StableDiffusionControlNetInpaintPipeline as controlnetPipeline
14
+ from utils.utils import TokenizerWrapper, add_tokens
15
+
16
+
17
+ torch.set_grad_enabled(False)
18
+
19
+ weight_dtype = torch.float16
20
+ global pipe
21
+ pipe = Pipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=weight_dtype)
22
+ pipe.tokenizer = TokenizerWrapper(
23
+ from_pretrained="runwayml/stable-diffusion-v1-5", subfolder="tokenizer", revision=None
24
+ )
25
+
26
+ add_tokens(
27
+ tokenizer=pipe.tokenizer,
28
+ text_encoder=pipe.text_encoder,
29
+ placeholder_tokens=["P_ctxt", "P_shape", "P_obj"],
30
+ initialize_tokens=["a", "a", "a"],
31
+ num_vectors_per_token=10,
32
+ )
33
+
34
+ from safetensors.torch import load_model
35
+
36
+
37
+ load_model(pipe.unet, "./models/unet/unet.safetensors")
38
+ load_model(pipe.text_encoder, "./models/unet/text_encoder.safetensors")
39
+ pipe = pipe.to("cuda")
40
+
41
+
42
+ depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
43
+ feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")
44
+ openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
45
+ hed = HEDdetector.from_pretrained("lllyasviel/ControlNet")
46
+
47
+ global current_control
48
+ current_control = "canny"
49
+ # controlnet_conditioning_scale = 0.8
50
+
51
+
52
+ def set_seed(seed):
53
+ torch.manual_seed(seed)
54
+ torch.cuda.manual_seed(seed)
55
+ torch.cuda.manual_seed_all(seed)
56
+ np.random.seed(seed)
57
+ random.seed(seed)
58
+
59
+
60
+ def get_depth_map(image):
61
+ image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
62
+ with torch.no_grad(), torch.autocast("cuda"):
63
+ depth_map = depth_estimator(image).predicted_depth
64
+
65
+ depth_map = torch.nn.functional.interpolate(
66
+ depth_map.unsqueeze(1),
67
+ size=(1024, 1024),
68
+ mode="bicubic",
69
+ align_corners=False,
70
+ )
71
+ depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
72
+ depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
73
+ depth_map = (depth_map - depth_min) / (depth_max - depth_min)
74
+ image = torch.cat([depth_map] * 3, dim=1)
75
+
76
+ image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
77
+ image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
78
+ return image
79
+
80
+
81
+ def add_task(prompt, negative_prompt, control_type):
82
+ # print(control_type)
83
+ if control_type == "object-removal":
84
+ promptA = "empty scene blur " + prompt + " P_ctxt"
85
+ promptB = "empty scene blur " + prompt + " P_ctxt"
86
+ negative_promptA = negative_prompt + " P_obj"
87
+ negative_promptB = negative_prompt + " P_obj"
88
+ elif control_type == "shape-guided":
89
+ promptA = prompt + " P_shape"
90
+ promptB = prompt + " P_ctxt"
91
+ negative_promptA = (
92
+ negative_prompt + ", worst quality, low quality, normal quality, bad quality, blurry P_shape"
93
+ )
94
+ negative_promptB = negative_prompt + ", worst quality, low quality, normal quality, bad quality, blurry P_ctxt"
95
+ elif control_type == "image-outpainting":
96
+ promptA = "empty scene " + prompt + " P_ctxt"
97
+ promptB = "empty scene " + prompt + " P_ctxt"
98
+ negative_promptA = negative_prompt + " P_obj"
99
+ negative_promptB = negative_prompt + " P_obj"
100
+ else:
101
+ promptA = prompt + " P_obj"
102
+ promptB = prompt + " P_obj"
103
+ negative_promptA = negative_prompt + ", worst quality, low quality, normal quality, bad quality, blurry, P_obj"
104
+ negative_promptB = negative_prompt + ", worst quality, low quality, normal quality, bad quality, blurry, P_obj"
105
+
106
+ return promptA, promptB, negative_promptA, negative_promptB
107
+
108
+
109
+ def predict(
110
+ input_image,
111
+ prompt,
112
+ fitting_degree,
113
+ ddim_steps,
114
+ scale,
115
+ seed,
116
+ negative_prompt,
117
+ task,
118
+ vertical_expansion_ratio,
119
+ horizontal_expansion_ratio,
120
+ ):
121
+ size1, size2 = input_image["image"].convert("RGB").size
122
+
123
+ if task != "image-outpainting":
124
+ if size1 < size2:
125
+ input_image["image"] = input_image["image"].convert("RGB").resize((640, int(size2 / size1 * 640)))
126
+ else:
127
+ input_image["image"] = input_image["image"].convert("RGB").resize((int(size1 / size2 * 640), 640))
128
+ else:
129
+ if size1 < size2:
130
+ input_image["image"] = input_image["image"].convert("RGB").resize((512, int(size2 / size1 * 512)))
131
+ else:
132
+ input_image["image"] = input_image["image"].convert("RGB").resize((int(size1 / size2 * 512), 512))
133
+
134
+ if vertical_expansion_ratio != None and horizontal_expansion_ratio != None:
135
+ o_W, o_H = input_image["image"].convert("RGB").size
136
+ c_W = int(horizontal_expansion_ratio * o_W)
137
+ c_H = int(vertical_expansion_ratio * o_H)
138
+
139
+ expand_img = np.ones((c_H, c_W, 3), dtype=np.uint8) * 127
140
+ original_img = np.array(input_image["image"])
141
+ expand_img[
142
+ int((c_H - o_H) / 2.0) : int((c_H - o_H) / 2.0) + o_H,
143
+ int((c_W - o_W) / 2.0) : int((c_W - o_W) / 2.0) + o_W,
144
+ :,
145
+ ] = original_img
146
+
147
+ blurry_gap = 10
148
+
149
+ expand_mask = np.ones((c_H, c_W, 3), dtype=np.uint8) * 255
150
+ if vertical_expansion_ratio == 1 and horizontal_expansion_ratio != 1:
151
+ expand_mask[
152
+ int((c_H - o_H) / 2.0) : int((c_H - o_H) / 2.0) + o_H,
153
+ int((c_W - o_W) / 2.0) + blurry_gap : int((c_W - o_W) / 2.0) + o_W - blurry_gap,
154
+ :,
155
+ ] = 0
156
+ elif vertical_expansion_ratio != 1 and horizontal_expansion_ratio != 1:
157
+ expand_mask[
158
+ int((c_H - o_H) / 2.0) + blurry_gap : int((c_H - o_H) / 2.0) + o_H - blurry_gap,
159
+ int((c_W - o_W) / 2.0) + blurry_gap : int((c_W - o_W) / 2.0) + o_W - blurry_gap,
160
+ :,
161
+ ] = 0
162
+ elif vertical_expansion_ratio != 1 and horizontal_expansion_ratio == 1:
163
+ expand_mask[
164
+ int((c_H - o_H) / 2.0) + blurry_gap : int((c_H - o_H) / 2.0) + o_H - blurry_gap,
165
+ int((c_W - o_W) / 2.0) : int((c_W - o_W) / 2.0) + o_W,
166
+ :,
167
+ ] = 0
168
+
169
+ input_image["image"] = Image.fromarray(expand_img)
170
+ input_image["mask"] = Image.fromarray(expand_mask)
171
+
172
+ promptA, promptB, negative_promptA, negative_promptB = add_task(prompt, negative_prompt, task)
173
+ print(promptA, promptB, negative_promptA, negative_promptB)
174
+ img = np.array(input_image["image"].convert("RGB"))
175
+
176
+ W = int(np.shape(img)[0] - np.shape(img)[0] % 8)
177
+ H = int(np.shape(img)[1] - np.shape(img)[1] % 8)
178
+ input_image["image"] = input_image["image"].resize((H, W))
179
+ input_image["mask"] = input_image["mask"].resize((H, W))
180
+ set_seed(seed)
181
+ global pipe
182
+ result = pipe(
183
+ promptA=promptA,
184
+ promptB=promptB,
185
+ tradoff=fitting_degree,
186
+ tradoff_nag=fitting_degree,
187
+ negative_promptA=negative_promptA,
188
+ negative_promptB=negative_promptB,
189
+ image=input_image["image"].convert("RGB"),
190
+ mask_image=input_image["mask"].convert("RGB"),
191
+ width=H,
192
+ height=W,
193
+ guidance_scale=scale,
194
+ num_inference_steps=ddim_steps,
195
+ ).images[0]
196
+ mask_np = np.array(input_image["mask"].convert("RGB"))
197
+ red = np.array(result).astype("float") * 1
198
+ red[:, :, 0] = 180.0
199
+ red[:, :, 2] = 0
200
+ red[:, :, 1] = 0
201
+ result_m = np.array(result)
202
+ result_m = Image.fromarray(
203
+ (
204
+ result_m.astype("float") * (1 - mask_np.astype("float") / 512.0) + mask_np.astype("float") / 512.0 * red
205
+ ).astype("uint8")
206
+ )
207
+ m_img = input_image["mask"].convert("RGB").filter(ImageFilter.GaussianBlur(radius=3))
208
+ m_img = np.asarray(m_img) / 255.0
209
+ img_np = np.asarray(input_image["image"].convert("RGB")) / 255.0
210
+ ours_np = np.asarray(result) / 255.0
211
+ ours_np = ours_np * m_img + (1 - m_img) * img_np
212
+ result_paste = Image.fromarray(np.uint8(ours_np * 255))
213
+
214
+ dict_res = [input_image["mask"].convert("RGB"), result_m]
215
+
216
+ dict_out = [input_image["image"].convert("RGB"), result_paste]
217
+
218
+ return dict_out, dict_res
219
+
220
+
221
+ def predict_controlnet(
222
+ input_image,
223
+ input_control_image,
224
+ control_type,
225
+ prompt,
226
+ ddim_steps,
227
+ scale,
228
+ seed,
229
+ negative_prompt,
230
+ controlnet_conditioning_scale,
231
+ ):
232
+ promptA = prompt + " P_obj"
233
+ promptB = prompt + " P_obj"
234
+ negative_promptA = negative_prompt
235
+ negative_promptB = negative_prompt
236
+ size1, size2 = input_image["image"].convert("RGB").size
237
+
238
+ if size1 < size2:
239
+ input_image["image"] = input_image["image"].convert("RGB").resize((640, int(size2 / size1 * 640)))
240
+ else:
241
+ input_image["image"] = input_image["image"].convert("RGB").resize((int(size1 / size2 * 640), 640))
242
+ img = np.array(input_image["image"].convert("RGB"))
243
+ W = int(np.shape(img)[0] - np.shape(img)[0] % 8)
244
+ H = int(np.shape(img)[1] - np.shape(img)[1] % 8)
245
+ input_image["image"] = input_image["image"].resize((H, W))
246
+ input_image["mask"] = input_image["mask"].resize((H, W))
247
+
248
+ global current_control
249
+ global pipe
250
+
251
+ base_control = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=weight_dtype)
252
+ control_pipe = controlnetPipeline(
253
+ pipe.vae, pipe.text_encoder, pipe.tokenizer, pipe.unet, base_control, pipe.scheduler, None, None, False
254
+ )
255
+ control_pipe = control_pipe.to("cuda")
256
+ current_control = "canny"
257
+ if current_control != control_type:
258
+ if control_type == "canny" or control_type is None:
259
+ control_pipe.controlnet = ControlNetModel.from_pretrained(
260
+ "lllyasviel/sd-controlnet-canny", torch_dtype=weight_dtype
261
+ )
262
+ elif control_type == "pose":
263
+ control_pipe.controlnet = ControlNetModel.from_pretrained(
264
+ "lllyasviel/sd-controlnet-openpose", torch_dtype=weight_dtype
265
+ )
266
+ elif control_type == "depth":
267
+ control_pipe.controlnet = ControlNetModel.from_pretrained(
268
+ "lllyasviel/sd-controlnet-depth", torch_dtype=weight_dtype
269
+ )
270
+ else:
271
+ control_pipe.controlnet = ControlNetModel.from_pretrained(
272
+ "lllyasviel/sd-controlnet-hed", torch_dtype=weight_dtype
273
+ )
274
+ control_pipe = control_pipe.to("cuda")
275
+ current_control = control_type
276
+
277
+ controlnet_image = input_control_image
278
+ if current_control == "canny":
279
+ controlnet_image = controlnet_image.resize((H, W))
280
+ controlnet_image = np.array(controlnet_image)
281
+ controlnet_image = cv2.Canny(controlnet_image, 100, 200)
282
+ controlnet_image = controlnet_image[:, :, None]
283
+ controlnet_image = np.concatenate([controlnet_image, controlnet_image, controlnet_image], axis=2)
284
+ controlnet_image = Image.fromarray(controlnet_image)
285
+ elif current_control == "pose":
286
+ controlnet_image = openpose(controlnet_image)
287
+ elif current_control == "depth":
288
+ controlnet_image = controlnet_image.resize((H, W))
289
+ controlnet_image = get_depth_map(controlnet_image)
290
+ else:
291
+ controlnet_image = hed(controlnet_image)
292
+
293
+ mask_np = np.array(input_image["mask"].convert("RGB"))
294
+ controlnet_image = controlnet_image.resize((H, W))
295
+ set_seed(seed)
296
+ result = control_pipe(
297
+ promptA=promptB,
298
+ promptB=promptA,
299
+ tradoff=1.0,
300
+ tradoff_nag=1.0,
301
+ negative_promptA=negative_promptA,
302
+ negative_promptB=negative_promptB,
303
+ image=input_image["image"].convert("RGB"),
304
+ mask_image=input_image["mask"].convert("RGB"),
305
+ control_image=controlnet_image,
306
+ width=H,
307
+ height=W,
308
+ guidance_scale=scale,
309
+ controlnet_conditioning_scale=controlnet_conditioning_scale,
310
+ num_inference_steps=ddim_steps,
311
+ ).images[0]
312
+ red = np.array(result).astype("float") * 1
313
+ red[:, :, 0] = 180.0
314
+ red[:, :, 2] = 0
315
+ red[:, :, 1] = 0
316
+ result_m = np.array(result)
317
+ result_m = Image.fromarray(
318
+ (
319
+ result_m.astype("float") * (1 - mask_np.astype("float") / 512.0) + mask_np.astype("float") / 512.0 * red
320
+ ).astype("uint8")
321
+ )
322
+
323
+ mask_np = np.array(input_image["mask"].convert("RGB"))
324
+ m_img = input_image["mask"].convert("RGB").filter(ImageFilter.GaussianBlur(radius=4))
325
+ m_img = np.asarray(m_img) / 255.0
326
+ img_np = np.asarray(input_image["image"].convert("RGB")) / 255.0
327
+ ours_np = np.asarray(result) / 255.0
328
+ ours_np = ours_np * m_img + (1 - m_img) * img_np
329
+ result_paste = Image.fromarray(np.uint8(ours_np * 255))
330
+ return [input_image["image"].convert("RGB"), result_paste], [controlnet_image, result_m]
331
+
332
+
333
+ def infer(
334
+ input_image,
335
+ text_guided_prompt,
336
+ text_guided_negative_prompt,
337
+ shape_guided_prompt,
338
+ shape_guided_negative_prompt,
339
+ fitting_degree,
340
+ ddim_steps,
341
+ scale,
342
+ seed,
343
+ task,
344
+ enable_control,
345
+ input_control_image,
346
+ control_type,
347
+ vertical_expansion_ratio,
348
+ horizontal_expansion_ratio,
349
+ outpaint_prompt,
350
+ outpaint_negative_prompt,
351
+ controlnet_conditioning_scale,
352
+ removal_prompt,
353
+ removal_negative_prompt,
354
+ ):
355
+ if task == "text-guided":
356
+ prompt = text_guided_prompt
357
+ negative_prompt = text_guided_negative_prompt
358
+ elif task == "shape-guided":
359
+ prompt = shape_guided_prompt
360
+ negative_prompt = shape_guided_negative_prompt
361
+ elif task == "object-removal":
362
+ prompt = removal_prompt
363
+ negative_prompt = removal_negative_prompt
364
+ elif task == "image-outpainting":
365
+ prompt = outpaint_prompt
366
+ negative_prompt = outpaint_negative_prompt
367
+ return predict(
368
+ input_image,
369
+ prompt,
370
+ fitting_degree,
371
+ ddim_steps,
372
+ scale,
373
+ seed,
374
+ negative_prompt,
375
+ task,
376
+ vertical_expansion_ratio,
377
+ horizontal_expansion_ratio,
378
+ )
379
+ else:
380
+ task = "text-guided"
381
+ prompt = text_guided_prompt
382
+ negative_prompt = text_guided_negative_prompt
383
+
384
+ if enable_control and task == "text-guided":
385
+ return predict_controlnet(
386
+ input_image,
387
+ input_control_image,
388
+ control_type,
389
+ prompt,
390
+ ddim_steps,
391
+ scale,
392
+ seed,
393
+ negative_prompt,
394
+ controlnet_conditioning_scale,
395
+ )
396
+ else:
397
+ return predict(input_image, prompt, fitting_degree, ddim_steps, scale, seed, negative_prompt, task, None, None)
398
+
399
+
400
+ def select_tab_text_guided():
401
+ return "text-guided"
402
+
403
+
404
+ def select_tab_object_removal():
405
+ return "object-removal"
406
+
407
+
408
+ def select_tab_image_outpainting():
409
+ return "image-outpainting"
410
+
411
+
412
+ def select_tab_shape_guided():
413
+ return "shape-guided"
414
+
415
+
416
+ with gr.Blocks(css="style.css") as demo:
417
+ with gr.Row():
418
+ gr.Markdown(
419
+ "<div align='center'><font size='18'>PowerPaint: High-Quality Versatile Image Inpainting</font></div>" # noqa
420
+ )
421
+ with gr.Row():
422
+ gr.Markdown(
423
+ "<div align='center'><font size='5'><a href='https://powerpaint.github.io/'>Project Page</a> &ensp;" # noqa
424
+ "<a href='https://arxiv.org/abs/2312.03594/'>Paper</a> &ensp;"
425
+ "<a href='https://github.com/open-mmlab/mmagic/tree/main/projects/powerpaint'>Code</a> </font></div>" # noqa
426
+ )
427
+ with gr.Row():
428
+ gr.Markdown(
429
+ "**Note:** Due to network-related factors, the page may experience occasional bugs! If the inpainting results deviate significantly from expectations, consider toggling between task options to refresh the content." # noqa
430
+ )
431
+ # Attention: Due to network-related factors, the page may experience occasional bugs. If the inpainting results deviate significantly from expectations, consider toggling between task options to refresh the content.
432
+ with gr.Row():
433
+ with gr.Column():
434
+ gr.Markdown("### Input image and draw mask")
435
+ input_image = gr.Image(source="upload", tool="sketch", type="pil")
436
+
437
+ task = gr.Radio(
438
+ ["text-guided", "object-removal", "shape-guided", "image-outpainting"], show_label=False, visible=False
439
+ )
440
+
441
+ # Text-guided object inpainting
442
+ with gr.Tab("Text-guided object inpainting") as tab_text_guided:
443
+ enable_text_guided = gr.Checkbox(
444
+ label="Enable text-guided object inpainting", value=True, interactive=False
445
+ )
446
+ text_guided_prompt = gr.Textbox(label="Prompt")
447
+ text_guided_negative_prompt = gr.Textbox(label="negative_prompt")
448
+ gr.Markdown("### Controlnet setting")
449
+ enable_control = gr.Checkbox(
450
+ label="Enable controlnet", info="Enable this if you want to use controlnet"
451
+ )
452
+ controlnet_conditioning_scale = gr.Slider(
453
+ label="controlnet conditioning scale",
454
+ minimum=0,
455
+ maximum=1,
456
+ step=0.05,
457
+ value=0.5,
458
+ )
459
+ control_type = gr.Radio(["canny", "pose", "depth", "hed"], label="Control type")
460
+ input_control_image = gr.Image(source="upload", type="pil")
461
+ tab_text_guided.select(fn=select_tab_text_guided, inputs=None, outputs=task)
462
+
463
+ # Object removal inpainting
464
+ with gr.Tab("Object removal inpainting") as tab_object_removal:
465
+ enable_object_removal = gr.Checkbox(
466
+ label="Enable object removal inpainting",
467
+ value=True,
468
+ info="The recommended configuration for the Guidance Scale is 10 or higher. \
469
+ If undesired objects appear in the masked area, \
470
+ you can address this by specifically increasing the Guidance Scale.",
471
+ interactive=False,
472
+ )
473
+ removal_prompt = gr.Textbox(label="Prompt")
474
+ removal_negative_prompt = gr.Textbox(label="negative_prompt")
475
+ tab_object_removal.select(fn=select_tab_object_removal, inputs=None, outputs=task)
476
+
477
+ # Object image outpainting
478
+ with gr.Tab("Image outpainting") as tab_image_outpainting:
479
+ enable_object_removal = gr.Checkbox(
480
+ label="Enable image outpainting",
481
+ value=True,
482
+ info="The recommended configuration for the Guidance Scale is 10 or higher. \
483
+ If unwanted random objects appear in the extended image region, \
484
+ you can enhance the cleanliness of the extension area by increasing the Guidance Scale.",
485
+ interactive=False,
486
+ )
487
+ outpaint_prompt = gr.Textbox(label="Outpainting_prompt")
488
+ outpaint_negative_prompt = gr.Textbox(label="Outpainting_negative_prompt")
489
+ horizontal_expansion_ratio = gr.Slider(
490
+ label="horizontal expansion ratio",
491
+ minimum=1,
492
+ maximum=4,
493
+ step=0.05,
494
+ value=1,
495
+ )
496
+ vertical_expansion_ratio = gr.Slider(
497
+ label="vertical expansion ratio",
498
+ minimum=1,
499
+ maximum=4,
500
+ step=0.05,
501
+ value=1,
502
+ )
503
+ tab_image_outpainting.select(fn=select_tab_image_outpainting, inputs=None, outputs=task)
504
+
505
+ # Shape-guided object inpainting
506
+ with gr.Tab("Shape-guided object inpainting") as tab_shape_guided:
507
+ enable_shape_guided = gr.Checkbox(
508
+ label="Enable shape-guided object inpainting", value=True, interactive=False
509
+ )
510
+ shape_guided_prompt = gr.Textbox(label="shape_guided_prompt")
511
+ shape_guided_negative_prompt = gr.Textbox(label="shape_guided_negative_prompt")
512
+ fitting_degree = gr.Slider(
513
+ label="fitting degree",
514
+ minimum=0,
515
+ maximum=1,
516
+ step=0.05,
517
+ value=1,
518
+ )
519
+ tab_shape_guided.select(fn=select_tab_shape_guided, inputs=None, outputs=task)
520
+
521
+ run_button = gr.Button(label="Run")
522
+ with gr.Accordion("Advanced options", open=False):
523
+ ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1)
524
+ scale = gr.Slider(
525
+ label="Guidance Scale",
526
+ info="For object removal and image outpainting, it is recommended to set the value at 10 or above.",
527
+ minimum=0.1,
528
+ maximum=30.0,
529
+ value=7.5,
530
+ step=0.1,
531
+ )
532
+ seed = gr.Slider(
533
+ label="Seed",
534
+ minimum=0,
535
+ maximum=2147483647,
536
+ step=1,
537
+ randomize=True,
538
+ )
539
+ with gr.Column():
540
+ gr.Markdown("### Inpainting result")
541
+ inpaint_result = gr.Gallery(label="Generated images", show_label=False, columns=2)
542
+ gr.Markdown("### Mask")
543
+ gallery = gr.Gallery(label="Generated masks", show_label=False, columns=2)
544
+
545
+ run_button.click(
546
+ fn=infer,
547
+ inputs=[
548
+ input_image,
549
+ text_guided_prompt,
550
+ text_guided_negative_prompt,
551
+ shape_guided_prompt,
552
+ shape_guided_negative_prompt,
553
+ fitting_degree,
554
+ ddim_steps,
555
+ scale,
556
+ seed,
557
+ task,
558
+ enable_control,
559
+ input_control_image,
560
+ control_type,
561
+ vertical_expansion_ratio,
562
+ horizontal_expansion_ratio,
563
+ outpaint_prompt,
564
+ outpaint_negative_prompt,
565
+ controlnet_conditioning_scale,
566
+ removal_prompt,
567
+ removal_negative_prompt,
568
+ ],
569
+ outputs=[inpaint_result, gallery],
570
+ )
571
+
572
+ demo.queue()
573
+ demo.launch(share=False, server_name="0.0.0.0", server_port=7860)
model/BrushNet_CA.py ADDED
@@ -0,0 +1,960 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sys
2
+ from dataclasses import dataclass
3
+ from typing import Any, Dict, List, Optional, Tuple, Union
4
+
5
+
6
+ sys.path.append(".model")
7
+ import torch
8
+ from torch import nn
9
+
10
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
11
+ from diffusers.models.attention_processor import (
12
+ ADDED_KV_ATTENTION_PROCESSORS,
13
+ CROSS_ATTENTION_PROCESSORS,
14
+ AttentionProcessor,
15
+ AttnAddedKVProcessor,
16
+ AttnProcessor,
17
+ )
18
+ from diffusers.models.embeddings import (
19
+ TextImageProjection,
20
+ TextImageTimeEmbedding,
21
+ TextTimeEmbedding,
22
+ TimestepEmbedding,
23
+ Timesteps,
24
+ )
25
+ from diffusers.models.modeling_utils import ModelMixin
26
+ from diffusers.utils import BaseOutput, logging
27
+ from model.diffusers_c.models.unets.unet_2d_blocks import (
28
+ CrossAttnDownBlock2D,
29
+ DownBlock2D,
30
+ get_down_block,
31
+ get_mid_block,
32
+ get_up_block,
33
+ )
34
+ from model.diffusers_c.models.unets.unet_2d_condition import UNet2DConditionModel
35
+
36
+
37
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
38
+
39
+
40
+ @dataclass
41
+ class BrushNetOutput(BaseOutput):
42
+ """
43
+ The output of [`BrushNetModel`].
44
+
45
+ Args:
46
+ up_block_res_samples (`tuple[torch.Tensor]`):
47
+ A tuple of upsample activations at different resolutions for each upsampling block. Each tensor should
48
+ be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be
49
+ used to condition the original UNet's upsampling activations.
50
+ down_block_res_samples (`tuple[torch.Tensor]`):
51
+ A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should
52
+ be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be
53
+ used to condition the original UNet's downsampling activations.
54
+ mid_down_block_re_sample (`torch.Tensor`):
55
+ The activation of the midde block (the lowest sample resolution). Each tensor should be of shape
56
+ `(batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution)`.
57
+ Output can be used to condition the original UNet's middle block activation.
58
+ """
59
+
60
+ up_block_res_samples: Tuple[torch.Tensor]
61
+ down_block_res_samples: Tuple[torch.Tensor]
62
+ mid_block_res_sample: torch.Tensor
63
+
64
+
65
+ class BrushNetModel(ModelMixin, ConfigMixin):
66
+ """
67
+ A BrushNet model.
68
+
69
+ Args:
70
+ in_channels (`int`, defaults to 4):
71
+ The number of channels in the input sample.
72
+ flip_sin_to_cos (`bool`, defaults to `True`):
73
+ Whether to flip the sin to cos in the time embedding.
74
+ freq_shift (`int`, defaults to 0):
75
+ The frequency shift to apply to the time embedding.
76
+ down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
77
+ The tuple of downsample blocks to use.
78
+ mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
79
+ Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or
80
+ `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
81
+ up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
82
+ The tuple of upsample blocks to use.
83
+ only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`):
84
+ block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`):
85
+ The tuple of output channels for each block.
86
+ layers_per_block (`int`, defaults to 2):
87
+ The number of layers per block.
88
+ downsample_padding (`int`, defaults to 1):
89
+ The padding to use for the downsampling convolution.
90
+ mid_block_scale_factor (`float`, defaults to 1):
91
+ The scale factor to use for the mid block.
92
+ act_fn (`str`, defaults to "silu"):
93
+ The activation function to use.
94
+ norm_num_groups (`int`, *optional*, defaults to 32):
95
+ The number of groups to use for the normalization. If None, normalization and activation layers is skipped
96
+ in post-processing.
97
+ norm_eps (`float`, defaults to 1e-5):
98
+ The epsilon to use for the normalization.
99
+ cross_attention_dim (`int`, defaults to 1280):
100
+ The dimension of the cross attention features.
101
+ transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
102
+ The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
103
+ [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
104
+ [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
105
+ encoder_hid_dim (`int`, *optional*, defaults to None):
106
+ If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
107
+ dimension to `cross_attention_dim`.
108
+ encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
109
+ If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
110
+ embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
111
+ attention_head_dim (`Union[int, Tuple[int]]`, defaults to 8):
112
+ The dimension of the attention heads.
113
+ use_linear_projection (`bool`, defaults to `False`):
114
+ class_embed_type (`str`, *optional*, defaults to `None`):
115
+ The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None,
116
+ `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
117
+ addition_embed_type (`str`, *optional*, defaults to `None`):
118
+ Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
119
+ "text". "text" will use the `TextTimeEmbedding` layer.
120
+ num_class_embeds (`int`, *optional*, defaults to 0):
121
+ Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
122
+ class conditioning with `class_embed_type` equal to `None`.
123
+ upcast_attention (`bool`, defaults to `False`):
124
+ resnet_time_scale_shift (`str`, defaults to `"default"`):
125
+ Time scale shift config for ResNet blocks (see `ResnetBlock2D`). Choose from `default` or `scale_shift`.
126
+ projection_class_embeddings_input_dim (`int`, *optional*, defaults to `None`):
127
+ The dimension of the `class_labels` input when `class_embed_type="projection"`. Required when
128
+ `class_embed_type="projection"`.
129
+ brushnet_conditioning_channel_order (`str`, defaults to `"rgb"`):
130
+ The channel order of conditional image. Will convert to `rgb` if it's `bgr`.
131
+ conditioning_embedding_out_channels (`tuple[int]`, *optional*, defaults to `(16, 32, 96, 256)`):
132
+ The tuple of output channel for each block in the `conditioning_embedding` layer.
133
+ global_pool_conditions (`bool`, defaults to `False`):
134
+ TODO(Patrick) - unused parameter.
135
+ addition_embed_type_num_heads (`int`, defaults to 64):
136
+ The number of heads to use for the `TextTimeEmbedding` layer.
137
+ """
138
+
139
+ _supports_gradient_checkpointing = True
140
+
141
+ @register_to_config
142
+ def __init__(
143
+ self,
144
+ in_channels: int = 4,
145
+ conditioning_channels: int = 5,
146
+ flip_sin_to_cos: bool = True,
147
+ freq_shift: int = 0,
148
+ down_block_types: Tuple[str, ...] = (
149
+ "CrossAttnDownBlock2D",
150
+ "CrossAttnDownBlock2D",
151
+ "CrossAttnDownBlock2D",
152
+ "DownBlock2D",
153
+ ),
154
+ mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
155
+ up_block_types: Tuple[str, ...] = (
156
+ "UpBlock2D",
157
+ "CrossAttnUpBlock2D",
158
+ "CrossAttnUpBlock2D",
159
+ "CrossAttnUpBlock2D",
160
+ ),
161
+ only_cross_attention: Union[bool, Tuple[bool]] = False,
162
+ block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
163
+ layers_per_block: int = 2,
164
+ downsample_padding: int = 1,
165
+ mid_block_scale_factor: float = 1,
166
+ act_fn: str = "silu",
167
+ norm_num_groups: Optional[int] = 32,
168
+ norm_eps: float = 1e-5,
169
+ cross_attention_dim: int = 1280,
170
+ transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
171
+ encoder_hid_dim: Optional[int] = None,
172
+ encoder_hid_dim_type: Optional[str] = None,
173
+ attention_head_dim: Union[int, Tuple[int, ...]] = 8,
174
+ num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
175
+ use_linear_projection: bool = False,
176
+ class_embed_type: Optional[str] = None,
177
+ addition_embed_type: Optional[str] = None,
178
+ addition_time_embed_dim: Optional[int] = None,
179
+ num_class_embeds: Optional[int] = None,
180
+ upcast_attention: bool = False,
181
+ resnet_time_scale_shift: str = "default",
182
+ projection_class_embeddings_input_dim: Optional[int] = None,
183
+ brushnet_conditioning_channel_order: str = "rgb",
184
+ conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
185
+ global_pool_conditions: bool = False,
186
+ addition_embed_type_num_heads: int = 64,
187
+ ):
188
+ super().__init__()
189
+
190
+ # If `num_attention_heads` is not defined (which is the case for most models)
191
+ # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
192
+ # The reason for this behavior is to correct for incorrectly named variables that were introduced
193
+ # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
194
+ # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
195
+ # which is why we correct for the naming here.
196
+ num_attention_heads = num_attention_heads or attention_head_dim
197
+
198
+ # Check inputs
199
+ if len(down_block_types) != len(up_block_types):
200
+ raise ValueError(
201
+ f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
202
+ )
203
+
204
+ if len(block_out_channels) != len(down_block_types):
205
+ raise ValueError(
206
+ f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
207
+ )
208
+
209
+ if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
210
+ raise ValueError(
211
+ f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
212
+ )
213
+
214
+ if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
215
+ raise ValueError(
216
+ f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
217
+ )
218
+
219
+ if isinstance(transformer_layers_per_block, int):
220
+ transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
221
+
222
+ # input
223
+ conv_in_kernel = 3
224
+ conv_in_padding = (conv_in_kernel - 1) // 2
225
+ self.conv_in_condition = nn.Conv2d(
226
+ in_channels + conditioning_channels,
227
+ block_out_channels[0],
228
+ kernel_size=conv_in_kernel,
229
+ padding=conv_in_padding,
230
+ )
231
+
232
+ # time
233
+ time_embed_dim = block_out_channels[0] * 4
234
+ self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
235
+ timestep_input_dim = block_out_channels[0]
236
+ self.time_embedding = TimestepEmbedding(
237
+ timestep_input_dim,
238
+ time_embed_dim,
239
+ act_fn=act_fn,
240
+ )
241
+
242
+ if encoder_hid_dim_type is None and encoder_hid_dim is not None:
243
+ encoder_hid_dim_type = "text_proj"
244
+ self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
245
+ logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
246
+
247
+ if encoder_hid_dim is None and encoder_hid_dim_type is not None:
248
+ raise ValueError(
249
+ f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
250
+ )
251
+
252
+ if encoder_hid_dim_type == "text_proj":
253
+ self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
254
+ elif encoder_hid_dim_type == "text_image_proj":
255
+ # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
256
+ # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
257
+ # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
258
+ self.encoder_hid_proj = TextImageProjection(
259
+ text_embed_dim=encoder_hid_dim,
260
+ image_embed_dim=cross_attention_dim,
261
+ cross_attention_dim=cross_attention_dim,
262
+ )
263
+
264
+ elif encoder_hid_dim_type is not None:
265
+ raise ValueError(
266
+ f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
267
+ )
268
+ else:
269
+ self.encoder_hid_proj = None
270
+
271
+ # class embedding
272
+ if class_embed_type is None and num_class_embeds is not None:
273
+ self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
274
+ elif class_embed_type == "timestep":
275
+ self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
276
+ elif class_embed_type == "identity":
277
+ self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
278
+ elif class_embed_type == "projection":
279
+ if projection_class_embeddings_input_dim is None:
280
+ raise ValueError(
281
+ "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
282
+ )
283
+ # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
284
+ # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
285
+ # 2. it projects from an arbitrary input dimension.
286
+ #
287
+ # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
288
+ # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
289
+ # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
290
+ self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
291
+ else:
292
+ self.class_embedding = None
293
+
294
+ if addition_embed_type == "text":
295
+ if encoder_hid_dim is not None:
296
+ text_time_embedding_from_dim = encoder_hid_dim
297
+ else:
298
+ text_time_embedding_from_dim = cross_attention_dim
299
+
300
+ self.add_embedding = TextTimeEmbedding(
301
+ text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
302
+ )
303
+ elif addition_embed_type == "text_image":
304
+ # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
305
+ # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
306
+ # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`
307
+ self.add_embedding = TextImageTimeEmbedding(
308
+ text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
309
+ )
310
+ elif addition_embed_type == "text_time":
311
+ self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
312
+ self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
313
+
314
+ elif addition_embed_type is not None:
315
+ raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
316
+
317
+ self.down_blocks = nn.ModuleList([])
318
+ self.brushnet_down_blocks = nn.ModuleList([])
319
+
320
+ if isinstance(only_cross_attention, bool):
321
+ only_cross_attention = [only_cross_attention] * len(down_block_types)
322
+
323
+ if isinstance(attention_head_dim, int):
324
+ attention_head_dim = (attention_head_dim,) * len(down_block_types)
325
+
326
+ if isinstance(num_attention_heads, int):
327
+ num_attention_heads = (num_attention_heads,) * len(down_block_types)
328
+
329
+ # down
330
+ output_channel = block_out_channels[0]
331
+
332
+ brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
333
+ brushnet_block = zero_module(brushnet_block)
334
+ self.brushnet_down_blocks.append(brushnet_block)
335
+
336
+ for i, down_block_type in enumerate(down_block_types):
337
+ input_channel = output_channel
338
+ output_channel = block_out_channels[i]
339
+ is_final_block = i == len(block_out_channels) - 1
340
+
341
+ down_block = get_down_block(
342
+ down_block_type,
343
+ num_layers=layers_per_block,
344
+ transformer_layers_per_block=transformer_layers_per_block[i],
345
+ in_channels=input_channel,
346
+ out_channels=output_channel,
347
+ temb_channels=time_embed_dim,
348
+ add_downsample=not is_final_block,
349
+ resnet_eps=norm_eps,
350
+ resnet_act_fn=act_fn,
351
+ resnet_groups=norm_num_groups,
352
+ cross_attention_dim=cross_attention_dim,
353
+ num_attention_heads=num_attention_heads[i],
354
+ attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
355
+ downsample_padding=downsample_padding,
356
+ use_linear_projection=use_linear_projection,
357
+ only_cross_attention=only_cross_attention[i],
358
+ upcast_attention=upcast_attention,
359
+ resnet_time_scale_shift=resnet_time_scale_shift,
360
+ )
361
+ self.down_blocks.append(down_block)
362
+
363
+ for _ in range(layers_per_block):
364
+ brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
365
+ brushnet_block = zero_module(brushnet_block)
366
+ self.brushnet_down_blocks.append(brushnet_block)
367
+
368
+ if not is_final_block:
369
+ brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
370
+ brushnet_block = zero_module(brushnet_block)
371
+ self.brushnet_down_blocks.append(brushnet_block)
372
+
373
+ # mid
374
+ mid_block_channel = block_out_channels[-1]
375
+
376
+ brushnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1)
377
+ brushnet_block = zero_module(brushnet_block)
378
+ self.brushnet_mid_block = brushnet_block
379
+
380
+ self.mid_block = get_mid_block(
381
+ mid_block_type,
382
+ transformer_layers_per_block=transformer_layers_per_block[-1],
383
+ in_channels=mid_block_channel,
384
+ temb_channels=time_embed_dim,
385
+ resnet_eps=norm_eps,
386
+ resnet_act_fn=act_fn,
387
+ output_scale_factor=mid_block_scale_factor,
388
+ resnet_time_scale_shift=resnet_time_scale_shift,
389
+ cross_attention_dim=cross_attention_dim,
390
+ num_attention_heads=num_attention_heads[-1],
391
+ resnet_groups=norm_num_groups,
392
+ use_linear_projection=use_linear_projection,
393
+ upcast_attention=upcast_attention,
394
+ )
395
+
396
+ # count how many layers upsample the images
397
+ self.num_upsamplers = 0
398
+
399
+ # up
400
+ reversed_block_out_channels = list(reversed(block_out_channels))
401
+ reversed_num_attention_heads = list(reversed(num_attention_heads))
402
+ reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
403
+ only_cross_attention = list(reversed(only_cross_attention))
404
+
405
+ output_channel = reversed_block_out_channels[0]
406
+
407
+ self.up_blocks = nn.ModuleList([])
408
+ self.brushnet_up_blocks = nn.ModuleList([])
409
+
410
+ for i, up_block_type in enumerate(up_block_types):
411
+ is_final_block = i == len(block_out_channels) - 1
412
+
413
+ prev_output_channel = output_channel
414
+ output_channel = reversed_block_out_channels[i]
415
+ input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
416
+
417
+ # add upsample block for all BUT final layer
418
+ if not is_final_block:
419
+ add_upsample = True
420
+ self.num_upsamplers += 1
421
+ else:
422
+ add_upsample = False
423
+
424
+ up_block = get_up_block(
425
+ up_block_type,
426
+ num_layers=layers_per_block + 1,
427
+ transformer_layers_per_block=reversed_transformer_layers_per_block[i],
428
+ in_channels=input_channel,
429
+ out_channels=output_channel,
430
+ prev_output_channel=prev_output_channel,
431
+ temb_channels=time_embed_dim,
432
+ add_upsample=add_upsample,
433
+ resnet_eps=norm_eps,
434
+ resnet_act_fn=act_fn,
435
+ resolution_idx=i,
436
+ resnet_groups=norm_num_groups,
437
+ cross_attention_dim=cross_attention_dim,
438
+ num_attention_heads=reversed_num_attention_heads[i],
439
+ use_linear_projection=use_linear_projection,
440
+ only_cross_attention=only_cross_attention[i],
441
+ upcast_attention=upcast_attention,
442
+ resnet_time_scale_shift=resnet_time_scale_shift,
443
+ attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
444
+ )
445
+ self.up_blocks.append(up_block)
446
+ prev_output_channel = output_channel
447
+
448
+ for _ in range(layers_per_block + 1):
449
+ brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
450
+ brushnet_block = zero_module(brushnet_block)
451
+ self.brushnet_up_blocks.append(brushnet_block)
452
+
453
+ if not is_final_block:
454
+ brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
455
+ brushnet_block = zero_module(brushnet_block)
456
+ self.brushnet_up_blocks.append(brushnet_block)
457
+
458
+ @classmethod
459
+ def from_unet(
460
+ cls,
461
+ unet: UNet2DConditionModel,
462
+ brushnet_conditioning_channel_order: str = "rgb",
463
+ conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
464
+ load_weights_from_unet: bool = True,
465
+ conditioning_channels: int = 5,
466
+ ):
467
+ r"""
468
+ Instantiate a [`BrushNetModel`] from [`UNet2DConditionModel`].
469
+
470
+ Parameters:
471
+ unet (`UNet2DConditionModel`):
472
+ The UNet model weights to copy to the [`BrushNetModel`]. All configuration options are also copied
473
+ where applicable.
474
+ """
475
+ transformer_layers_per_block = (
476
+ unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1
477
+ )
478
+ encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None
479
+ encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None
480
+ addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None
481
+ addition_time_embed_dim = (
482
+ unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None
483
+ )
484
+
485
+ brushnet = cls(
486
+ in_channels=unet.config.in_channels,
487
+ conditioning_channels=conditioning_channels,
488
+ flip_sin_to_cos=unet.config.flip_sin_to_cos,
489
+ freq_shift=unet.config.freq_shift,
490
+ # down_block_types=['DownBlock2D','DownBlock2D','DownBlock2D','DownBlock2D'],
491
+ down_block_types=[
492
+ "CrossAttnDownBlock2D",
493
+ "CrossAttnDownBlock2D",
494
+ "CrossAttnDownBlock2D",
495
+ "DownBlock2D",
496
+ ],
497
+ # mid_block_type='MidBlock2D',
498
+ mid_block_type="UNetMidBlock2DCrossAttn",
499
+ # up_block_types=['UpBlock2D','UpBlock2D','UpBlock2D','UpBlock2D'],
500
+ up_block_types=["UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"],
501
+ only_cross_attention=unet.config.only_cross_attention,
502
+ block_out_channels=unet.config.block_out_channels,
503
+ layers_per_block=unet.config.layers_per_block,
504
+ downsample_padding=unet.config.downsample_padding,
505
+ mid_block_scale_factor=unet.config.mid_block_scale_factor,
506
+ act_fn=unet.config.act_fn,
507
+ norm_num_groups=unet.config.norm_num_groups,
508
+ norm_eps=unet.config.norm_eps,
509
+ cross_attention_dim=unet.config.cross_attention_dim,
510
+ transformer_layers_per_block=transformer_layers_per_block,
511
+ encoder_hid_dim=encoder_hid_dim,
512
+ encoder_hid_dim_type=encoder_hid_dim_type,
513
+ attention_head_dim=unet.config.attention_head_dim,
514
+ num_attention_heads=unet.config.num_attention_heads,
515
+ use_linear_projection=unet.config.use_linear_projection,
516
+ class_embed_type=unet.config.class_embed_type,
517
+ addition_embed_type=addition_embed_type,
518
+ addition_time_embed_dim=addition_time_embed_dim,
519
+ num_class_embeds=unet.config.num_class_embeds,
520
+ upcast_attention=unet.config.upcast_attention,
521
+ resnet_time_scale_shift=unet.config.resnet_time_scale_shift,
522
+ projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,
523
+ brushnet_conditioning_channel_order=brushnet_conditioning_channel_order,
524
+ conditioning_embedding_out_channels=conditioning_embedding_out_channels,
525
+ )
526
+
527
+ if load_weights_from_unet:
528
+ conv_in_condition_weight = torch.zeros_like(brushnet.conv_in_condition.weight)
529
+ conv_in_condition_weight[:, :4, ...] = unet.conv_in.weight
530
+ conv_in_condition_weight[:, 4:8, ...] = unet.conv_in.weight
531
+ brushnet.conv_in_condition.weight = torch.nn.Parameter(conv_in_condition_weight)
532
+ brushnet.conv_in_condition.bias = unet.conv_in.bias
533
+
534
+ brushnet.time_proj.load_state_dict(unet.time_proj.state_dict())
535
+ brushnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())
536
+
537
+ if brushnet.class_embedding:
538
+ brushnet.class_embedding.load_state_dict(unet.class_embedding.state_dict())
539
+
540
+ brushnet.down_blocks.load_state_dict(unet.down_blocks.state_dict(), strict=False)
541
+ brushnet.mid_block.load_state_dict(unet.mid_block.state_dict(), strict=False)
542
+ brushnet.up_blocks.load_state_dict(unet.up_blocks.state_dict(), strict=False)
543
+
544
+ return brushnet.to(unet.dtype)
545
+
546
+ @property
547
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
548
+ def attn_processors(self) -> Dict[str, AttentionProcessor]:
549
+ r"""
550
+ Returns:
551
+ `dict` of attention processors: A dictionary containing all attention processors used in the model with
552
+ indexed by its weight name.
553
+ """
554
+ # set recursively
555
+ processors = {}
556
+
557
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
558
+ if hasattr(module, "get_processor"):
559
+ processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
560
+
561
+ for sub_name, child in module.named_children():
562
+ fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
563
+
564
+ return processors
565
+
566
+ for name, module in self.named_children():
567
+ fn_recursive_add_processors(name, module, processors)
568
+
569
+ return processors
570
+
571
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
572
+ def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
573
+ r"""
574
+ Sets the attention processor to use to compute attention.
575
+
576
+ Parameters:
577
+ processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
578
+ The instantiated processor class or a dictionary of processor classes that will be set as the processor
579
+ for **all** `Attention` layers.
580
+
581
+ If `processor` is a dict, the key needs to define the path to the corresponding cross attention
582
+ processor. This is strongly recommended when setting trainable attention processors.
583
+
584
+ """
585
+ count = len(self.attn_processors.keys())
586
+
587
+ if isinstance(processor, dict) and len(processor) != count:
588
+ raise ValueError(
589
+ f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
590
+ f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
591
+ )
592
+
593
+ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
594
+ if hasattr(module, "set_processor"):
595
+ if not isinstance(processor, dict):
596
+ module.set_processor(processor)
597
+ else:
598
+ module.set_processor(processor.pop(f"{name}.processor"))
599
+
600
+ for sub_name, child in module.named_children():
601
+ fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
602
+
603
+ for name, module in self.named_children():
604
+ fn_recursive_attn_processor(name, module, processor)
605
+
606
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
607
+ def set_default_attn_processor(self):
608
+ """
609
+ Disables custom attention processors and sets the default attention implementation.
610
+ """
611
+ if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
612
+ processor = AttnAddedKVProcessor()
613
+ elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
614
+ processor = AttnProcessor()
615
+ else:
616
+ raise ValueError(
617
+ f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
618
+ )
619
+
620
+ self.set_attn_processor(processor)
621
+
622
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice
623
+ def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None:
624
+ r"""
625
+ Enable sliced attention computation.
626
+
627
+ When this option is enabled, the attention module splits the input tensor in slices to compute attention in
628
+ several steps. This is useful for saving some memory in exchange for a small decrease in speed.
629
+
630
+ Args:
631
+ slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
632
+ When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
633
+ `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
634
+ provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
635
+ must be a multiple of `slice_size`.
636
+ """
637
+ sliceable_head_dims = []
638
+
639
+ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
640
+ if hasattr(module, "set_attention_slice"):
641
+ sliceable_head_dims.append(module.sliceable_head_dim)
642
+
643
+ for child in module.children():
644
+ fn_recursive_retrieve_sliceable_dims(child)
645
+
646
+ # retrieve number of attention layers
647
+ for module in self.children():
648
+ fn_recursive_retrieve_sliceable_dims(module)
649
+
650
+ num_sliceable_layers = len(sliceable_head_dims)
651
+
652
+ if slice_size == "auto":
653
+ # half the attention head size is usually a good trade-off between
654
+ # speed and memory
655
+ slice_size = [dim // 2 for dim in sliceable_head_dims]
656
+ elif slice_size == "max":
657
+ # make smallest slice possible
658
+ slice_size = num_sliceable_layers * [1]
659
+
660
+ slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
661
+
662
+ if len(slice_size) != len(sliceable_head_dims):
663
+ raise ValueError(
664
+ f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
665
+ f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
666
+ )
667
+
668
+ for i in range(len(slice_size)):
669
+ size = slice_size[i]
670
+ dim = sliceable_head_dims[i]
671
+ if size is not None and size > dim:
672
+ raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
673
+
674
+ # Recursively walk through all the children.
675
+ # Any children which exposes the set_attention_slice method
676
+ # gets the message
677
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
678
+ if hasattr(module, "set_attention_slice"):
679
+ module.set_attention_slice(slice_size.pop())
680
+
681
+ for child in module.children():
682
+ fn_recursive_set_attention_slice(child, slice_size)
683
+
684
+ reversed_slice_size = list(reversed(slice_size))
685
+ for module in self.children():
686
+ fn_recursive_set_attention_slice(module, reversed_slice_size)
687
+
688
+ def _set_gradient_checkpointing(self, module, value: bool = False) -> None:
689
+ if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):
690
+ module.gradient_checkpointing = value
691
+
692
+ def forward(
693
+ self,
694
+ sample: torch.FloatTensor,
695
+ timestep: Union[torch.Tensor, float, int],
696
+ encoder_hidden_states: torch.Tensor,
697
+ brushnet_cond: torch.FloatTensor,
698
+ conditioning_scale: float = 1.0,
699
+ class_labels: Optional[torch.Tensor] = None,
700
+ timestep_cond: Optional[torch.Tensor] = None,
701
+ attention_mask: Optional[torch.Tensor] = None,
702
+ added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
703
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
704
+ guess_mode: bool = False,
705
+ return_dict: bool = True,
706
+ ) -> Union[BrushNetOutput, Tuple[Tuple[torch.FloatTensor, ...], torch.FloatTensor]]:
707
+ """
708
+ The [`BrushNetModel`] forward method.
709
+
710
+ Args:
711
+ sample (`torch.FloatTensor`):
712
+ The noisy input tensor.
713
+ timestep (`Union[torch.Tensor, float, int]`):
714
+ The number of timesteps to denoise an input.
715
+ encoder_hidden_states (`torch.Tensor`):
716
+ The encoder hidden states.
717
+ brushnet_cond (`torch.FloatTensor`):
718
+ The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`.
719
+ conditioning_scale (`float`, defaults to `1.0`):
720
+ The scale factor for BrushNet outputs.
721
+ class_labels (`torch.Tensor`, *optional*, defaults to `None`):
722
+ Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
723
+ timestep_cond (`torch.Tensor`, *optional*, defaults to `None`):
724
+ Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the
725
+ timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep
726
+ embeddings.
727
+ attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
728
+ An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
729
+ is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
730
+ negative values to the attention scores corresponding to "discard" tokens.
731
+ added_cond_kwargs (`dict`):
732
+ Additional conditions for the Stable Diffusion XL UNet.
733
+ cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`):
734
+ A kwargs dictionary that if specified is passed along to the `AttnProcessor`.
735
+ guess_mode (`bool`, defaults to `False`):
736
+ In this mode, the BrushNet encoder tries its best to recognize the input content of the input even if
737
+ you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended.
738
+ return_dict (`bool`, defaults to `True`):
739
+ Whether or not to return a [`~models.brushnet.BrushNetOutput`] instead of a plain tuple.
740
+
741
+ Returns:
742
+ [`~models.brushnet.BrushNetOutput`] **or** `tuple`:
743
+ If `return_dict` is `True`, a [`~models.brushnet.BrushNetOutput`] is returned, otherwise a tuple is
744
+ returned where the first element is the sample tensor.
745
+ """
746
+ # check channel order
747
+ channel_order = self.config.brushnet_conditioning_channel_order
748
+
749
+ if channel_order == "rgb":
750
+ # in rgb order by default
751
+ ...
752
+ elif channel_order == "bgr":
753
+ brushnet_cond = torch.flip(brushnet_cond, dims=[1])
754
+ else:
755
+ raise ValueError(f"unknown `brushnet_conditioning_channel_order`: {channel_order}")
756
+
757
+ # prepare attention_mask
758
+ if attention_mask is not None:
759
+ attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
760
+ attention_mask = attention_mask.unsqueeze(1)
761
+
762
+ # 1. time
763
+ timesteps = timestep
764
+ if not torch.is_tensor(timesteps):
765
+ # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
766
+ # This would be a good case for the `match` statement (Python 3.10+)
767
+ is_mps = sample.device.type == "mps"
768
+ if isinstance(timestep, float):
769
+ dtype = torch.float32 if is_mps else torch.float64
770
+ else:
771
+ dtype = torch.int32 if is_mps else torch.int64
772
+ timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
773
+ elif len(timesteps.shape) == 0:
774
+ timesteps = timesteps[None].to(sample.device)
775
+
776
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
777
+ timesteps = timesteps.expand(sample.shape[0])
778
+
779
+ t_emb = self.time_proj(timesteps)
780
+
781
+ # timesteps does not contain any weights and will always return f32 tensors
782
+ # but time_embedding might actually be running in fp16. so we need to cast here.
783
+ # there might be better ways to encapsulate this.
784
+ t_emb = t_emb.to(dtype=sample.dtype)
785
+
786
+ emb = self.time_embedding(t_emb, timestep_cond)
787
+ aug_emb = None
788
+
789
+ if self.class_embedding is not None:
790
+ if class_labels is None:
791
+ raise ValueError("class_labels should be provided when num_class_embeds > 0")
792
+
793
+ if self.config.class_embed_type == "timestep":
794
+ class_labels = self.time_proj(class_labels)
795
+
796
+ class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
797
+ emb = emb + class_emb
798
+
799
+ if self.config.addition_embed_type is not None:
800
+ if self.config.addition_embed_type == "text":
801
+ aug_emb = self.add_embedding(encoder_hidden_states)
802
+
803
+ elif self.config.addition_embed_type == "text_time":
804
+ if "text_embeds" not in added_cond_kwargs:
805
+ raise ValueError(
806
+ f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
807
+ )
808
+ text_embeds = added_cond_kwargs.get("text_embeds")
809
+ if "time_ids" not in added_cond_kwargs:
810
+ raise ValueError(
811
+ f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
812
+ )
813
+ time_ids = added_cond_kwargs.get("time_ids")
814
+ time_embeds = self.add_time_proj(time_ids.flatten())
815
+ time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
816
+
817
+ add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
818
+ add_embeds = add_embeds.to(emb.dtype)
819
+ aug_emb = self.add_embedding(add_embeds)
820
+
821
+ emb = emb + aug_emb if aug_emb is not None else emb
822
+
823
+ # 2. pre-process
824
+ brushnet_cond = torch.concat([sample, brushnet_cond], 1)
825
+ sample = self.conv_in_condition(brushnet_cond)
826
+
827
+ # 3. down
828
+ down_block_res_samples = (sample,)
829
+ for downsample_block in self.down_blocks:
830
+ if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
831
+ sample, res_samples = downsample_block(
832
+ hidden_states=sample,
833
+ temb=emb,
834
+ encoder_hidden_states=encoder_hidden_states,
835
+ attention_mask=attention_mask,
836
+ cross_attention_kwargs=cross_attention_kwargs,
837
+ )
838
+ else:
839
+ sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
840
+
841
+ down_block_res_samples += res_samples
842
+
843
+ # 4. PaintingNet down blocks
844
+ brushnet_down_block_res_samples = ()
845
+ for down_block_res_sample, brushnet_down_block in zip(down_block_res_samples, self.brushnet_down_blocks):
846
+ down_block_res_sample = brushnet_down_block(down_block_res_sample)
847
+ brushnet_down_block_res_samples = brushnet_down_block_res_samples + (down_block_res_sample,)
848
+
849
+ # 5. mid
850
+ if self.mid_block is not None:
851
+ if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
852
+ sample = self.mid_block(
853
+ sample,
854
+ emb,
855
+ encoder_hidden_states=encoder_hidden_states,
856
+ attention_mask=attention_mask,
857
+ cross_attention_kwargs=cross_attention_kwargs,
858
+ )
859
+ else:
860
+ sample = self.mid_block(sample, emb)
861
+
862
+ # 6. BrushNet mid blocks
863
+ brushnet_mid_block_res_sample = self.brushnet_mid_block(sample)
864
+
865
+ # 7. up
866
+ up_block_res_samples = ()
867
+ for i, upsample_block in enumerate(self.up_blocks):
868
+ is_final_block = i == len(self.up_blocks) - 1
869
+
870
+ res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
871
+ down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
872
+
873
+ # if we have not reached the final block and need to forward the
874
+ # upsample size, we do it here
875
+ if not is_final_block:
876
+ upsample_size = down_block_res_samples[-1].shape[2:]
877
+
878
+ if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
879
+ sample, up_res_samples = upsample_block(
880
+ hidden_states=sample,
881
+ temb=emb,
882
+ res_hidden_states_tuple=res_samples,
883
+ encoder_hidden_states=encoder_hidden_states,
884
+ cross_attention_kwargs=cross_attention_kwargs,
885
+ upsample_size=upsample_size,
886
+ attention_mask=attention_mask,
887
+ return_res_samples=True,
888
+ )
889
+ else:
890
+ sample, up_res_samples = upsample_block(
891
+ hidden_states=sample,
892
+ temb=emb,
893
+ res_hidden_states_tuple=res_samples,
894
+ upsample_size=upsample_size,
895
+ return_res_samples=True,
896
+ )
897
+
898
+ up_block_res_samples += up_res_samples
899
+
900
+ # 8. BrushNet up blocks
901
+ brushnet_up_block_res_samples = ()
902
+ for up_block_res_sample, brushnet_up_block in zip(up_block_res_samples, self.brushnet_up_blocks):
903
+ up_block_res_sample = brushnet_up_block(up_block_res_sample)
904
+ brushnet_up_block_res_samples = brushnet_up_block_res_samples + (up_block_res_sample,)
905
+
906
+ # 6. scaling
907
+ if guess_mode and not self.config.global_pool_conditions:
908
+ scales = torch.logspace(
909
+ -1,
910
+ 0,
911
+ len(brushnet_down_block_res_samples) + 1 + len(brushnet_up_block_res_samples),
912
+ device=sample.device,
913
+ ) # 0.1 to 1.0
914
+ scales = scales * conditioning_scale
915
+
916
+ brushnet_down_block_res_samples = [
917
+ sample * scale
918
+ for sample, scale in zip(
919
+ brushnet_down_block_res_samples, scales[: len(brushnet_down_block_res_samples)]
920
+ )
921
+ ]
922
+ brushnet_mid_block_res_sample = (
923
+ brushnet_mid_block_res_sample * scales[len(brushnet_down_block_res_samples)]
924
+ )
925
+ brushnet_up_block_res_samples = [
926
+ sample * scale
927
+ for sample, scale in zip(
928
+ brushnet_up_block_res_samples, scales[len(brushnet_down_block_res_samples) + 1 :]
929
+ )
930
+ ]
931
+ else:
932
+ brushnet_down_block_res_samples = [
933
+ sample * conditioning_scale for sample in brushnet_down_block_res_samples
934
+ ]
935
+ brushnet_mid_block_res_sample = brushnet_mid_block_res_sample * conditioning_scale
936
+ brushnet_up_block_res_samples = [sample * conditioning_scale for sample in brushnet_up_block_res_samples]
937
+
938
+ if self.config.global_pool_conditions:
939
+ brushnet_down_block_res_samples = [
940
+ torch.mean(sample, dim=(2, 3), keepdim=True) for sample in brushnet_down_block_res_samples
941
+ ]
942
+ brushnet_mid_block_res_sample = torch.mean(brushnet_mid_block_res_sample, dim=(2, 3), keepdim=True)
943
+ brushnet_up_block_res_samples = [
944
+ torch.mean(sample, dim=(2, 3), keepdim=True) for sample in brushnet_up_block_res_samples
945
+ ]
946
+
947
+ if not return_dict:
948
+ return (brushnet_down_block_res_samples, brushnet_mid_block_res_sample, brushnet_up_block_res_samples)
949
+
950
+ return BrushNetOutput(
951
+ down_block_res_samples=brushnet_down_block_res_samples,
952
+ mid_block_res_sample=brushnet_mid_block_res_sample,
953
+ up_block_res_samples=brushnet_up_block_res_samples,
954
+ )
955
+
956
+
957
+ def zero_module(module):
958
+ for p in module.parameters():
959
+ nn.init.zeros_(p)
960
+ return module
model/__init__.py ADDED
File without changes
model/__pycache__/BrushNet_CA.cpython-310.pyc ADDED
Binary file (29.1 kB). View file
 
model/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (130 Bytes). View file
 
model/diffusers_c/__init__.py ADDED
@@ -0,0 +1,789 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ __version__ = "0.27.0.dev0"
2
+
3
+ from typing import TYPE_CHECKING
4
+
5
+ from .utils import (
6
+ DIFFUSERS_SLOW_IMPORT,
7
+ OptionalDependencyNotAvailable,
8
+ _LazyModule,
9
+ is_flax_available,
10
+ is_k_diffusion_available,
11
+ is_librosa_available,
12
+ is_note_seq_available,
13
+ is_onnx_available,
14
+ is_scipy_available,
15
+ is_torch_available,
16
+ is_torchsde_available,
17
+ is_transformers_available,
18
+ )
19
+
20
+
21
+ # Lazy Import based on
22
+ # https://github.com/huggingface/transformers/blob/main/src/transformers/__init__.py
23
+
24
+ # When adding a new object to this init, please add it to `_import_structure`. The `_import_structure` is a dictionary submodule to list of object names,
25
+ # and is used to defer the actual importing for when the objects are requested.
26
+ # This way `import diffusers` provides the names in the namespace without actually importing anything (and especially none of the backends).
27
+
28
+ _import_structure = {
29
+ "configuration_utils": ["ConfigMixin"],
30
+ "models": [],
31
+ "pipelines": [],
32
+ "schedulers": [],
33
+ "utils": [
34
+ "OptionalDependencyNotAvailable",
35
+ "is_flax_available",
36
+ "is_inflect_available",
37
+ "is_invisible_watermark_available",
38
+ "is_k_diffusion_available",
39
+ "is_k_diffusion_version",
40
+ "is_librosa_available",
41
+ "is_note_seq_available",
42
+ "is_onnx_available",
43
+ "is_scipy_available",
44
+ "is_torch_available",
45
+ "is_torchsde_available",
46
+ "is_transformers_available",
47
+ "is_transformers_version",
48
+ "is_unidecode_available",
49
+ "logging",
50
+ ],
51
+ }
52
+
53
+ try:
54
+ if not is_onnx_available():
55
+ raise OptionalDependencyNotAvailable()
56
+ except OptionalDependencyNotAvailable:
57
+ from .utils import dummy_onnx_objects # noqa F403
58
+
59
+ _import_structure["utils.dummy_onnx_objects"] = [
60
+ name for name in dir(dummy_onnx_objects) if not name.startswith("_")
61
+ ]
62
+
63
+ else:
64
+ _import_structure["pipelines"].extend(["OnnxRuntimeModel"])
65
+
66
+ try:
67
+ if not is_torch_available():
68
+ raise OptionalDependencyNotAvailable()
69
+ except OptionalDependencyNotAvailable:
70
+ from .utils import dummy_pt_objects # noqa F403
71
+
72
+ _import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")]
73
+
74
+ else:
75
+ _import_structure["models"].extend(
76
+ [
77
+ "AsymmetricAutoencoderKL",
78
+ "AutoencoderKL",
79
+ "AutoencoderKLTemporalDecoder",
80
+ "AutoencoderTiny",
81
+ "ConsistencyDecoderVAE",
82
+ "BrushNetModel",
83
+ "ControlNetModel",
84
+ "I2VGenXLUNet",
85
+ "Kandinsky3UNet",
86
+ "ModelMixin",
87
+ "MotionAdapter",
88
+ "MultiAdapter",
89
+ "PriorTransformer",
90
+ "StableCascadeUNet",
91
+ "T2IAdapter",
92
+ "T5FilmDecoder",
93
+ "Transformer2DModel",
94
+ "UNet1DModel",
95
+ "UNet2DConditionModel",
96
+ "UNet2DModel",
97
+ "UNet3DConditionModel",
98
+ "UNetMotionModel",
99
+ "UNetSpatioTemporalConditionModel",
100
+ "UVit2DModel",
101
+ "VQModel",
102
+ ]
103
+ )
104
+
105
+ _import_structure["optimization"] = [
106
+ "get_constant_schedule",
107
+ "get_constant_schedule_with_warmup",
108
+ "get_cosine_schedule_with_warmup",
109
+ "get_cosine_with_hard_restarts_schedule_with_warmup",
110
+ "get_linear_schedule_with_warmup",
111
+ "get_polynomial_decay_schedule_with_warmup",
112
+ "get_scheduler",
113
+ ]
114
+ _import_structure["pipelines"].extend(
115
+ [
116
+ "AudioPipelineOutput",
117
+ "AutoPipelineForImage2Image",
118
+ "AutoPipelineForInpainting",
119
+ "AutoPipelineForText2Image",
120
+ "ConsistencyModelPipeline",
121
+ "DanceDiffusionPipeline",
122
+ "DDIMPipeline",
123
+ "DDPMPipeline",
124
+ "DiffusionPipeline",
125
+ "DiTPipeline",
126
+ "ImagePipelineOutput",
127
+ "KarrasVePipeline",
128
+ "LDMPipeline",
129
+ "LDMSuperResolutionPipeline",
130
+ "PNDMPipeline",
131
+ "RePaintPipeline",
132
+ "ScoreSdeVePipeline",
133
+ "StableDiffusionMixin",
134
+ ]
135
+ )
136
+ _import_structure["schedulers"].extend(
137
+ [
138
+ "AmusedScheduler",
139
+ "CMStochasticIterativeScheduler",
140
+ "DDIMInverseScheduler",
141
+ "DDIMParallelScheduler",
142
+ "DDIMScheduler",
143
+ "DDPMParallelScheduler",
144
+ "DDPMScheduler",
145
+ "DDPMWuerstchenScheduler",
146
+ "DEISMultistepScheduler",
147
+ "DPMSolverMultistepInverseScheduler",
148
+ "DPMSolverMultistepScheduler",
149
+ "DPMSolverSinglestepScheduler",
150
+ "EDMDPMSolverMultistepScheduler",
151
+ "EDMEulerScheduler",
152
+ "EulerAncestralDiscreteScheduler",
153
+ "EulerDiscreteScheduler",
154
+ "HeunDiscreteScheduler",
155
+ "IPNDMScheduler",
156
+ "KarrasVeScheduler",
157
+ "KDPM2AncestralDiscreteScheduler",
158
+ "KDPM2DiscreteScheduler",
159
+ "LCMScheduler",
160
+ "PNDMScheduler",
161
+ "RePaintScheduler",
162
+ "SASolverScheduler",
163
+ "SchedulerMixin",
164
+ "ScoreSdeVeScheduler",
165
+ "TCDScheduler",
166
+ "UnCLIPScheduler",
167
+ "UniPCMultistepScheduler",
168
+ "VQDiffusionScheduler",
169
+ ]
170
+ )
171
+ _import_structure["training_utils"] = ["EMAModel"]
172
+
173
+ try:
174
+ if not (is_torch_available() and is_scipy_available()):
175
+ raise OptionalDependencyNotAvailable()
176
+ except OptionalDependencyNotAvailable:
177
+ from .utils import dummy_torch_and_scipy_objects # noqa F403
178
+
179
+ _import_structure["utils.dummy_torch_and_scipy_objects"] = [
180
+ name for name in dir(dummy_torch_and_scipy_objects) if not name.startswith("_")
181
+ ]
182
+
183
+ else:
184
+ _import_structure["schedulers"].extend(["LMSDiscreteScheduler"])
185
+
186
+ try:
187
+ if not (is_torch_available() and is_torchsde_available()):
188
+ raise OptionalDependencyNotAvailable()
189
+ except OptionalDependencyNotAvailable:
190
+ from .utils import dummy_torch_and_torchsde_objects # noqa F403
191
+
192
+ _import_structure["utils.dummy_torch_and_torchsde_objects"] = [
193
+ name for name in dir(dummy_torch_and_torchsde_objects) if not name.startswith("_")
194
+ ]
195
+
196
+ else:
197
+ _import_structure["schedulers"].extend(["DPMSolverSDEScheduler"])
198
+
199
+ try:
200
+ if not (is_torch_available() and is_transformers_available()):
201
+ raise OptionalDependencyNotAvailable()
202
+ except OptionalDependencyNotAvailable:
203
+ from .utils import dummy_torch_and_transformers_objects # noqa F403
204
+
205
+ _import_structure["utils.dummy_torch_and_transformers_objects"] = [
206
+ name for name in dir(dummy_torch_and_transformers_objects) if not name.startswith("_")
207
+ ]
208
+
209
+ else:
210
+ _import_structure["pipelines"].extend(
211
+ [
212
+ "AltDiffusionImg2ImgPipeline",
213
+ "AltDiffusionPipeline",
214
+ "AmusedImg2ImgPipeline",
215
+ "AmusedInpaintPipeline",
216
+ "AmusedPipeline",
217
+ "AnimateDiffPipeline",
218
+ "AnimateDiffVideoToVideoPipeline",
219
+ "AudioLDM2Pipeline",
220
+ "AudioLDM2ProjectionModel",
221
+ "AudioLDM2UNet2DConditionModel",
222
+ "AudioLDMPipeline",
223
+ "BlipDiffusionControlNetPipeline",
224
+ "BlipDiffusionPipeline",
225
+ "CLIPImageProjection",
226
+ "CycleDiffusionPipeline",
227
+ "I2VGenXLPipeline",
228
+ "IFImg2ImgPipeline",
229
+ "IFImg2ImgSuperResolutionPipeline",
230
+ "IFInpaintingPipeline",
231
+ "IFInpaintingSuperResolutionPipeline",
232
+ "IFPipeline",
233
+ "IFSuperResolutionPipeline",
234
+ "ImageTextPipelineOutput",
235
+ "Kandinsky3Img2ImgPipeline",
236
+ "Kandinsky3Pipeline",
237
+ "KandinskyCombinedPipeline",
238
+ "KandinskyImg2ImgCombinedPipeline",
239
+ "KandinskyImg2ImgPipeline",
240
+ "KandinskyInpaintCombinedPipeline",
241
+ "KandinskyInpaintPipeline",
242
+ "KandinskyPipeline",
243
+ "KandinskyPriorPipeline",
244
+ "KandinskyV22CombinedPipeline",
245
+ "KandinskyV22ControlnetImg2ImgPipeline",
246
+ "KandinskyV22ControlnetPipeline",
247
+ "KandinskyV22Img2ImgCombinedPipeline",
248
+ "KandinskyV22Img2ImgPipeline",
249
+ "KandinskyV22InpaintCombinedPipeline",
250
+ "KandinskyV22InpaintPipeline",
251
+ "KandinskyV22Pipeline",
252
+ "KandinskyV22PriorEmb2EmbPipeline",
253
+ "KandinskyV22PriorPipeline",
254
+ "LatentConsistencyModelImg2ImgPipeline",
255
+ "LatentConsistencyModelPipeline",
256
+ "LDMTextToImagePipeline",
257
+ "MusicLDMPipeline",
258
+ "PaintByExamplePipeline",
259
+ "PIAPipeline",
260
+ "PixArtAlphaPipeline",
261
+ "SemanticStableDiffusionPipeline",
262
+ "ShapEImg2ImgPipeline",
263
+ "ShapEPipeline",
264
+ "StableCascadeCombinedPipeline",
265
+ "StableCascadeDecoderPipeline",
266
+ "StableCascadePriorPipeline",
267
+ "StableDiffusionAdapterPipeline",
268
+ "StableDiffusionAttendAndExcitePipeline",
269
+ "StableDiffusionBrushNetPipeline",
270
+ "StableDiffusionBrushNetPowerPaintPipeline",
271
+ "StableDiffusionControlNetImg2ImgPipeline",
272
+ "StableDiffusionControlNetInpaintPipeline",
273
+ "StableDiffusionControlNetPipeline",
274
+ "StableDiffusionDepth2ImgPipeline",
275
+ "StableDiffusionDiffEditPipeline",
276
+ "StableDiffusionGLIGENPipeline",
277
+ "StableDiffusionGLIGENTextImagePipeline",
278
+ "StableDiffusionImageVariationPipeline",
279
+ "StableDiffusionImg2ImgPipeline",
280
+ "StableDiffusionInpaintPipeline",
281
+ "StableDiffusionInpaintPipelineLegacy",
282
+ "StableDiffusionInstructPix2PixPipeline",
283
+ "StableDiffusionLatentUpscalePipeline",
284
+ "StableDiffusionLDM3DPipeline",
285
+ "StableDiffusionModelEditingPipeline",
286
+ "StableDiffusionPanoramaPipeline",
287
+ "StableDiffusionParadigmsPipeline",
288
+ "StableDiffusionPipeline",
289
+ "StableDiffusionPipelineSafe",
290
+ "StableDiffusionPix2PixZeroPipeline",
291
+ "StableDiffusionSAGPipeline",
292
+ "StableDiffusionUpscalePipeline",
293
+ "StableDiffusionXLAdapterPipeline",
294
+ "StableDiffusionXLControlNetImg2ImgPipeline",
295
+ "StableDiffusionXLControlNetInpaintPipeline",
296
+ "StableDiffusionXLControlNetPipeline",
297
+ "StableDiffusionXLImg2ImgPipeline",
298
+ "StableDiffusionXLInpaintPipeline",
299
+ "StableDiffusionXLInstructPix2PixPipeline",
300
+ "StableDiffusionXLPipeline",
301
+ "StableUnCLIPImg2ImgPipeline",
302
+ "StableUnCLIPPipeline",
303
+ "StableVideoDiffusionPipeline",
304
+ "TextToVideoSDPipeline",
305
+ "TextToVideoZeroPipeline",
306
+ "TextToVideoZeroSDXLPipeline",
307
+ "UnCLIPImageVariationPipeline",
308
+ "UnCLIPPipeline",
309
+ "UniDiffuserModel",
310
+ "UniDiffuserPipeline",
311
+ "UniDiffuserTextDecoder",
312
+ "VersatileDiffusionDualGuidedPipeline",
313
+ "VersatileDiffusionImageVariationPipeline",
314
+ "VersatileDiffusionPipeline",
315
+ "VersatileDiffusionTextToImagePipeline",
316
+ "VideoToVideoSDPipeline",
317
+ "VQDiffusionPipeline",
318
+ "WuerstchenCombinedPipeline",
319
+ "WuerstchenDecoderPipeline",
320
+ "WuerstchenPriorPipeline",
321
+ ]
322
+ )
323
+
324
+ try:
325
+ if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
326
+ raise OptionalDependencyNotAvailable()
327
+ except OptionalDependencyNotAvailable:
328
+ from .utils import dummy_torch_and_transformers_and_k_diffusion_objects # noqa F403
329
+
330
+ _import_structure["utils.dummy_torch_and_transformers_and_k_diffusion_objects"] = [
331
+ name for name in dir(dummy_torch_and_transformers_and_k_diffusion_objects) if not name.startswith("_")
332
+ ]
333
+
334
+ else:
335
+ _import_structure["pipelines"].extend(["StableDiffusionKDiffusionPipeline", "StableDiffusionXLKDiffusionPipeline"])
336
+
337
+ try:
338
+ if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
339
+ raise OptionalDependencyNotAvailable()
340
+ except OptionalDependencyNotAvailable:
341
+ from .utils import dummy_torch_and_transformers_and_onnx_objects # noqa F403
342
+
343
+ _import_structure["utils.dummy_torch_and_transformers_and_onnx_objects"] = [
344
+ name for name in dir(dummy_torch_and_transformers_and_onnx_objects) if not name.startswith("_")
345
+ ]
346
+
347
+ else:
348
+ _import_structure["pipelines"].extend(
349
+ [
350
+ "OnnxStableDiffusionImg2ImgPipeline",
351
+ "OnnxStableDiffusionInpaintPipeline",
352
+ "OnnxStableDiffusionInpaintPipelineLegacy",
353
+ "OnnxStableDiffusionPipeline",
354
+ "OnnxStableDiffusionUpscalePipeline",
355
+ "StableDiffusionOnnxPipeline",
356
+ ]
357
+ )
358
+
359
+ try:
360
+ if not (is_torch_available() and is_librosa_available()):
361
+ raise OptionalDependencyNotAvailable()
362
+ except OptionalDependencyNotAvailable:
363
+ from .utils import dummy_torch_and_librosa_objects # noqa F403
364
+
365
+ _import_structure["utils.dummy_torch_and_librosa_objects"] = [
366
+ name for name in dir(dummy_torch_and_librosa_objects) if not name.startswith("_")
367
+ ]
368
+
369
+ else:
370
+ _import_structure["pipelines"].extend(["AudioDiffusionPipeline", "Mel"])
371
+
372
+ try:
373
+ if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
374
+ raise OptionalDependencyNotAvailable()
375
+ except OptionalDependencyNotAvailable:
376
+ from .utils import dummy_transformers_and_torch_and_note_seq_objects # noqa F403
377
+
378
+ _import_structure["utils.dummy_transformers_and_torch_and_note_seq_objects"] = [
379
+ name for name in dir(dummy_transformers_and_torch_and_note_seq_objects) if not name.startswith("_")
380
+ ]
381
+
382
+
383
+ else:
384
+ _import_structure["pipelines"].extend(["SpectrogramDiffusionPipeline"])
385
+
386
+ try:
387
+ if not is_flax_available():
388
+ raise OptionalDependencyNotAvailable()
389
+ except OptionalDependencyNotAvailable:
390
+ from .utils import dummy_flax_objects # noqa F403
391
+
392
+ _import_structure["utils.dummy_flax_objects"] = [
393
+ name for name in dir(dummy_flax_objects) if not name.startswith("_")
394
+ ]
395
+
396
+
397
+ else:
398
+ _import_structure["models.controlnet_flax"] = ["FlaxControlNetModel"]
399
+ _import_structure["models.modeling_flax_utils"] = ["FlaxModelMixin"]
400
+ _import_structure["models.unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"]
401
+ _import_structure["models.vae_flax"] = ["FlaxAutoencoderKL"]
402
+ _import_structure["pipelines"].extend(["FlaxDiffusionPipeline"])
403
+ _import_structure["schedulers"].extend(
404
+ [
405
+ "FlaxDDIMScheduler",
406
+ "FlaxDDPMScheduler",
407
+ "FlaxDPMSolverMultistepScheduler",
408
+ "FlaxEulerDiscreteScheduler",
409
+ "FlaxKarrasVeScheduler",
410
+ "FlaxLMSDiscreteScheduler",
411
+ "FlaxPNDMScheduler",
412
+ "FlaxSchedulerMixin",
413
+ "FlaxScoreSdeVeScheduler",
414
+ ]
415
+ )
416
+
417
+
418
+ try:
419
+ if not (is_flax_available() and is_transformers_available()):
420
+ raise OptionalDependencyNotAvailable()
421
+ except OptionalDependencyNotAvailable:
422
+ from .utils import dummy_flax_and_transformers_objects # noqa F403
423
+
424
+ _import_structure["utils.dummy_flax_and_transformers_objects"] = [
425
+ name for name in dir(dummy_flax_and_transformers_objects) if not name.startswith("_")
426
+ ]
427
+
428
+
429
+ else:
430
+ _import_structure["pipelines"].extend(
431
+ [
432
+ "FlaxStableDiffusionControlNetPipeline",
433
+ "FlaxStableDiffusionImg2ImgPipeline",
434
+ "FlaxStableDiffusionInpaintPipeline",
435
+ "FlaxStableDiffusionPipeline",
436
+ "FlaxStableDiffusionXLPipeline",
437
+ ]
438
+ )
439
+
440
+ try:
441
+ if not (is_note_seq_available()):
442
+ raise OptionalDependencyNotAvailable()
443
+ except OptionalDependencyNotAvailable:
444
+ from .utils import dummy_note_seq_objects # noqa F403
445
+
446
+ _import_structure["utils.dummy_note_seq_objects"] = [
447
+ name for name in dir(dummy_note_seq_objects) if not name.startswith("_")
448
+ ]
449
+
450
+
451
+ else:
452
+ _import_structure["pipelines"].extend(["MidiProcessor"])
453
+
454
+ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
455
+ from .configuration_utils import ConfigMixin
456
+
457
+ try:
458
+ if not is_onnx_available():
459
+ raise OptionalDependencyNotAvailable()
460
+ except OptionalDependencyNotAvailable:
461
+ from .utils.dummy_onnx_objects import * # noqa F403
462
+ else:
463
+ from .pipelines import OnnxRuntimeModel
464
+
465
+ try:
466
+ if not is_torch_available():
467
+ raise OptionalDependencyNotAvailable()
468
+ except OptionalDependencyNotAvailable:
469
+ from .utils.dummy_pt_objects import * # noqa F403
470
+ else:
471
+ from .models import (
472
+ AsymmetricAutoencoderKL,
473
+ AutoencoderKL,
474
+ AutoencoderKLTemporalDecoder,
475
+ AutoencoderTiny,
476
+ BrushNetModel,
477
+ ConsistencyDecoderVAE,
478
+ ControlNetModel,
479
+ I2VGenXLUNet,
480
+ Kandinsky3UNet,
481
+ ModelMixin,
482
+ MotionAdapter,
483
+ MultiAdapter,
484
+ PriorTransformer,
485
+ T2IAdapter,
486
+ T5FilmDecoder,
487
+ Transformer2DModel,
488
+ UNet1DModel,
489
+ UNet2DConditionModel,
490
+ UNet2DModel,
491
+ UNet3DConditionModel,
492
+ UNetMotionModel,
493
+ UNetSpatioTemporalConditionModel,
494
+ UVit2DModel,
495
+ VQModel,
496
+ )
497
+ from .optimization import (
498
+ get_constant_schedule,
499
+ get_constant_schedule_with_warmup,
500
+ get_cosine_schedule_with_warmup,
501
+ get_cosine_with_hard_restarts_schedule_with_warmup,
502
+ get_linear_schedule_with_warmup,
503
+ get_polynomial_decay_schedule_with_warmup,
504
+ get_scheduler,
505
+ )
506
+ from .pipelines import (
507
+ AudioPipelineOutput,
508
+ AutoPipelineForImage2Image,
509
+ AutoPipelineForInpainting,
510
+ AutoPipelineForText2Image,
511
+ BlipDiffusionControlNetPipeline,
512
+ BlipDiffusionPipeline,
513
+ CLIPImageProjection,
514
+ ConsistencyModelPipeline,
515
+ DanceDiffusionPipeline,
516
+ DDIMPipeline,
517
+ DDPMPipeline,
518
+ DiffusionPipeline,
519
+ DiTPipeline,
520
+ ImagePipelineOutput,
521
+ KarrasVePipeline,
522
+ LDMPipeline,
523
+ LDMSuperResolutionPipeline,
524
+ PNDMPipeline,
525
+ RePaintPipeline,
526
+ ScoreSdeVePipeline,
527
+ StableDiffusionMixin,
528
+ )
529
+ from .schedulers import (
530
+ AmusedScheduler,
531
+ CMStochasticIterativeScheduler,
532
+ DDIMInverseScheduler,
533
+ DDIMParallelScheduler,
534
+ DDIMScheduler,
535
+ DDPMParallelScheduler,
536
+ DDPMScheduler,
537
+ DDPMWuerstchenScheduler,
538
+ DEISMultistepScheduler,
539
+ DPMSolverMultistepInverseScheduler,
540
+ DPMSolverMultistepScheduler,
541
+ DPMSolverSinglestepScheduler,
542
+ EDMDPMSolverMultistepScheduler,
543
+ EDMEulerScheduler,
544
+ EulerAncestralDiscreteScheduler,
545
+ EulerDiscreteScheduler,
546
+ HeunDiscreteScheduler,
547
+ IPNDMScheduler,
548
+ KarrasVeScheduler,
549
+ KDPM2AncestralDiscreteScheduler,
550
+ KDPM2DiscreteScheduler,
551
+ LCMScheduler,
552
+ PNDMScheduler,
553
+ RePaintScheduler,
554
+ SASolverScheduler,
555
+ SchedulerMixin,
556
+ ScoreSdeVeScheduler,
557
+ TCDScheduler,
558
+ UnCLIPScheduler,
559
+ UniPCMultistepScheduler,
560
+ VQDiffusionScheduler,
561
+ )
562
+ from .training_utils import EMAModel
563
+
564
+ try:
565
+ if not (is_torch_available() and is_scipy_available()):
566
+ raise OptionalDependencyNotAvailable()
567
+ except OptionalDependencyNotAvailable:
568
+ from .utils.dummy_torch_and_scipy_objects import * # noqa F403
569
+ else:
570
+ from .schedulers import LMSDiscreteScheduler
571
+
572
+ try:
573
+ if not (is_torch_available() and is_torchsde_available()):
574
+ raise OptionalDependencyNotAvailable()
575
+ except OptionalDependencyNotAvailable:
576
+ from .utils.dummy_torch_and_torchsde_objects import * # noqa F403
577
+ else:
578
+ from .schedulers import DPMSolverSDEScheduler
579
+
580
+ try:
581
+ if not (is_torch_available() and is_transformers_available()):
582
+ raise OptionalDependencyNotAvailable()
583
+ except OptionalDependencyNotAvailable:
584
+ from .utils.dummy_torch_and_transformers_objects import * # noqa F403
585
+ else:
586
+ from .pipelines import (
587
+ AltDiffusionImg2ImgPipeline,
588
+ AltDiffusionPipeline,
589
+ AmusedImg2ImgPipeline,
590
+ AmusedInpaintPipeline,
591
+ AmusedPipeline,
592
+ AnimateDiffPipeline,
593
+ AnimateDiffVideoToVideoPipeline,
594
+ AudioLDM2Pipeline,
595
+ AudioLDM2ProjectionModel,
596
+ AudioLDM2UNet2DConditionModel,
597
+ AudioLDMPipeline,
598
+ CLIPImageProjection,
599
+ CycleDiffusionPipeline,
600
+ I2VGenXLPipeline,
601
+ IFImg2ImgPipeline,
602
+ IFImg2ImgSuperResolutionPipeline,
603
+ IFInpaintingPipeline,
604
+ IFInpaintingSuperResolutionPipeline,
605
+ IFPipeline,
606
+ IFSuperResolutionPipeline,
607
+ ImageTextPipelineOutput,
608
+ Kandinsky3Img2ImgPipeline,
609
+ Kandinsky3Pipeline,
610
+ KandinskyCombinedPipeline,
611
+ KandinskyImg2ImgCombinedPipeline,
612
+ KandinskyImg2ImgPipeline,
613
+ KandinskyInpaintCombinedPipeline,
614
+ KandinskyInpaintPipeline,
615
+ KandinskyPipeline,
616
+ KandinskyPriorPipeline,
617
+ KandinskyV22CombinedPipeline,
618
+ KandinskyV22ControlnetImg2ImgPipeline,
619
+ KandinskyV22ControlnetPipeline,
620
+ KandinskyV22Img2ImgCombinedPipeline,
621
+ KandinskyV22Img2ImgPipeline,
622
+ KandinskyV22InpaintCombinedPipeline,
623
+ KandinskyV22InpaintPipeline,
624
+ KandinskyV22Pipeline,
625
+ KandinskyV22PriorEmb2EmbPipeline,
626
+ KandinskyV22PriorPipeline,
627
+ LatentConsistencyModelImg2ImgPipeline,
628
+ LatentConsistencyModelPipeline,
629
+ LDMTextToImagePipeline,
630
+ MusicLDMPipeline,
631
+ PaintByExamplePipeline,
632
+ PIAPipeline,
633
+ PixArtAlphaPipeline,
634
+ SemanticStableDiffusionPipeline,
635
+ ShapEImg2ImgPipeline,
636
+ ShapEPipeline,
637
+ StableCascadeCombinedPipeline,
638
+ StableCascadeDecoderPipeline,
639
+ StableCascadePriorPipeline,
640
+ StableDiffusionAdapterPipeline,
641
+ StableDiffusionAttendAndExcitePipeline,
642
+ StableDiffusionBrushNetPipeline,
643
+ StableDiffusionBrushNetPowerPaintPipeline,
644
+ StableDiffusionControlNetImg2ImgPipeline,
645
+ StableDiffusionControlNetInpaintPipeline,
646
+ StableDiffusionControlNetPipeline,
647
+ StableDiffusionDepth2ImgPipeline,
648
+ StableDiffusionDiffEditPipeline,
649
+ StableDiffusionGLIGENPipeline,
650
+ StableDiffusionGLIGENTextImagePipeline,
651
+ StableDiffusionImageVariationPipeline,
652
+ StableDiffusionImg2ImgPipeline,
653
+ StableDiffusionInpaintPipeline,
654
+ StableDiffusionInpaintPipelineLegacy,
655
+ StableDiffusionInstructPix2PixPipeline,
656
+ StableDiffusionLatentUpscalePipeline,
657
+ StableDiffusionLDM3DPipeline,
658
+ StableDiffusionModelEditingPipeline,
659
+ StableDiffusionPanoramaPipeline,
660
+ StableDiffusionParadigmsPipeline,
661
+ StableDiffusionPipeline,
662
+ StableDiffusionPipelineSafe,
663
+ StableDiffusionPix2PixZeroPipeline,
664
+ StableDiffusionSAGPipeline,
665
+ StableDiffusionUpscalePipeline,
666
+ StableDiffusionXLAdapterPipeline,
667
+ StableDiffusionXLControlNetImg2ImgPipeline,
668
+ StableDiffusionXLControlNetInpaintPipeline,
669
+ StableDiffusionXLControlNetPipeline,
670
+ StableDiffusionXLImg2ImgPipeline,
671
+ StableDiffusionXLInpaintPipeline,
672
+ StableDiffusionXLInstructPix2PixPipeline,
673
+ StableDiffusionXLPipeline,
674
+ StableUnCLIPImg2ImgPipeline,
675
+ StableUnCLIPPipeline,
676
+ StableVideoDiffusionPipeline,
677
+ TextToVideoSDPipeline,
678
+ TextToVideoZeroPipeline,
679
+ TextToVideoZeroSDXLPipeline,
680
+ UnCLIPImageVariationPipeline,
681
+ UnCLIPPipeline,
682
+ UniDiffuserModel,
683
+ UniDiffuserPipeline,
684
+ UniDiffuserTextDecoder,
685
+ VersatileDiffusionDualGuidedPipeline,
686
+ VersatileDiffusionImageVariationPipeline,
687
+ VersatileDiffusionPipeline,
688
+ VersatileDiffusionTextToImagePipeline,
689
+ VideoToVideoSDPipeline,
690
+ VQDiffusionPipeline,
691
+ WuerstchenCombinedPipeline,
692
+ WuerstchenDecoderPipeline,
693
+ WuerstchenPriorPipeline,
694
+ )
695
+
696
+ try:
697
+ if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
698
+ raise OptionalDependencyNotAvailable()
699
+ except OptionalDependencyNotAvailable:
700
+ from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import * # noqa F403
701
+ else:
702
+ from .pipelines import StableDiffusionKDiffusionPipeline, StableDiffusionXLKDiffusionPipeline
703
+
704
+ try:
705
+ if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
706
+ raise OptionalDependencyNotAvailable()
707
+ except OptionalDependencyNotAvailable:
708
+ from .utils.dummy_torch_and_transformers_and_onnx_objects import * # noqa F403
709
+ else:
710
+ from .pipelines import (
711
+ OnnxStableDiffusionImg2ImgPipeline,
712
+ OnnxStableDiffusionInpaintPipeline,
713
+ OnnxStableDiffusionInpaintPipelineLegacy,
714
+ OnnxStableDiffusionPipeline,
715
+ OnnxStableDiffusionUpscalePipeline,
716
+ StableDiffusionOnnxPipeline,
717
+ )
718
+
719
+ try:
720
+ if not (is_torch_available() and is_librosa_available()):
721
+ raise OptionalDependencyNotAvailable()
722
+ except OptionalDependencyNotAvailable:
723
+ from .utils.dummy_torch_and_librosa_objects import * # noqa F403
724
+ else:
725
+ from .pipelines import AudioDiffusionPipeline, Mel
726
+
727
+ try:
728
+ if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
729
+ raise OptionalDependencyNotAvailable()
730
+ except OptionalDependencyNotAvailable:
731
+ from .utils.dummy_transformers_and_torch_and_note_seq_objects import * # noqa F403
732
+ else:
733
+ from .pipelines import SpectrogramDiffusionPipeline
734
+
735
+ try:
736
+ if not is_flax_available():
737
+ raise OptionalDependencyNotAvailable()
738
+ except OptionalDependencyNotAvailable:
739
+ from .utils.dummy_flax_objects import * # noqa F403
740
+ else:
741
+ from .models.controlnet_flax import FlaxControlNetModel
742
+ from .models.modeling_flax_utils import FlaxModelMixin
743
+ from .models.unets.unet_2d_condition_flax import FlaxUNet2DConditionModel
744
+ from .models.vae_flax import FlaxAutoencoderKL
745
+ from .pipelines import FlaxDiffusionPipeline
746
+ from .schedulers import (
747
+ FlaxDDIMScheduler,
748
+ FlaxDDPMScheduler,
749
+ FlaxDPMSolverMultistepScheduler,
750
+ FlaxEulerDiscreteScheduler,
751
+ FlaxKarrasVeScheduler,
752
+ FlaxLMSDiscreteScheduler,
753
+ FlaxPNDMScheduler,
754
+ FlaxSchedulerMixin,
755
+ FlaxScoreSdeVeScheduler,
756
+ )
757
+
758
+ try:
759
+ if not (is_flax_available() and is_transformers_available()):
760
+ raise OptionalDependencyNotAvailable()
761
+ except OptionalDependencyNotAvailable:
762
+ from .utils.dummy_flax_and_transformers_objects import * # noqa F403
763
+ else:
764
+ from .pipelines import (
765
+ FlaxStableDiffusionControlNetPipeline,
766
+ FlaxStableDiffusionImg2ImgPipeline,
767
+ FlaxStableDiffusionInpaintPipeline,
768
+ FlaxStableDiffusionPipeline,
769
+ FlaxStableDiffusionXLPipeline,
770
+ )
771
+
772
+ try:
773
+ if not (is_note_seq_available()):
774
+ raise OptionalDependencyNotAvailable()
775
+ except OptionalDependencyNotAvailable:
776
+ from .utils.dummy_note_seq_objects import * # noqa F403
777
+ else:
778
+ from .pipelines import MidiProcessor
779
+
780
+ else:
781
+ import sys
782
+
783
+ sys.modules[__name__] = _LazyModule(
784
+ __name__,
785
+ globals()["__file__"],
786
+ _import_structure,
787
+ module_spec=__spec__,
788
+ extra_objects={"__version__": __version__},
789
+ )
model/diffusers_c/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (15.3 kB). View file
 
model/diffusers_c/__pycache__/configuration_utils.cpython-310.pyc ADDED
Binary file (24 kB). View file
 
model/diffusers_c/__pycache__/dependency_versions_check.cpython-310.pyc ADDED
Binary file (665 Bytes). View file
 
model/diffusers_c/__pycache__/dependency_versions_table.cpython-310.pyc ADDED
Binary file (1.37 kB). View file
 
model/diffusers_c/__pycache__/image_processor.cpython-310.pyc ADDED
Binary file (30.2 kB). View file
 
model/diffusers_c/commands/__init__.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from abc import ABC, abstractmethod
16
+ from argparse import ArgumentParser
17
+
18
+
19
+ class BaseDiffusersCLICommand(ABC):
20
+ @staticmethod
21
+ @abstractmethod
22
+ def register_subcommand(parser: ArgumentParser):
23
+ raise NotImplementedError()
24
+
25
+ @abstractmethod
26
+ def run(self):
27
+ raise NotImplementedError()
model/diffusers_c/commands/diffusers_cli.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ from argparse import ArgumentParser
17
+
18
+ from .env import EnvironmentCommand
19
+ from .fp16_safetensors import FP16SafetensorsCommand
20
+
21
+
22
+ def main():
23
+ parser = ArgumentParser("Diffusers CLI tool", usage="diffusers-cli <command> [<args>]")
24
+ commands_parser = parser.add_subparsers(help="diffusers-cli command helpers")
25
+
26
+ # Register commands
27
+ EnvironmentCommand.register_subcommand(commands_parser)
28
+ FP16SafetensorsCommand.register_subcommand(commands_parser)
29
+
30
+ # Let's go
31
+ args = parser.parse_args()
32
+
33
+ if not hasattr(args, "func"):
34
+ parser.print_help()
35
+ exit(1)
36
+
37
+ # Run
38
+ service = args.func(args)
39
+ service.run()
40
+
41
+
42
+ if __name__ == "__main__":
43
+ main()
model/diffusers_c/commands/env.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import platform
16
+ from argparse import ArgumentParser
17
+
18
+ import huggingface_hub
19
+
20
+ from .. import __version__ as version
21
+ from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
22
+ from . import BaseDiffusersCLICommand
23
+
24
+
25
+ def info_command_factory(_):
26
+ return EnvironmentCommand()
27
+
28
+
29
+ class EnvironmentCommand(BaseDiffusersCLICommand):
30
+ @staticmethod
31
+ def register_subcommand(parser: ArgumentParser):
32
+ download_parser = parser.add_parser("env")
33
+ download_parser.set_defaults(func=info_command_factory)
34
+
35
+ def run(self):
36
+ hub_version = huggingface_hub.__version__
37
+
38
+ pt_version = "not installed"
39
+ pt_cuda_available = "NA"
40
+ if is_torch_available():
41
+ import torch
42
+
43
+ pt_version = torch.__version__
44
+ pt_cuda_available = torch.cuda.is_available()
45
+
46
+ transformers_version = "not installed"
47
+ if is_transformers_available():
48
+ import transformers
49
+
50
+ transformers_version = transformers.__version__
51
+
52
+ accelerate_version = "not installed"
53
+ if is_accelerate_available():
54
+ import accelerate
55
+
56
+ accelerate_version = accelerate.__version__
57
+
58
+ xformers_version = "not installed"
59
+ if is_xformers_available():
60
+ import xformers
61
+
62
+ xformers_version = xformers.__version__
63
+
64
+ info = {
65
+ "`diffusers` version": version,
66
+ "Platform": platform.platform(),
67
+ "Python version": platform.python_version(),
68
+ "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
69
+ "Huggingface_hub version": hub_version,
70
+ "Transformers version": transformers_version,
71
+ "Accelerate version": accelerate_version,
72
+ "xFormers version": xformers_version,
73
+ "Using GPU in script?": "<fill in>",
74
+ "Using distributed or parallel set-up in script?": "<fill in>",
75
+ }
76
+
77
+ print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n")
78
+ print(self.format_dict(info))
79
+
80
+ return info
81
+
82
+ @staticmethod
83
+ def format_dict(d):
84
+ return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"
model/diffusers_c/commands/fp16_safetensors.py ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """
16
+ Usage example:
17
+ diffusers-cli fp16_safetensors --ckpt_id=openai/shap-e --fp16 --use_safetensors
18
+ """
19
+
20
+ import glob
21
+ import json
22
+ import warnings
23
+ from argparse import ArgumentParser, Namespace
24
+ from importlib import import_module
25
+
26
+ import huggingface_hub
27
+ import torch
28
+ from huggingface_hub import hf_hub_download
29
+ from packaging import version
30
+
31
+ from ..utils import logging
32
+ from . import BaseDiffusersCLICommand
33
+
34
+
35
+ def conversion_command_factory(args: Namespace):
36
+ if args.use_auth_token:
37
+ warnings.warn(
38
+ "The `--use_auth_token` flag is deprecated and will be removed in a future version. Authentication is now"
39
+ " handled automatically if user is logged in."
40
+ )
41
+ return FP16SafetensorsCommand(args.ckpt_id, args.fp16, args.use_safetensors)
42
+
43
+
44
+ class FP16SafetensorsCommand(BaseDiffusersCLICommand):
45
+ @staticmethod
46
+ def register_subcommand(parser: ArgumentParser):
47
+ conversion_parser = parser.add_parser("fp16_safetensors")
48
+ conversion_parser.add_argument(
49
+ "--ckpt_id",
50
+ type=str,
51
+ help="Repo id of the checkpoints on which to run the conversion. Example: 'openai/shap-e'.",
52
+ )
53
+ conversion_parser.add_argument(
54
+ "--fp16", action="store_true", help="If serializing the variables in FP16 precision."
55
+ )
56
+ conversion_parser.add_argument(
57
+ "--use_safetensors", action="store_true", help="If serializing in the safetensors format."
58
+ )
59
+ conversion_parser.add_argument(
60
+ "--use_auth_token",
61
+ action="store_true",
62
+ help="When working with checkpoints having private visibility. When used `huggingface-cli login` needs to be run beforehand.",
63
+ )
64
+ conversion_parser.set_defaults(func=conversion_command_factory)
65
+
66
+ def __init__(self, ckpt_id: str, fp16: bool, use_safetensors: bool):
67
+ self.logger = logging.get_logger("diffusers-cli/fp16_safetensors")
68
+ self.ckpt_id = ckpt_id
69
+ self.local_ckpt_dir = f"/tmp/{ckpt_id}"
70
+ self.fp16 = fp16
71
+
72
+ self.use_safetensors = use_safetensors
73
+
74
+ if not self.use_safetensors and not self.fp16:
75
+ raise NotImplementedError(
76
+ "When `use_safetensors` and `fp16` both are False, then this command is of no use."
77
+ )
78
+
79
+ def run(self):
80
+ if version.parse(huggingface_hub.__version__) < version.parse("0.9.0"):
81
+ raise ImportError(
82
+ "The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub"
83
+ " installation."
84
+ )
85
+ else:
86
+ from huggingface_hub import create_commit
87
+ from huggingface_hub._commit_api import CommitOperationAdd
88
+
89
+ model_index = hf_hub_download(repo_id=self.ckpt_id, filename="model_index.json")
90
+ with open(model_index, "r") as f:
91
+ pipeline_class_name = json.load(f)["_class_name"]
92
+ pipeline_class = getattr(import_module("diffusers"), pipeline_class_name)
93
+ self.logger.info(f"Pipeline class imported: {pipeline_class_name}.")
94
+
95
+ # Load the appropriate pipeline. We could have use `DiffusionPipeline`
96
+ # here, but just to avoid any rough edge cases.
97
+ pipeline = pipeline_class.from_pretrained(
98
+ self.ckpt_id, torch_dtype=torch.float16 if self.fp16 else torch.float32
99
+ )
100
+ pipeline.save_pretrained(
101
+ self.local_ckpt_dir,
102
+ safe_serialization=True if self.use_safetensors else False,
103
+ variant="fp16" if self.fp16 else None,
104
+ )
105
+ self.logger.info(f"Pipeline locally saved to {self.local_ckpt_dir}.")
106
+
107
+ # Fetch all the paths.
108
+ if self.fp16:
109
+ modified_paths = glob.glob(f"{self.local_ckpt_dir}/*/*.fp16.*")
110
+ elif self.use_safetensors:
111
+ modified_paths = glob.glob(f"{self.local_ckpt_dir}/*/*.safetensors")
112
+
113
+ # Prepare for the PR.
114
+ commit_message = f"Serialize variables with FP16: {self.fp16} and safetensors: {self.use_safetensors}."
115
+ operations = []
116
+ for path in modified_paths:
117
+ operations.append(CommitOperationAdd(path_in_repo="/".join(path.split("/")[4:]), path_or_fileobj=path))
118
+
119
+ # Open the PR.
120
+ commit_description = (
121
+ "Variables converted by the [`diffusers`' `fp16_safetensors`"
122
+ " CLI](https://github.com/huggingface/diffusers/blob/main/src/diffusers/commands/fp16_safetensors.py)."
123
+ )
124
+ hub_pr_url = create_commit(
125
+ repo_id=self.ckpt_id,
126
+ operations=operations,
127
+ commit_message=commit_message,
128
+ commit_description=commit_description,
129
+ repo_type="model",
130
+ create_pr=True,
131
+ ).pr_url
132
+ self.logger.info(f"PR created here: {hub_pr_url}.")
model/diffusers_c/configuration_utils.py ADDED
@@ -0,0 +1,703 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Inc. team.
2
+ # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """ConfigMixin base class and utilities."""
16
+
17
+ import dataclasses
18
+ import functools
19
+ import importlib
20
+ import inspect
21
+ import json
22
+ import os
23
+ import re
24
+ from collections import OrderedDict
25
+ from pathlib import PosixPath
26
+ from typing import Any, Dict, Tuple, Union
27
+
28
+ import numpy as np
29
+ from huggingface_hub import create_repo, hf_hub_download
30
+ from huggingface_hub.utils import (
31
+ EntryNotFoundError,
32
+ RepositoryNotFoundError,
33
+ RevisionNotFoundError,
34
+ validate_hf_hub_args,
35
+ )
36
+ from requests import HTTPError
37
+
38
+ from . import __version__
39
+ from .utils import (
40
+ HUGGINGFACE_CO_RESOLVE_ENDPOINT,
41
+ DummyObject,
42
+ deprecate,
43
+ extract_commit_hash,
44
+ http_user_agent,
45
+ logging,
46
+ )
47
+
48
+
49
+ logger = logging.get_logger(__name__)
50
+
51
+ _re_configuration_file = re.compile(r"config\.(.*)\.json")
52
+
53
+
54
+ class FrozenDict(OrderedDict):
55
+ def __init__(self, *args, **kwargs):
56
+ super().__init__(*args, **kwargs)
57
+
58
+ for key, value in self.items():
59
+ setattr(self, key, value)
60
+
61
+ self.__frozen = True
62
+
63
+ def __delitem__(self, *args, **kwargs):
64
+ raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
65
+
66
+ def setdefault(self, *args, **kwargs):
67
+ raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
68
+
69
+ def pop(self, *args, **kwargs):
70
+ raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
71
+
72
+ def update(self, *args, **kwargs):
73
+ raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
74
+
75
+ def __setattr__(self, name, value):
76
+ if hasattr(self, "__frozen") and self.__frozen:
77
+ raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
78
+ super().__setattr__(name, value)
79
+
80
+ def __setitem__(self, name, value):
81
+ if hasattr(self, "__frozen") and self.__frozen:
82
+ raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
83
+ super().__setitem__(name, value)
84
+
85
+
86
+ class ConfigMixin:
87
+ r"""
88
+ Base class for all configuration classes. All configuration parameters are stored under `self.config`. Also
89
+ provides the [`~ConfigMixin.from_config`] and [`~ConfigMixin.save_config`] methods for loading, downloading, and
90
+ saving classes that inherit from [`ConfigMixin`].
91
+
92
+ Class attributes:
93
+ - **config_name** (`str`) -- A filename under which the config should stored when calling
94
+ [`~ConfigMixin.save_config`] (should be overridden by parent class).
95
+ - **ignore_for_config** (`List[str]`) -- A list of attributes that should not be saved in the config (should be
96
+ overridden by subclass).
97
+ - **has_compatibles** (`bool`) -- Whether the class has compatible classes (should be overridden by subclass).
98
+ - **_deprecated_kwargs** (`List[str]`) -- Keyword arguments that are deprecated. Note that the `init` function
99
+ should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by
100
+ subclass).
101
+ """
102
+
103
+ config_name = None
104
+ ignore_for_config = []
105
+ has_compatibles = False
106
+
107
+ _deprecated_kwargs = []
108
+
109
+ def register_to_config(self, **kwargs):
110
+ if self.config_name is None:
111
+ raise NotImplementedError(f"Make sure that {self.__class__} has defined a class name `config_name`")
112
+ # Special case for `kwargs` used in deprecation warning added to schedulers
113
+ # TODO: remove this when we remove the deprecation warning, and the `kwargs` argument,
114
+ # or solve in a more general way.
115
+ kwargs.pop("kwargs", None)
116
+
117
+ if not hasattr(self, "_internal_dict"):
118
+ internal_dict = kwargs
119
+ else:
120
+ previous_dict = dict(self._internal_dict)
121
+ internal_dict = {**self._internal_dict, **kwargs}
122
+ logger.debug(f"Updating config from {previous_dict} to {internal_dict}")
123
+
124
+ self._internal_dict = FrozenDict(internal_dict)
125
+
126
+ def __getattr__(self, name: str) -> Any:
127
+ """The only reason we overwrite `getattr` here is to gracefully deprecate accessing
128
+ config attributes directly. See https://github.com/huggingface/diffusers/pull/3129
129
+
130
+ This function is mostly copied from PyTorch's __getattr__ overwrite:
131
+ https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
132
+ """
133
+
134
+ is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
135
+ is_attribute = name in self.__dict__
136
+
137
+ if is_in_config and not is_attribute:
138
+ deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'scheduler.config.{name}'."
139
+ deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False)
140
+ return self._internal_dict[name]
141
+
142
+ raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
143
+
144
+ def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
145
+ """
146
+ Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the
147
+ [`~ConfigMixin.from_config`] class method.
148
+
149
+ Args:
150
+ save_directory (`str` or `os.PathLike`):
151
+ Directory where the configuration JSON file is saved (will be created if it does not exist).
152
+ push_to_hub (`bool`, *optional*, defaults to `False`):
153
+ Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
154
+ repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
155
+ namespace).
156
+ kwargs (`Dict[str, Any]`, *optional*):
157
+ Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
158
+ """
159
+ if os.path.isfile(save_directory):
160
+ raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
161
+
162
+ os.makedirs(save_directory, exist_ok=True)
163
+
164
+ # If we save using the predefined names, we can load using `from_config`
165
+ output_config_file = os.path.join(save_directory, self.config_name)
166
+
167
+ self.to_json_file(output_config_file)
168
+ logger.info(f"Configuration saved in {output_config_file}")
169
+
170
+ if push_to_hub:
171
+ commit_message = kwargs.pop("commit_message", None)
172
+ private = kwargs.pop("private", False)
173
+ create_pr = kwargs.pop("create_pr", False)
174
+ token = kwargs.pop("token", None)
175
+ repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
176
+ repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
177
+
178
+ self._upload_folder(
179
+ save_directory,
180
+ repo_id,
181
+ token=token,
182
+ commit_message=commit_message,
183
+ create_pr=create_pr,
184
+ )
185
+
186
+ @classmethod
187
+ def from_config(cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs):
188
+ r"""
189
+ Instantiate a Python class from a config dictionary.
190
+
191
+ Parameters:
192
+ config (`Dict[str, Any]`):
193
+ A config dictionary from which the Python class is instantiated. Make sure to only load configuration
194
+ files of compatible classes.
195
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
196
+ Whether kwargs that are not consumed by the Python class should be returned or not.
197
+ kwargs (remaining dictionary of keyword arguments, *optional*):
198
+ Can be used to update the configuration object (after it is loaded) and initiate the Python class.
199
+ `**kwargs` are passed directly to the underlying scheduler/model's `__init__` method and eventually
200
+ overwrite the same named arguments in `config`.
201
+
202
+ Returns:
203
+ [`ModelMixin`] or [`SchedulerMixin`]:
204
+ A model or scheduler object instantiated from a config dictionary.
205
+
206
+ Examples:
207
+
208
+ ```python
209
+ >>> from diffusers import DDPMScheduler, DDIMScheduler, PNDMScheduler
210
+
211
+ >>> # Download scheduler from huggingface.co and cache.
212
+ >>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cifar10-32")
213
+
214
+ >>> # Instantiate DDIM scheduler class with same config as DDPM
215
+ >>> scheduler = DDIMScheduler.from_config(scheduler.config)
216
+
217
+ >>> # Instantiate PNDM scheduler class with same config as DDPM
218
+ >>> scheduler = PNDMScheduler.from_config(scheduler.config)
219
+ ```
220
+ """
221
+ # <===== TO BE REMOVED WITH DEPRECATION
222
+ # TODO(Patrick) - make sure to remove the following lines when config=="model_path" is deprecated
223
+ if "pretrained_model_name_or_path" in kwargs:
224
+ config = kwargs.pop("pretrained_model_name_or_path")
225
+
226
+ if config is None:
227
+ raise ValueError("Please make sure to provide a config as the first positional argument.")
228
+ # ======>
229
+
230
+ if not isinstance(config, dict):
231
+ deprecation_message = "It is deprecated to pass a pretrained model name or path to `from_config`."
232
+ if "Scheduler" in cls.__name__:
233
+ deprecation_message += (
234
+ f"If you were trying to load a scheduler, please use {cls}.from_pretrained(...) instead."
235
+ " Otherwise, please make sure to pass a configuration dictionary instead. This functionality will"
236
+ " be removed in v1.0.0."
237
+ )
238
+ elif "Model" in cls.__name__:
239
+ deprecation_message += (
240
+ f"If you were trying to load a model, please use {cls}.load_config(...) followed by"
241
+ f" {cls}.from_config(...) instead. Otherwise, please make sure to pass a configuration dictionary"
242
+ " instead. This functionality will be removed in v1.0.0."
243
+ )
244
+ deprecate("config-passed-as-path", "1.0.0", deprecation_message, standard_warn=False)
245
+ config, kwargs = cls.load_config(pretrained_model_name_or_path=config, return_unused_kwargs=True, **kwargs)
246
+
247
+ init_dict, unused_kwargs, hidden_dict = cls.extract_init_dict(config, **kwargs)
248
+
249
+ # Allow dtype to be specified on initialization
250
+ if "dtype" in unused_kwargs:
251
+ init_dict["dtype"] = unused_kwargs.pop("dtype")
252
+
253
+ # add possible deprecated kwargs
254
+ for deprecated_kwarg in cls._deprecated_kwargs:
255
+ if deprecated_kwarg in unused_kwargs:
256
+ init_dict[deprecated_kwarg] = unused_kwargs.pop(deprecated_kwarg)
257
+
258
+ # Return model and optionally state and/or unused_kwargs
259
+ model = cls(**init_dict)
260
+
261
+ # make sure to also save config parameters that might be used for compatible classes
262
+ # update _class_name
263
+ if "_class_name" in hidden_dict:
264
+ hidden_dict["_class_name"] = cls.__name__
265
+
266
+ model.register_to_config(**hidden_dict)
267
+
268
+ # add hidden kwargs of compatible classes to unused_kwargs
269
+ unused_kwargs = {**unused_kwargs, **hidden_dict}
270
+
271
+ if return_unused_kwargs:
272
+ return (model, unused_kwargs)
273
+ else:
274
+ return model
275
+
276
+ @classmethod
277
+ def get_config_dict(cls, *args, **kwargs):
278
+ deprecation_message = (
279
+ f" The function get_config_dict is deprecated. Please use {cls}.load_config instead. This function will be"
280
+ " removed in version v1.0.0"
281
+ )
282
+ deprecate("get_config_dict", "1.0.0", deprecation_message, standard_warn=False)
283
+ return cls.load_config(*args, **kwargs)
284
+
285
+ @classmethod
286
+ @validate_hf_hub_args
287
+ def load_config(
288
+ cls,
289
+ pretrained_model_name_or_path: Union[str, os.PathLike],
290
+ return_unused_kwargs=False,
291
+ return_commit_hash=False,
292
+ **kwargs,
293
+ ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
294
+ r"""
295
+ Load a model or scheduler configuration.
296
+
297
+ Parameters:
298
+ pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
299
+ Can be either:
300
+
301
+ - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
302
+ the Hub.
303
+ - A path to a *directory* (for example `./my_model_directory`) containing model weights saved with
304
+ [`~ConfigMixin.save_config`].
305
+
306
+ cache_dir (`Union[str, os.PathLike]`, *optional*):
307
+ Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
308
+ is not used.
309
+ force_download (`bool`, *optional*, defaults to `False`):
310
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
311
+ cached versions if they exist.
312
+ resume_download (`bool`, *optional*, defaults to `False`):
313
+ Whether or not to resume downloading the model weights and configuration files. If set to `False`, any
314
+ incompletely downloaded files are deleted.
315
+ proxies (`Dict[str, str]`, *optional*):
316
+ A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
317
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
318
+ output_loading_info(`bool`, *optional*, defaults to `False`):
319
+ Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
320
+ local_files_only (`bool`, *optional*, defaults to `False`):
321
+ Whether to only load local model weights and configuration files or not. If set to `True`, the model
322
+ won't be downloaded from the Hub.
323
+ token (`str` or *bool*, *optional*):
324
+ The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
325
+ `diffusers-cli login` (stored in `~/.huggingface`) is used.
326
+ revision (`str`, *optional*, defaults to `"main"`):
327
+ The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
328
+ allowed by Git.
329
+ subfolder (`str`, *optional*, defaults to `""`):
330
+ The subfolder location of a model file within a larger model repository on the Hub or locally.
331
+ return_unused_kwargs (`bool`, *optional*, defaults to `False):
332
+ Whether unused keyword arguments of the config are returned.
333
+ return_commit_hash (`bool`, *optional*, defaults to `False):
334
+ Whether the `commit_hash` of the loaded configuration are returned.
335
+
336
+ Returns:
337
+ `dict`:
338
+ A dictionary of all the parameters stored in a JSON configuration file.
339
+
340
+ """
341
+ cache_dir = kwargs.pop("cache_dir", None)
342
+ force_download = kwargs.pop("force_download", False)
343
+ resume_download = kwargs.pop("resume_download", False)
344
+ proxies = kwargs.pop("proxies", None)
345
+ token = kwargs.pop("token", None)
346
+ local_files_only = kwargs.pop("local_files_only", False)
347
+ revision = kwargs.pop("revision", None)
348
+ _ = kwargs.pop("mirror", None)
349
+ subfolder = kwargs.pop("subfolder", None)
350
+ user_agent = kwargs.pop("user_agent", {})
351
+
352
+ user_agent = {**user_agent, "file_type": "config"}
353
+ user_agent = http_user_agent(user_agent)
354
+
355
+ pretrained_model_name_or_path = str(pretrained_model_name_or_path)
356
+
357
+ if cls.config_name is None:
358
+ raise ValueError(
359
+ "`self.config_name` is not defined. Note that one should not load a config from "
360
+ "`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`"
361
+ )
362
+
363
+ if os.path.isfile(pretrained_model_name_or_path):
364
+ config_file = pretrained_model_name_or_path
365
+ elif os.path.isdir(pretrained_model_name_or_path):
366
+ if os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
367
+ # Load from a PyTorch checkpoint
368
+ config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
369
+ elif subfolder is not None and os.path.isfile(
370
+ os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
371
+ ):
372
+ config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
373
+ else:
374
+ raise EnvironmentError(
375
+ f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}."
376
+ )
377
+ else:
378
+ try:
379
+ # Load from URL or cache if already cached
380
+ config_file = hf_hub_download(
381
+ pretrained_model_name_or_path,
382
+ filename=cls.config_name,
383
+ cache_dir=cache_dir,
384
+ force_download=force_download,
385
+ proxies=proxies,
386
+ resume_download=resume_download,
387
+ local_files_only=local_files_only,
388
+ token=token,
389
+ user_agent=user_agent,
390
+ subfolder=subfolder,
391
+ revision=revision,
392
+ )
393
+ except RepositoryNotFoundError:
394
+ raise EnvironmentError(
395
+ f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier"
396
+ " listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a"
397
+ " token having permission to this repo with `token` or log in with `huggingface-cli login`."
398
+ )
399
+ except RevisionNotFoundError:
400
+ raise EnvironmentError(
401
+ f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for"
402
+ " this model name. Check the model page at"
403
+ f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
404
+ )
405
+ except EntryNotFoundError:
406
+ raise EnvironmentError(
407
+ f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}."
408
+ )
409
+ except HTTPError as err:
410
+ raise EnvironmentError(
411
+ "There was a specific connection error when trying to load"
412
+ f" {pretrained_model_name_or_path}:\n{err}"
413
+ )
414
+ except ValueError:
415
+ raise EnvironmentError(
416
+ f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
417
+ f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
418
+ f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to"
419
+ " run the library in offline mode at"
420
+ " 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
421
+ )
422
+ except EnvironmentError:
423
+ raise EnvironmentError(
424
+ f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from "
425
+ "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
426
+ f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
427
+ f"containing a {cls.config_name} file"
428
+ )
429
+
430
+ try:
431
+ # Load config dict
432
+ config_dict = cls._dict_from_json_file(config_file)
433
+
434
+ commit_hash = extract_commit_hash(config_file)
435
+ except (json.JSONDecodeError, UnicodeDecodeError):
436
+ raise EnvironmentError(f"It looks like the config file at '{config_file}' is not a valid JSON file.")
437
+
438
+ if not (return_unused_kwargs or return_commit_hash):
439
+ return config_dict
440
+
441
+ outputs = (config_dict,)
442
+
443
+ if return_unused_kwargs:
444
+ outputs += (kwargs,)
445
+
446
+ if return_commit_hash:
447
+ outputs += (commit_hash,)
448
+
449
+ return outputs
450
+
451
+ @staticmethod
452
+ def _get_init_keys(cls):
453
+ return set(dict(inspect.signature(cls.__init__).parameters).keys())
454
+
455
+ @classmethod
456
+ def extract_init_dict(cls, config_dict, **kwargs):
457
+ # Skip keys that were not present in the original config, so default __init__ values were used
458
+ used_defaults = config_dict.get("_use_default_values", [])
459
+ config_dict = {k: v for k, v in config_dict.items() if k not in used_defaults and k != "_use_default_values"}
460
+
461
+ # 0. Copy origin config dict
462
+ original_dict = dict(config_dict.items())
463
+
464
+ # 1. Retrieve expected config attributes from __init__ signature
465
+ expected_keys = cls._get_init_keys(cls)
466
+ expected_keys.remove("self")
467
+ # remove general kwargs if present in dict
468
+ if "kwargs" in expected_keys:
469
+ expected_keys.remove("kwargs")
470
+ # remove flax internal keys
471
+ if hasattr(cls, "_flax_internal_args"):
472
+ for arg in cls._flax_internal_args:
473
+ expected_keys.remove(arg)
474
+
475
+ # 2. Remove attributes that cannot be expected from expected config attributes
476
+ # remove keys to be ignored
477
+ if len(cls.ignore_for_config) > 0:
478
+ expected_keys = expected_keys - set(cls.ignore_for_config)
479
+
480
+ # load diffusers library to import compatible and original scheduler
481
+ diffusers_library = importlib.import_module(__name__.split(".")[0])
482
+
483
+ if cls.has_compatibles:
484
+ compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)]
485
+ else:
486
+ compatible_classes = []
487
+
488
+ expected_keys_comp_cls = set()
489
+ for c in compatible_classes:
490
+ expected_keys_c = cls._get_init_keys(c)
491
+ expected_keys_comp_cls = expected_keys_comp_cls.union(expected_keys_c)
492
+ expected_keys_comp_cls = expected_keys_comp_cls - cls._get_init_keys(cls)
493
+ config_dict = {k: v for k, v in config_dict.items() if k not in expected_keys_comp_cls}
494
+
495
+ # remove attributes from orig class that cannot be expected
496
+ orig_cls_name = config_dict.pop("_class_name", cls.__name__)
497
+ if (
498
+ isinstance(orig_cls_name, str)
499
+ and orig_cls_name != cls.__name__
500
+ and hasattr(diffusers_library, orig_cls_name)
501
+ ):
502
+ orig_cls = getattr(diffusers_library, orig_cls_name)
503
+ unexpected_keys_from_orig = cls._get_init_keys(orig_cls) - expected_keys
504
+ config_dict = {k: v for k, v in config_dict.items() if k not in unexpected_keys_from_orig}
505
+ elif not isinstance(orig_cls_name, str) and not isinstance(orig_cls_name, (list, tuple)):
506
+ raise ValueError(
507
+ "Make sure that the `_class_name` is of type string or list of string (for custom pipelines)."
508
+ )
509
+
510
+ # remove private attributes
511
+ config_dict = {k: v for k, v in config_dict.items() if not k.startswith("_")}
512
+
513
+ # 3. Create keyword arguments that will be passed to __init__ from expected keyword arguments
514
+ init_dict = {}
515
+ for key in expected_keys:
516
+ # if config param is passed to kwarg and is present in config dict
517
+ # it should overwrite existing config dict key
518
+ if key in kwargs and key in config_dict:
519
+ config_dict[key] = kwargs.pop(key)
520
+
521
+ if key in kwargs:
522
+ # overwrite key
523
+ init_dict[key] = kwargs.pop(key)
524
+ elif key in config_dict:
525
+ # use value from config dict
526
+ init_dict[key] = config_dict.pop(key)
527
+
528
+ # 4. Give nice warning if unexpected values have been passed
529
+ if len(config_dict) > 0:
530
+ logger.warning(
531
+ f"The config attributes {config_dict} were passed to {cls.__name__}, "
532
+ "but are not expected and will be ignored. Please verify your "
533
+ f"{cls.config_name} configuration file."
534
+ )
535
+
536
+ # 5. Give nice info if config attributes are initialized to default because they have not been passed
537
+ passed_keys = set(init_dict.keys())
538
+ if len(expected_keys - passed_keys) > 0:
539
+ logger.info(
540
+ f"{expected_keys - passed_keys} was not found in config. Values will be initialized to default values."
541
+ )
542
+
543
+ # 6. Define unused keyword arguments
544
+ unused_kwargs = {**config_dict, **kwargs}
545
+
546
+ # 7. Define "hidden" config parameters that were saved for compatible classes
547
+ hidden_config_dict = {k: v for k, v in original_dict.items() if k not in init_dict}
548
+
549
+ return init_dict, unused_kwargs, hidden_config_dict
550
+
551
+ @classmethod
552
+ def _dict_from_json_file(cls, json_file: Union[str, os.PathLike]):
553
+ with open(json_file, "r", encoding="utf-8") as reader:
554
+ text = reader.read()
555
+ return json.loads(text)
556
+
557
+ def __repr__(self):
558
+ return f"{self.__class__.__name__} {self.to_json_string()}"
559
+
560
+ @property
561
+ def config(self) -> Dict[str, Any]:
562
+ """
563
+ Returns the config of the class as a frozen dictionary
564
+
565
+ Returns:
566
+ `Dict[str, Any]`: Config of the class.
567
+ """
568
+ return self._internal_dict
569
+
570
+ def to_json_string(self) -> str:
571
+ """
572
+ Serializes the configuration instance to a JSON string.
573
+
574
+ Returns:
575
+ `str`:
576
+ String containing all the attributes that make up the configuration instance in JSON format.
577
+ """
578
+ config_dict = self._internal_dict if hasattr(self, "_internal_dict") else {}
579
+ config_dict["_class_name"] = self.__class__.__name__
580
+ config_dict["_diffusers_version"] = __version__
581
+
582
+ def to_json_saveable(value):
583
+ if isinstance(value, np.ndarray):
584
+ value = value.tolist()
585
+ elif isinstance(value, PosixPath):
586
+ value = str(value)
587
+ return value
588
+
589
+ config_dict = {k: to_json_saveable(v) for k, v in config_dict.items()}
590
+ # Don't save "_ignore_files" or "_use_default_values"
591
+ config_dict.pop("_ignore_files", None)
592
+ config_dict.pop("_use_default_values", None)
593
+
594
+ return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
595
+
596
+ def to_json_file(self, json_file_path: Union[str, os.PathLike]):
597
+ """
598
+ Save the configuration instance's parameters to a JSON file.
599
+
600
+ Args:
601
+ json_file_path (`str` or `os.PathLike`):
602
+ Path to the JSON file to save a configuration instance's parameters.
603
+ """
604
+ with open(json_file_path, "w", encoding="utf-8") as writer:
605
+ writer.write(self.to_json_string())
606
+
607
+
608
+ def register_to_config(init):
609
+ r"""
610
+ Decorator to apply on the init of classes inheriting from [`ConfigMixin`] so that all the arguments are
611
+ automatically sent to `self.register_for_config`. To ignore a specific argument accepted by the init but that
612
+ shouldn't be registered in the config, use the `ignore_for_config` class variable
613
+
614
+ Warning: Once decorated, all private arguments (beginning with an underscore) are trashed and not sent to the init!
615
+ """
616
+
617
+ @functools.wraps(init)
618
+ def inner_init(self, *args, **kwargs):
619
+ # Ignore private kwargs in the init.
620
+ init_kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
621
+ config_init_kwargs = {k: v for k, v in kwargs.items() if k.startswith("_")}
622
+ if not isinstance(self, ConfigMixin):
623
+ raise RuntimeError(
624
+ f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does "
625
+ "not inherit from `ConfigMixin`."
626
+ )
627
+
628
+ ignore = getattr(self, "ignore_for_config", [])
629
+ # Get positional arguments aligned with kwargs
630
+ new_kwargs = {}
631
+ signature = inspect.signature(init)
632
+ parameters = {
633
+ name: p.default for i, (name, p) in enumerate(signature.parameters.items()) if i > 0 and name not in ignore
634
+ }
635
+ for arg, name in zip(args, parameters.keys()):
636
+ new_kwargs[name] = arg
637
+
638
+ # Then add all kwargs
639
+ new_kwargs.update(
640
+ {
641
+ k: init_kwargs.get(k, default)
642
+ for k, default in parameters.items()
643
+ if k not in ignore and k not in new_kwargs
644
+ }
645
+ )
646
+
647
+ # Take note of the parameters that were not present in the loaded config
648
+ if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0:
649
+ new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs))
650
+
651
+ new_kwargs = {**config_init_kwargs, **new_kwargs}
652
+ getattr(self, "register_to_config")(**new_kwargs)
653
+ init(self, *args, **init_kwargs)
654
+
655
+ return inner_init
656
+
657
+
658
+ def flax_register_to_config(cls):
659
+ original_init = cls.__init__
660
+
661
+ @functools.wraps(original_init)
662
+ def init(self, *args, **kwargs):
663
+ if not isinstance(self, ConfigMixin):
664
+ raise RuntimeError(
665
+ f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does "
666
+ "not inherit from `ConfigMixin`."
667
+ )
668
+
669
+ # Ignore private kwargs in the init. Retrieve all passed attributes
670
+ init_kwargs = dict(kwargs.items())
671
+
672
+ # Retrieve default values
673
+ fields = dataclasses.fields(self)
674
+ default_kwargs = {}
675
+ for field in fields:
676
+ # ignore flax specific attributes
677
+ if field.name in self._flax_internal_args:
678
+ continue
679
+ if type(field.default) == dataclasses._MISSING_TYPE:
680
+ default_kwargs[field.name] = None
681
+ else:
682
+ default_kwargs[field.name] = getattr(self, field.name)
683
+
684
+ # Make sure init_kwargs override default kwargs
685
+ new_kwargs = {**default_kwargs, **init_kwargs}
686
+ # dtype should be part of `init_kwargs`, but not `new_kwargs`
687
+ if "dtype" in new_kwargs:
688
+ new_kwargs.pop("dtype")
689
+
690
+ # Get positional arguments aligned with kwargs
691
+ for i, arg in enumerate(args):
692
+ name = fields[i].name
693
+ new_kwargs[name] = arg
694
+
695
+ # Take note of the parameters that were not present in the loaded config
696
+ if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0:
697
+ new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs))
698
+
699
+ getattr(self, "register_to_config")(**new_kwargs)
700
+ original_init(self, *args, **kwargs)
701
+
702
+ cls.__init__ = init
703
+ return cls
model/diffusers_c/dependency_versions_check.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from .dependency_versions_table import deps
16
+ from .utils.versions import require_version, require_version_core
17
+
18
+
19
+ # define which module versions we always want to check at run time
20
+ # (usually the ones defined in `install_requires` in setup.py)
21
+ #
22
+ # order specific notes:
23
+ # - tqdm must be checked before tokenizers
24
+
25
+ pkgs_to_check_at_runtime = "python requests filelock numpy".split()
26
+ for pkg in pkgs_to_check_at_runtime:
27
+ if pkg in deps:
28
+ require_version_core(deps[pkg])
29
+ else:
30
+ raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")
31
+
32
+
33
+ def dep_version_check(pkg, hint=None):
34
+ require_version(deps[pkg], hint)
model/diffusers_c/dependency_versions_table.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # THIS FILE HAS BEEN AUTOGENERATED. To update:
2
+ # 1. modify the `_deps` dict in setup.py
3
+ # 2. run `make deps_table_update`
4
+ deps = {
5
+ "Pillow": "Pillow",
6
+ "accelerate": "accelerate>=0.11.0",
7
+ "compel": "compel==0.1.8",
8
+ "datasets": "datasets",
9
+ "filelock": "filelock",
10
+ "flax": "flax>=0.4.1",
11
+ "hf-doc-builder": "hf-doc-builder>=0.3.0",
12
+ "huggingface-hub": "huggingface-hub",
13
+ "requests-mock": "requests-mock==1.10.0",
14
+ "importlib_metadata": "importlib_metadata",
15
+ "invisible-watermark": "invisible-watermark>=0.2.0",
16
+ "isort": "isort>=5.5.4",
17
+ "jax": "jax>=0.4.1",
18
+ "jaxlib": "jaxlib>=0.4.1",
19
+ "Jinja2": "Jinja2",
20
+ "k-diffusion": "k-diffusion>=0.0.12",
21
+ "torchsde": "torchsde",
22
+ "note_seq": "note_seq",
23
+ "librosa": "librosa",
24
+ "numpy": "numpy",
25
+ "parameterized": "parameterized",
26
+ "peft": "peft>=0.6.0",
27
+ "protobuf": "protobuf>=3.20.3,<4",
28
+ "pytest": "pytest",
29
+ "pytest-timeout": "pytest-timeout",
30
+ "pytest-xdist": "pytest-xdist",
31
+ "python": "python>=3.8.0",
32
+ "ruff": "ruff==0.1.5",
33
+ "safetensors": "safetensors>=0.3.1",
34
+ "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
35
+ "GitPython": "GitPython<3.1.19",
36
+ "scipy": "scipy",
37
+ "onnx": "onnx",
38
+ "regex": "regex!=2019.12.17",
39
+ "requests": "requests",
40
+ "tensorboard": "tensorboard",
41
+ "torch": "torch>=1.4",
42
+ "torchvision": "torchvision",
43
+ "transformers": "transformers>=4.25.1",
44
+ "urllib3": "urllib3<=2.0.0",
45
+ }
model/diffusers_c/experimental/README.md ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ # 🧨 Diffusers Experimental
2
+
3
+ We are adding experimental code to support novel applications and usages of the Diffusers library.
4
+ Currently, the following experiments are supported:
5
+ * Reinforcement learning via an implementation of the [Diffuser](https://arxiv.org/abs/2205.09991) model.
model/diffusers_c/experimental/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from .rl import ValueGuidedRLPipeline
model/diffusers_c/experimental/rl/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from .value_guided_sampling import ValueGuidedRLPipeline
model/diffusers_c/experimental/rl/value_guided_sampling.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import numpy as np
16
+ import torch
17
+ import tqdm
18
+
19
+ from ...models.unets.unet_1d import UNet1DModel
20
+ from ...pipelines import DiffusionPipeline
21
+ from ...utils.dummy_pt_objects import DDPMScheduler
22
+ from ...utils.torch_utils import randn_tensor
23
+
24
+
25
+ class ValueGuidedRLPipeline(DiffusionPipeline):
26
+ r"""
27
+ Pipeline for value-guided sampling from a diffusion model trained to predict sequences of states.
28
+
29
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
30
+ implemented for all pipelines (downloading, saving, running on a particular device, etc.).
31
+
32
+ Parameters:
33
+ value_function ([`UNet1DModel`]):
34
+ A specialized UNet for fine-tuning trajectories base on reward.
35
+ unet ([`UNet1DModel`]):
36
+ UNet architecture to denoise the encoded trajectories.
37
+ scheduler ([`SchedulerMixin`]):
38
+ A scheduler to be used in combination with `unet` to denoise the encoded trajectories. Default for this
39
+ application is [`DDPMScheduler`].
40
+ env ():
41
+ An environment following the OpenAI gym API to act in. For now only Hopper has pretrained models.
42
+ """
43
+
44
+ def __init__(
45
+ self,
46
+ value_function: UNet1DModel,
47
+ unet: UNet1DModel,
48
+ scheduler: DDPMScheduler,
49
+ env,
50
+ ):
51
+ super().__init__()
52
+
53
+ self.register_modules(value_function=value_function, unet=unet, scheduler=scheduler, env=env)
54
+
55
+ self.data = env.get_dataset()
56
+ self.means = {}
57
+ for key in self.data.keys():
58
+ try:
59
+ self.means[key] = self.data[key].mean()
60
+ except: # noqa: E722
61
+ pass
62
+ self.stds = {}
63
+ for key in self.data.keys():
64
+ try:
65
+ self.stds[key] = self.data[key].std()
66
+ except: # noqa: E722
67
+ pass
68
+ self.state_dim = env.observation_space.shape[0]
69
+ self.action_dim = env.action_space.shape[0]
70
+
71
+ def normalize(self, x_in, key):
72
+ return (x_in - self.means[key]) / self.stds[key]
73
+
74
+ def de_normalize(self, x_in, key):
75
+ return x_in * self.stds[key] + self.means[key]
76
+
77
+ def to_torch(self, x_in):
78
+ if isinstance(x_in, dict):
79
+ return {k: self.to_torch(v) for k, v in x_in.items()}
80
+ elif torch.is_tensor(x_in):
81
+ return x_in.to(self.unet.device)
82
+ return torch.tensor(x_in, device=self.unet.device)
83
+
84
+ def reset_x0(self, x_in, cond, act_dim):
85
+ for key, val in cond.items():
86
+ x_in[:, key, act_dim:] = val.clone()
87
+ return x_in
88
+
89
+ def run_diffusion(self, x, conditions, n_guide_steps, scale):
90
+ batch_size = x.shape[0]
91
+ y = None
92
+ for i in tqdm.tqdm(self.scheduler.timesteps):
93
+ # create batch of timesteps to pass into model
94
+ timesteps = torch.full((batch_size,), i, device=self.unet.device, dtype=torch.long)
95
+ for _ in range(n_guide_steps):
96
+ with torch.enable_grad():
97
+ x.requires_grad_()
98
+
99
+ # permute to match dimension for pre-trained models
100
+ y = self.value_function(x.permute(0, 2, 1), timesteps).sample
101
+ grad = torch.autograd.grad([y.sum()], [x])[0]
102
+
103
+ posterior_variance = self.scheduler._get_variance(i)
104
+ model_std = torch.exp(0.5 * posterior_variance)
105
+ grad = model_std * grad
106
+
107
+ grad[timesteps < 2] = 0
108
+ x = x.detach()
109
+ x = x + scale * grad
110
+ x = self.reset_x0(x, conditions, self.action_dim)
111
+
112
+ prev_x = self.unet(x.permute(0, 2, 1), timesteps).sample.permute(0, 2, 1)
113
+
114
+ # TODO: verify deprecation of this kwarg
115
+ x = self.scheduler.step(prev_x, i, x)["prev_sample"]
116
+
117
+ # apply conditions to the trajectory (set the initial state)
118
+ x = self.reset_x0(x, conditions, self.action_dim)
119
+ x = self.to_torch(x)
120
+ return x, y
121
+
122
+ def __call__(self, obs, batch_size=64, planning_horizon=32, n_guide_steps=2, scale=0.1):
123
+ # normalize the observations and create batch dimension
124
+ obs = self.normalize(obs, "observations")
125
+ obs = obs[None].repeat(batch_size, axis=0)
126
+
127
+ conditions = {0: self.to_torch(obs)}
128
+ shape = (batch_size, planning_horizon, self.state_dim + self.action_dim)
129
+
130
+ # generate initial noise and apply our conditions (to make the trajectories start at current state)
131
+ x1 = randn_tensor(shape, device=self.unet.device)
132
+ x = self.reset_x0(x1, conditions, self.action_dim)
133
+ x = self.to_torch(x)
134
+
135
+ # run the diffusion process
136
+ x, y = self.run_diffusion(x, conditions, n_guide_steps, scale)
137
+
138
+ # sort output trajectories by value
139
+ sorted_idx = y.argsort(0, descending=True).squeeze()
140
+ sorted_values = x[sorted_idx]
141
+ actions = sorted_values[:, :, : self.action_dim]
142
+ actions = actions.detach().cpu().numpy()
143
+ denorm_actions = self.de_normalize(actions, key="actions")
144
+
145
+ # select the action with the highest value
146
+ if y is not None:
147
+ selected_index = 0
148
+ else:
149
+ # if we didn't run value guiding, select a random action
150
+ selected_index = np.random.randint(0, batch_size)
151
+
152
+ denorm_actions = denorm_actions[selected_index, 0]
153
+ return denorm_actions
model/diffusers_c/image_processor.py ADDED
@@ -0,0 +1,990 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import math
16
+ import warnings
17
+ from typing import List, Optional, Tuple, Union
18
+
19
+ import numpy as np
20
+ import PIL.Image
21
+ import torch
22
+ import torch.nn.functional as F
23
+ from PIL import Image, ImageFilter, ImageOps
24
+
25
+ from .configuration_utils import ConfigMixin, register_to_config
26
+ from .utils import CONFIG_NAME, PIL_INTERPOLATION, deprecate
27
+
28
+
29
+ PipelineImageInput = Union[
30
+ PIL.Image.Image,
31
+ np.ndarray,
32
+ torch.FloatTensor,
33
+ List[PIL.Image.Image],
34
+ List[np.ndarray],
35
+ List[torch.FloatTensor],
36
+ ]
37
+
38
+ PipelineDepthInput = PipelineImageInput
39
+
40
+
41
+ class VaeImageProcessor(ConfigMixin):
42
+ """
43
+ Image processor for VAE.
44
+
45
+ Args:
46
+ do_resize (`bool`, *optional*, defaults to `True`):
47
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
48
+ `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
49
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
50
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
51
+ resample (`str`, *optional*, defaults to `lanczos`):
52
+ Resampling filter to use when resizing the image.
53
+ do_normalize (`bool`, *optional*, defaults to `True`):
54
+ Whether to normalize the image to [-1,1].
55
+ do_binarize (`bool`, *optional*, defaults to `False`):
56
+ Whether to binarize the image to 0/1.
57
+ do_convert_rgb (`bool`, *optional*, defaults to be `False`):
58
+ Whether to convert the images to RGB format.
59
+ do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
60
+ Whether to convert the images to grayscale format.
61
+ """
62
+
63
+ config_name = CONFIG_NAME
64
+
65
+ @register_to_config
66
+ def __init__(
67
+ self,
68
+ do_resize: bool = True,
69
+ vae_scale_factor: int = 8,
70
+ resample: str = "lanczos",
71
+ do_normalize: bool = True,
72
+ do_binarize: bool = False,
73
+ do_convert_rgb: bool = False,
74
+ do_convert_grayscale: bool = False,
75
+ ):
76
+ super().__init__()
77
+ if do_convert_rgb and do_convert_grayscale:
78
+ raise ValueError(
79
+ "`do_convert_rgb` and `do_convert_grayscale` can not both be set to `True`,"
80
+ " if you intended to convert the image into RGB format, please set `do_convert_grayscale = False`.",
81
+ " if you intended to convert the image into grayscale format, please set `do_convert_rgb = False`",
82
+ )
83
+ self.config.do_convert_rgb = False
84
+
85
+ @staticmethod
86
+ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
87
+ """
88
+ Convert a numpy image or a batch of images to a PIL image.
89
+ """
90
+ if images.ndim == 3:
91
+ images = images[None, ...]
92
+ images = (images * 255).round().astype("uint8")
93
+ if images.shape[-1] == 1:
94
+ # special case for grayscale (single channel) images
95
+ pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
96
+ else:
97
+ pil_images = [Image.fromarray(image) for image in images]
98
+
99
+ return pil_images
100
+
101
+ @staticmethod
102
+ def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:
103
+ """
104
+ Convert a PIL image or a list of PIL images to NumPy arrays.
105
+ """
106
+ if not isinstance(images, list):
107
+ images = [images]
108
+ images = [np.array(image).astype(np.float32) / 255.0 for image in images]
109
+ images = np.stack(images, axis=0)
110
+
111
+ return images
112
+
113
+ @staticmethod
114
+ def numpy_to_pt(images: np.ndarray) -> torch.FloatTensor:
115
+ """
116
+ Convert a NumPy image to a PyTorch tensor.
117
+ """
118
+ if images.ndim == 3:
119
+ images = images[..., None]
120
+
121
+ images = torch.from_numpy(images.transpose(0, 3, 1, 2))
122
+ return images
123
+
124
+ @staticmethod
125
+ def pt_to_numpy(images: torch.FloatTensor) -> np.ndarray:
126
+ """
127
+ Convert a PyTorch tensor to a NumPy image.
128
+ """
129
+ images = images.cpu().permute(0, 2, 3, 1).float().numpy()
130
+ return images
131
+
132
+ @staticmethod
133
+ def normalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
134
+ """
135
+ Normalize an image array to [-1,1].
136
+ """
137
+ return 2.0 * images - 1.0
138
+
139
+ @staticmethod
140
+ def denormalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
141
+ """
142
+ Denormalize an image array to [0,1].
143
+ """
144
+ return (images / 2 + 0.5).clamp(0, 1)
145
+
146
+ @staticmethod
147
+ def convert_to_rgb(image: PIL.Image.Image) -> PIL.Image.Image:
148
+ """
149
+ Converts a PIL image to RGB format.
150
+ """
151
+ image = image.convert("RGB")
152
+
153
+ return image
154
+
155
+ @staticmethod
156
+ def convert_to_grayscale(image: PIL.Image.Image) -> PIL.Image.Image:
157
+ """
158
+ Converts a PIL image to grayscale format.
159
+ """
160
+ image = image.convert("L")
161
+
162
+ return image
163
+
164
+ @staticmethod
165
+ def blur(image: PIL.Image.Image, blur_factor: int = 4) -> PIL.Image.Image:
166
+ """
167
+ Applies Gaussian blur to an image.
168
+ """
169
+ image = image.filter(ImageFilter.GaussianBlur(blur_factor))
170
+
171
+ return image
172
+
173
+ @staticmethod
174
+ def get_crop_region(mask_image: PIL.Image.Image, width: int, height: int, pad=0):
175
+ """
176
+ Finds a rectangular region that contains all masked ares in an image, and expands region to match the aspect ratio of the original image;
177
+ for example, if user drew mask in a 128x32 region, and the dimensions for processing are 512x512, the region will be expanded to 128x128.
178
+
179
+ Args:
180
+ mask_image (PIL.Image.Image): Mask image.
181
+ width (int): Width of the image to be processed.
182
+ height (int): Height of the image to be processed.
183
+ pad (int, optional): Padding to be added to the crop region. Defaults to 0.
184
+
185
+ Returns:
186
+ tuple: (x1, y1, x2, y2) represent a rectangular region that contains all masked ares in an image and matches the original aspect ratio.
187
+ """
188
+
189
+ mask_image = mask_image.convert("L")
190
+ mask = np.array(mask_image)
191
+
192
+ # 1. find a rectangular region that contains all masked ares in an image
193
+ h, w = mask.shape
194
+ crop_left = 0
195
+ for i in range(w):
196
+ if not (mask[:, i] == 0).all():
197
+ break
198
+ crop_left += 1
199
+
200
+ crop_right = 0
201
+ for i in reversed(range(w)):
202
+ if not (mask[:, i] == 0).all():
203
+ break
204
+ crop_right += 1
205
+
206
+ crop_top = 0
207
+ for i in range(h):
208
+ if not (mask[i] == 0).all():
209
+ break
210
+ crop_top += 1
211
+
212
+ crop_bottom = 0
213
+ for i in reversed(range(h)):
214
+ if not (mask[i] == 0).all():
215
+ break
216
+ crop_bottom += 1
217
+
218
+ # 2. add padding to the crop region
219
+ x1, y1, x2, y2 = (
220
+ int(max(crop_left - pad, 0)),
221
+ int(max(crop_top - pad, 0)),
222
+ int(min(w - crop_right + pad, w)),
223
+ int(min(h - crop_bottom + pad, h)),
224
+ )
225
+
226
+ # 3. expands crop region to match the aspect ratio of the image to be processed
227
+ ratio_crop_region = (x2 - x1) / (y2 - y1)
228
+ ratio_processing = width / height
229
+
230
+ if ratio_crop_region > ratio_processing:
231
+ desired_height = (x2 - x1) / ratio_processing
232
+ desired_height_diff = int(desired_height - (y2 - y1))
233
+ y1 -= desired_height_diff // 2
234
+ y2 += desired_height_diff - desired_height_diff // 2
235
+ if y2 >= mask_image.height:
236
+ diff = y2 - mask_image.height
237
+ y2 -= diff
238
+ y1 -= diff
239
+ if y1 < 0:
240
+ y2 -= y1
241
+ y1 -= y1
242
+ if y2 >= mask_image.height:
243
+ y2 = mask_image.height
244
+ else:
245
+ desired_width = (y2 - y1) * ratio_processing
246
+ desired_width_diff = int(desired_width - (x2 - x1))
247
+ x1 -= desired_width_diff // 2
248
+ x2 += desired_width_diff - desired_width_diff // 2
249
+ if x2 >= mask_image.width:
250
+ diff = x2 - mask_image.width
251
+ x2 -= diff
252
+ x1 -= diff
253
+ if x1 < 0:
254
+ x2 -= x1
255
+ x1 -= x1
256
+ if x2 >= mask_image.width:
257
+ x2 = mask_image.width
258
+
259
+ return x1, y1, x2, y2
260
+
261
+ def _resize_and_fill(
262
+ self,
263
+ image: PIL.Image.Image,
264
+ width: int,
265
+ height: int,
266
+ ) -> PIL.Image.Image:
267
+ """
268
+ Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image.
269
+
270
+ Args:
271
+ image: The image to resize.
272
+ width: The width to resize the image to.
273
+ height: The height to resize the image to.
274
+ """
275
+
276
+ ratio = width / height
277
+ src_ratio = image.width / image.height
278
+
279
+ src_w = width if ratio < src_ratio else image.width * height // image.height
280
+ src_h = height if ratio >= src_ratio else image.height * width // image.width
281
+
282
+ resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION["lanczos"])
283
+ res = Image.new("RGB", (width, height))
284
+ res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
285
+
286
+ if ratio < src_ratio:
287
+ fill_height = height // 2 - src_h // 2
288
+ if fill_height > 0:
289
+ res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))
290
+ res.paste(
291
+ resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)),
292
+ box=(0, fill_height + src_h),
293
+ )
294
+ elif ratio > src_ratio:
295
+ fill_width = width // 2 - src_w // 2
296
+ if fill_width > 0:
297
+ res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
298
+ res.paste(
299
+ resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)),
300
+ box=(fill_width + src_w, 0),
301
+ )
302
+
303
+ return res
304
+
305
+ def _resize_and_crop(
306
+ self,
307
+ image: PIL.Image.Image,
308
+ width: int,
309
+ height: int,
310
+ ) -> PIL.Image.Image:
311
+ """
312
+ Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess.
313
+
314
+ Args:
315
+ image: The image to resize.
316
+ width: The width to resize the image to.
317
+ height: The height to resize the image to.
318
+ """
319
+ ratio = width / height
320
+ src_ratio = image.width / image.height
321
+
322
+ src_w = width if ratio > src_ratio else image.width * height // image.height
323
+ src_h = height if ratio <= src_ratio else image.height * width // image.width
324
+
325
+ resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION["lanczos"])
326
+ res = Image.new("RGB", (width, height))
327
+ res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
328
+ return res
329
+
330
+ def resize(
331
+ self,
332
+ image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
333
+ height: int,
334
+ width: int,
335
+ resize_mode: str = "default", # "default", "fill", "crop"
336
+ ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]:
337
+ """
338
+ Resize image.
339
+
340
+ Args:
341
+ image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`):
342
+ The image input, can be a PIL image, numpy array or pytorch tensor.
343
+ height (`int`):
344
+ The height to resize to.
345
+ width (`int`):
346
+ The width to resize to.
347
+ resize_mode (`str`, *optional*, defaults to `default`):
348
+ The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit
349
+ within the specified width and height, and it may not maintaining the original aspect ratio.
350
+ If `fill`, will resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image
351
+ within the dimensions, filling empty with data from image.
352
+ If `crop`, will resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image
353
+ within the dimensions, cropping the excess.
354
+ Note that resize_mode `fill` and `crop` are only supported for PIL image input.
355
+
356
+ Returns:
357
+ `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
358
+ The resized image.
359
+ """
360
+ if resize_mode != "default" and not isinstance(image, PIL.Image.Image):
361
+ raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}")
362
+ if isinstance(image, PIL.Image.Image):
363
+ if resize_mode == "default":
364
+ image = image.resize((width, height), resample=PIL_INTERPOLATION[self.config.resample])
365
+ elif resize_mode == "fill":
366
+ image = self._resize_and_fill(image, width, height)
367
+ elif resize_mode == "crop":
368
+ image = self._resize_and_crop(image, width, height)
369
+ else:
370
+ raise ValueError(f"resize_mode {resize_mode} is not supported")
371
+
372
+ elif isinstance(image, torch.Tensor):
373
+ image = torch.nn.functional.interpolate(
374
+ image,
375
+ size=(height, width),
376
+ )
377
+ elif isinstance(image, np.ndarray):
378
+ image = self.numpy_to_pt(image)
379
+ image = torch.nn.functional.interpolate(
380
+ image,
381
+ size=(height, width),
382
+ )
383
+ image = self.pt_to_numpy(image)
384
+ return image
385
+
386
+ def binarize(self, image: PIL.Image.Image) -> PIL.Image.Image:
387
+ """
388
+ Create a mask.
389
+
390
+ Args:
391
+ image (`PIL.Image.Image`):
392
+ The image input, should be a PIL image.
393
+
394
+ Returns:
395
+ `PIL.Image.Image`:
396
+ The binarized image. Values less than 0.5 are set to 0, values greater than 0.5 are set to 1.
397
+ """
398
+ image[image < 0.5] = 0
399
+ image[image >= 0.5] = 1
400
+
401
+ return image
402
+
403
+ def get_default_height_width(
404
+ self,
405
+ image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
406
+ height: Optional[int] = None,
407
+ width: Optional[int] = None,
408
+ ) -> Tuple[int, int]:
409
+ """
410
+ This function return the height and width that are downscaled to the next integer multiple of
411
+ `vae_scale_factor`.
412
+
413
+ Args:
414
+ image(`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`):
415
+ The image input, can be a PIL image, numpy array or pytorch tensor. if it is a numpy array, should have
416
+ shape `[batch, height, width]` or `[batch, height, width, channel]` if it is a pytorch tensor, should
417
+ have shape `[batch, channel, height, width]`.
418
+ height (`int`, *optional*, defaults to `None`):
419
+ The height in preprocessed image. If `None`, will use the height of `image` input.
420
+ width (`int`, *optional*`, defaults to `None`):
421
+ The width in preprocessed. If `None`, will use the width of the `image` input.
422
+ """
423
+
424
+ if height is None:
425
+ if isinstance(image, PIL.Image.Image):
426
+ height = image.height
427
+ elif isinstance(image, torch.Tensor):
428
+ height = image.shape[2]
429
+ else:
430
+ height = image.shape[1]
431
+
432
+ if width is None:
433
+ if isinstance(image, PIL.Image.Image):
434
+ width = image.width
435
+ elif isinstance(image, torch.Tensor):
436
+ width = image.shape[3]
437
+ else:
438
+ width = image.shape[2]
439
+
440
+ width, height = (
441
+ x - x % self.config.vae_scale_factor for x in (width, height)
442
+ ) # resize to integer multiple of vae_scale_factor
443
+
444
+ return height, width
445
+
446
+ def preprocess(
447
+ self,
448
+ image: PipelineImageInput,
449
+ height: Optional[int] = None,
450
+ width: Optional[int] = None,
451
+ resize_mode: str = "default", # "default", "fill", "crop"
452
+ crops_coords: Optional[Tuple[int, int, int, int]] = None,
453
+ ) -> torch.Tensor:
454
+ """
455
+ Preprocess the image input.
456
+
457
+ Args:
458
+ image (`pipeline_image_input`):
459
+ The image input, accepted formats are PIL images, NumPy arrays, PyTorch tensors; Also accept list of supported formats.
460
+ height (`int`, *optional*, defaults to `None`):
461
+ The height in preprocessed image. If `None`, will use the `get_default_height_width()` to get default height.
462
+ width (`int`, *optional*`, defaults to `None`):
463
+ The width in preprocessed. If `None`, will use get_default_height_width()` to get the default width.
464
+ resize_mode (`str`, *optional*, defaults to `default`):
465
+ The resize mode, can be one of `default` or `fill`. If `default`, will resize the image to fit
466
+ within the specified width and height, and it may not maintaining the original aspect ratio.
467
+ If `fill`, will resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image
468
+ within the dimensions, filling empty with data from image.
469
+ If `crop`, will resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image
470
+ within the dimensions, cropping the excess.
471
+ Note that resize_mode `fill` and `crop` are only supported for PIL image input.
472
+ crops_coords (`List[Tuple[int, int, int, int]]`, *optional*, defaults to `None`):
473
+ The crop coordinates for each image in the batch. If `None`, will not crop the image.
474
+ """
475
+ supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
476
+
477
+ # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
478
+ if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3:
479
+ if isinstance(image, torch.Tensor):
480
+ # if image is a pytorch tensor could have 2 possible shapes:
481
+ # 1. batch x height x width: we should insert the channel dimension at position 1
482
+ # 2. channel x height x width: we should insert batch dimension at position 0,
483
+ # however, since both channel and batch dimension has same size 1, it is same to insert at position 1
484
+ # for simplicity, we insert a dimension of size 1 at position 1 for both cases
485
+ image = image.unsqueeze(1)
486
+ else:
487
+ # if it is a numpy array, it could have 2 possible shapes:
488
+ # 1. batch x height x width: insert channel dimension on last position
489
+ # 2. height x width x channel: insert batch dimension on first position
490
+ if image.shape[-1] == 1:
491
+ image = np.expand_dims(image, axis=0)
492
+ else:
493
+ image = np.expand_dims(image, axis=-1)
494
+
495
+ if isinstance(image, supported_formats):
496
+ image = [image]
497
+ elif not (isinstance(image, list) and all(isinstance(i, supported_formats) for i in image)):
498
+ raise ValueError(
499
+ f"Input is in incorrect format: {[type(i) for i in image]}. Currently, we only support {', '.join(supported_formats)}"
500
+ )
501
+
502
+ if isinstance(image[0], PIL.Image.Image):
503
+ if crops_coords is not None:
504
+ image = [i.crop(crops_coords) for i in image]
505
+ if self.config.do_resize:
506
+ height, width = self.get_default_height_width(image[0], height, width)
507
+ image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image]
508
+ if self.config.do_convert_rgb:
509
+ image = [self.convert_to_rgb(i) for i in image]
510
+ elif self.config.do_convert_grayscale:
511
+ image = [self.convert_to_grayscale(i) for i in image]
512
+ image = self.pil_to_numpy(image) # to np
513
+ image = self.numpy_to_pt(image) # to pt
514
+
515
+ elif isinstance(image[0], np.ndarray):
516
+ image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0)
517
+
518
+ image = self.numpy_to_pt(image)
519
+
520
+ height, width = self.get_default_height_width(image, height, width)
521
+ if self.config.do_resize:
522
+ image = self.resize(image, height, width)
523
+
524
+ elif isinstance(image[0], torch.Tensor):
525
+ image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
526
+
527
+ if self.config.do_convert_grayscale and image.ndim == 3:
528
+ image = image.unsqueeze(1)
529
+
530
+ channel = image.shape[1]
531
+ # don't need any preprocess if the image is latents
532
+ if channel == 4:
533
+ return image
534
+
535
+ height, width = self.get_default_height_width(image, height, width)
536
+ if self.config.do_resize:
537
+ image = self.resize(image, height, width)
538
+
539
+ # expected range [0,1], normalize to [-1,1]
540
+ do_normalize = self.config.do_normalize
541
+ if do_normalize and image.min() < 0:
542
+ warnings.warn(
543
+ "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
544
+ f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]",
545
+ FutureWarning,
546
+ )
547
+ do_normalize = False
548
+
549
+ if do_normalize:
550
+ image = self.normalize(image)
551
+
552
+ if self.config.do_binarize:
553
+ image = self.binarize(image)
554
+
555
+ return image
556
+
557
+ def postprocess(
558
+ self,
559
+ image: torch.FloatTensor,
560
+ output_type: str = "pil",
561
+ do_denormalize: Optional[List[bool]] = None,
562
+ ) -> Union[PIL.Image.Image, np.ndarray, torch.FloatTensor]:
563
+ """
564
+ Postprocess the image output from tensor to `output_type`.
565
+
566
+ Args:
567
+ image (`torch.FloatTensor`):
568
+ The image input, should be a pytorch tensor with shape `B x C x H x W`.
569
+ output_type (`str`, *optional*, defaults to `pil`):
570
+ The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
571
+ do_denormalize (`List[bool]`, *optional*, defaults to `None`):
572
+ Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
573
+ `VaeImageProcessor` config.
574
+
575
+ Returns:
576
+ `PIL.Image.Image`, `np.ndarray` or `torch.FloatTensor`:
577
+ The postprocessed image.
578
+ """
579
+ if not isinstance(image, torch.Tensor):
580
+ raise ValueError(
581
+ f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
582
+ )
583
+ if output_type not in ["latent", "pt", "np", "pil"]:
584
+ deprecation_message = (
585
+ f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
586
+ "`pil`, `np`, `pt`, `latent`"
587
+ )
588
+ deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
589
+ output_type = "np"
590
+
591
+ if output_type == "latent":
592
+ return image
593
+
594
+ if do_denormalize is None:
595
+ do_denormalize = [self.config.do_normalize] * image.shape[0]
596
+
597
+ image = torch.stack(
598
+ [self.denormalize(image[i]) if do_denormalize[i] else image[i] for i in range(image.shape[0])]
599
+ )
600
+
601
+ if output_type == "pt":
602
+ return image
603
+
604
+ image = self.pt_to_numpy(image)
605
+
606
+ if output_type == "np":
607
+ return image
608
+
609
+ if output_type == "pil":
610
+ return self.numpy_to_pil(image)
611
+
612
+ def apply_overlay(
613
+ self,
614
+ mask: PIL.Image.Image,
615
+ init_image: PIL.Image.Image,
616
+ image: PIL.Image.Image,
617
+ crop_coords: Optional[Tuple[int, int, int, int]] = None,
618
+ ) -> PIL.Image.Image:
619
+ """
620
+ overlay the inpaint output to the original image
621
+ """
622
+
623
+ width, height = image.width, image.height
624
+
625
+ init_image = self.resize(init_image, width=width, height=height)
626
+ mask = self.resize(mask, width=width, height=height)
627
+
628
+ init_image_masked = PIL.Image.new("RGBa", (width, height))
629
+ init_image_masked.paste(init_image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert("L")))
630
+ init_image_masked = init_image_masked.convert("RGBA")
631
+
632
+ if crop_coords is not None:
633
+ x, y, x2, y2 = crop_coords
634
+ w = x2 - x
635
+ h = y2 - y
636
+ base_image = PIL.Image.new("RGBA", (width, height))
637
+ image = self.resize(image, height=h, width=w, resize_mode="crop")
638
+ base_image.paste(image, (x, y))
639
+ image = base_image.convert("RGB")
640
+
641
+ image = image.convert("RGBA")
642
+ image.alpha_composite(init_image_masked)
643
+ image = image.convert("RGB")
644
+
645
+ return image
646
+
647
+
648
+ class VaeImageProcessorLDM3D(VaeImageProcessor):
649
+ """
650
+ Image processor for VAE LDM3D.
651
+
652
+ Args:
653
+ do_resize (`bool`, *optional*, defaults to `True`):
654
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
655
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
656
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
657
+ resample (`str`, *optional*, defaults to `lanczos`):
658
+ Resampling filter to use when resizing the image.
659
+ do_normalize (`bool`, *optional*, defaults to `True`):
660
+ Whether to normalize the image to [-1,1].
661
+ """
662
+
663
+ config_name = CONFIG_NAME
664
+
665
+ @register_to_config
666
+ def __init__(
667
+ self,
668
+ do_resize: bool = True,
669
+ vae_scale_factor: int = 8,
670
+ resample: str = "lanczos",
671
+ do_normalize: bool = True,
672
+ ):
673
+ super().__init__()
674
+
675
+ @staticmethod
676
+ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
677
+ """
678
+ Convert a NumPy image or a batch of images to a PIL image.
679
+ """
680
+ if images.ndim == 3:
681
+ images = images[None, ...]
682
+ images = (images * 255).round().astype("uint8")
683
+ if images.shape[-1] == 1:
684
+ # special case for grayscale (single channel) images
685
+ pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
686
+ else:
687
+ pil_images = [Image.fromarray(image[:, :, :3]) for image in images]
688
+
689
+ return pil_images
690
+
691
+ @staticmethod
692
+ def depth_pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:
693
+ """
694
+ Convert a PIL image or a list of PIL images to NumPy arrays.
695
+ """
696
+ if not isinstance(images, list):
697
+ images = [images]
698
+
699
+ images = [np.array(image).astype(np.float32) / (2**16 - 1) for image in images]
700
+ images = np.stack(images, axis=0)
701
+ return images
702
+
703
+ @staticmethod
704
+ def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
705
+ """
706
+ Args:
707
+ image: RGB-like depth image
708
+
709
+ Returns: depth map
710
+
711
+ """
712
+ return image[:, :, 1] * 2**8 + image[:, :, 2]
713
+
714
+ def numpy_to_depth(self, images: np.ndarray) -> List[PIL.Image.Image]:
715
+ """
716
+ Convert a NumPy depth image or a batch of images to a PIL image.
717
+ """
718
+ if images.ndim == 3:
719
+ images = images[None, ...]
720
+ images_depth = images[:, :, :, 3:]
721
+ if images.shape[-1] == 6:
722
+ images_depth = (images_depth * 255).round().astype("uint8")
723
+ pil_images = [
724
+ Image.fromarray(self.rgblike_to_depthmap(image_depth), mode="I;16") for image_depth in images_depth
725
+ ]
726
+ elif images.shape[-1] == 4:
727
+ images_depth = (images_depth * 65535.0).astype(np.uint16)
728
+ pil_images = [Image.fromarray(image_depth, mode="I;16") for image_depth in images_depth]
729
+ else:
730
+ raise Exception("Not supported")
731
+
732
+ return pil_images
733
+
734
+ def postprocess(
735
+ self,
736
+ image: torch.FloatTensor,
737
+ output_type: str = "pil",
738
+ do_denormalize: Optional[List[bool]] = None,
739
+ ) -> Union[PIL.Image.Image, np.ndarray, torch.FloatTensor]:
740
+ """
741
+ Postprocess the image output from tensor to `output_type`.
742
+
743
+ Args:
744
+ image (`torch.FloatTensor`):
745
+ The image input, should be a pytorch tensor with shape `B x C x H x W`.
746
+ output_type (`str`, *optional*, defaults to `pil`):
747
+ The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
748
+ do_denormalize (`List[bool]`, *optional*, defaults to `None`):
749
+ Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
750
+ `VaeImageProcessor` config.
751
+
752
+ Returns:
753
+ `PIL.Image.Image`, `np.ndarray` or `torch.FloatTensor`:
754
+ The postprocessed image.
755
+ """
756
+ if not isinstance(image, torch.Tensor):
757
+ raise ValueError(
758
+ f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
759
+ )
760
+ if output_type not in ["latent", "pt", "np", "pil"]:
761
+ deprecation_message = (
762
+ f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
763
+ "`pil`, `np`, `pt`, `latent`"
764
+ )
765
+ deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
766
+ output_type = "np"
767
+
768
+ if do_denormalize is None:
769
+ do_denormalize = [self.config.do_normalize] * image.shape[0]
770
+
771
+ image = torch.stack(
772
+ [self.denormalize(image[i]) if do_denormalize[i] else image[i] for i in range(image.shape[0])]
773
+ )
774
+
775
+ image = self.pt_to_numpy(image)
776
+
777
+ if output_type == "np":
778
+ if image.shape[-1] == 6:
779
+ image_depth = np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0)
780
+ else:
781
+ image_depth = image[:, :, :, 3:]
782
+ return image[:, :, :, :3], image_depth
783
+
784
+ if output_type == "pil":
785
+ return self.numpy_to_pil(image), self.numpy_to_depth(image)
786
+ else:
787
+ raise Exception(f"This type {output_type} is not supported")
788
+
789
+ def preprocess(
790
+ self,
791
+ rgb: Union[torch.FloatTensor, PIL.Image.Image, np.ndarray],
792
+ depth: Union[torch.FloatTensor, PIL.Image.Image, np.ndarray],
793
+ height: Optional[int] = None,
794
+ width: Optional[int] = None,
795
+ target_res: Optional[int] = None,
796
+ ) -> torch.Tensor:
797
+ """
798
+ Preprocess the image input. Accepted formats are PIL images, NumPy arrays or PyTorch tensors.
799
+ """
800
+ supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
801
+
802
+ # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
803
+ if self.config.do_convert_grayscale and isinstance(rgb, (torch.Tensor, np.ndarray)) and rgb.ndim == 3:
804
+ raise Exception("This is not yet supported")
805
+
806
+ if isinstance(rgb, supported_formats):
807
+ rgb = [rgb]
808
+ depth = [depth]
809
+ elif not (isinstance(rgb, list) and all(isinstance(i, supported_formats) for i in rgb)):
810
+ raise ValueError(
811
+ f"Input is in incorrect format: {[type(i) for i in rgb]}. Currently, we only support {', '.join(supported_formats)}"
812
+ )
813
+
814
+ if isinstance(rgb[0], PIL.Image.Image):
815
+ if self.config.do_convert_rgb:
816
+ raise Exception("This is not yet supported")
817
+ # rgb = [self.convert_to_rgb(i) for i in rgb]
818
+ # depth = [self.convert_to_depth(i) for i in depth] #TODO define convert_to_depth
819
+ if self.config.do_resize or target_res:
820
+ height, width = self.get_default_height_width(rgb[0], height, width) if not target_res else target_res
821
+ rgb = [self.resize(i, height, width) for i in rgb]
822
+ depth = [self.resize(i, height, width) for i in depth]
823
+ rgb = self.pil_to_numpy(rgb) # to np
824
+ rgb = self.numpy_to_pt(rgb) # to pt
825
+
826
+ depth = self.depth_pil_to_numpy(depth) # to np
827
+ depth = self.numpy_to_pt(depth) # to pt
828
+
829
+ elif isinstance(rgb[0], np.ndarray):
830
+ rgb = np.concatenate(rgb, axis=0) if rgb[0].ndim == 4 else np.stack(rgb, axis=0)
831
+ rgb = self.numpy_to_pt(rgb)
832
+ height, width = self.get_default_height_width(rgb, height, width)
833
+ if self.config.do_resize:
834
+ rgb = self.resize(rgb, height, width)
835
+
836
+ depth = np.concatenate(depth, axis=0) if rgb[0].ndim == 4 else np.stack(depth, axis=0)
837
+ depth = self.numpy_to_pt(depth)
838
+ height, width = self.get_default_height_width(depth, height, width)
839
+ if self.config.do_resize:
840
+ depth = self.resize(depth, height, width)
841
+
842
+ elif isinstance(rgb[0], torch.Tensor):
843
+ raise Exception("This is not yet supported")
844
+ # rgb = torch.cat(rgb, axis=0) if rgb[0].ndim == 4 else torch.stack(rgb, axis=0)
845
+
846
+ # if self.config.do_convert_grayscale and rgb.ndim == 3:
847
+ # rgb = rgb.unsqueeze(1)
848
+
849
+ # channel = rgb.shape[1]
850
+
851
+ # height, width = self.get_default_height_width(rgb, height, width)
852
+ # if self.config.do_resize:
853
+ # rgb = self.resize(rgb, height, width)
854
+
855
+ # depth = torch.cat(depth, axis=0) if depth[0].ndim == 4 else torch.stack(depth, axis=0)
856
+
857
+ # if self.config.do_convert_grayscale and depth.ndim == 3:
858
+ # depth = depth.unsqueeze(1)
859
+
860
+ # channel = depth.shape[1]
861
+ # # don't need any preprocess if the image is latents
862
+ # if depth == 4:
863
+ # return rgb, depth
864
+
865
+ # height, width = self.get_default_height_width(depth, height, width)
866
+ # if self.config.do_resize:
867
+ # depth = self.resize(depth, height, width)
868
+ # expected range [0,1], normalize to [-1,1]
869
+ do_normalize = self.config.do_normalize
870
+ if rgb.min() < 0 and do_normalize:
871
+ warnings.warn(
872
+ "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
873
+ f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{rgb.min()},{rgb.max()}]",
874
+ FutureWarning,
875
+ )
876
+ do_normalize = False
877
+
878
+ if do_normalize:
879
+ rgb = self.normalize(rgb)
880
+ depth = self.normalize(depth)
881
+
882
+ if self.config.do_binarize:
883
+ rgb = self.binarize(rgb)
884
+ depth = self.binarize(depth)
885
+
886
+ return rgb, depth
887
+
888
+
889
+ class IPAdapterMaskProcessor(VaeImageProcessor):
890
+ """
891
+ Image processor for IP Adapter image masks.
892
+
893
+ Args:
894
+ do_resize (`bool`, *optional*, defaults to `True`):
895
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
896
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
897
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
898
+ resample (`str`, *optional*, defaults to `lanczos`):
899
+ Resampling filter to use when resizing the image.
900
+ do_normalize (`bool`, *optional*, defaults to `False`):
901
+ Whether to normalize the image to [-1,1].
902
+ do_binarize (`bool`, *optional*, defaults to `True`):
903
+ Whether to binarize the image to 0/1.
904
+ do_convert_grayscale (`bool`, *optional*, defaults to be `True`):
905
+ Whether to convert the images to grayscale format.
906
+
907
+ """
908
+
909
+ config_name = CONFIG_NAME
910
+
911
+ @register_to_config
912
+ def __init__(
913
+ self,
914
+ do_resize: bool = True,
915
+ vae_scale_factor: int = 8,
916
+ resample: str = "lanczos",
917
+ do_normalize: bool = False,
918
+ do_binarize: bool = True,
919
+ do_convert_grayscale: bool = True,
920
+ ):
921
+ super().__init__(
922
+ do_resize=do_resize,
923
+ vae_scale_factor=vae_scale_factor,
924
+ resample=resample,
925
+ do_normalize=do_normalize,
926
+ do_binarize=do_binarize,
927
+ do_convert_grayscale=do_convert_grayscale,
928
+ )
929
+
930
+ @staticmethod
931
+ def downsample(mask: torch.FloatTensor, batch_size: int, num_queries: int, value_embed_dim: int):
932
+ """
933
+ Downsamples the provided mask tensor to match the expected dimensions for scaled dot-product attention.
934
+ If the aspect ratio of the mask does not match the aspect ratio of the output image, a warning is issued.
935
+
936
+ Args:
937
+ mask (`torch.FloatTensor`):
938
+ The input mask tensor generated with `IPAdapterMaskProcessor.preprocess()`.
939
+ batch_size (`int`):
940
+ The batch size.
941
+ num_queries (`int`):
942
+ The number of queries.
943
+ value_embed_dim (`int`):
944
+ The dimensionality of the value embeddings.
945
+
946
+ Returns:
947
+ `torch.FloatTensor`:
948
+ The downsampled mask tensor.
949
+
950
+ """
951
+ o_h = mask.shape[1]
952
+ o_w = mask.shape[2]
953
+ ratio = o_w / o_h
954
+ mask_h = int(math.sqrt(num_queries / ratio))
955
+ mask_h = int(mask_h) + int((num_queries % int(mask_h)) != 0)
956
+ mask_w = num_queries // mask_h
957
+
958
+ mask_downsample = F.interpolate(mask.unsqueeze(0), size=(mask_h, mask_w), mode="bicubic").squeeze(0)
959
+
960
+ # Repeat batch_size times
961
+ if mask_downsample.shape[0] < batch_size:
962
+ mask_downsample = mask_downsample.repeat(batch_size, 1, 1)
963
+
964
+ mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1)
965
+
966
+ downsampled_area = mask_h * mask_w
967
+ # If the output image and the mask do not have the same aspect ratio, tensor shapes will not match
968
+ # Pad tensor if downsampled_mask.shape[1] is smaller than num_queries
969
+ if downsampled_area < num_queries:
970
+ warnings.warn(
971
+ "The aspect ratio of the mask does not match the aspect ratio of the output image. "
972
+ "Please update your masks or adjust the output size for optimal performance.",
973
+ UserWarning,
974
+ )
975
+ mask_downsample = F.pad(mask_downsample, (0, num_queries - mask_downsample.shape[1]), value=0.0)
976
+ # Discard last embeddings if downsampled_mask.shape[1] is bigger than num_queries
977
+ if downsampled_area > num_queries:
978
+ warnings.warn(
979
+ "The aspect ratio of the mask does not match the aspect ratio of the output image. "
980
+ "Please update your masks or adjust the output size for optimal performance.",
981
+ UserWarning,
982
+ )
983
+ mask_downsample = mask_downsample[:, :num_queries]
984
+
985
+ # Repeat last dimension to match SDPA output shape
986
+ mask_downsample = mask_downsample.view(mask_downsample.shape[0], mask_downsample.shape[1], 1).repeat(
987
+ 1, 1, value_embed_dim
988
+ )
989
+
990
+ return mask_downsample
model/diffusers_c/loaders/__init__.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import TYPE_CHECKING
2
+
3
+ from ..utils import DIFFUSERS_SLOW_IMPORT, _LazyModule, deprecate
4
+ from ..utils.import_utils import is_peft_available, is_torch_available, is_transformers_available
5
+
6
+
7
+ def text_encoder_lora_state_dict(text_encoder):
8
+ deprecate(
9
+ "text_encoder_load_state_dict in `models`",
10
+ "0.27.0",
11
+ "`text_encoder_lora_state_dict` is deprecated and will be removed in 0.27.0. Make sure to retrieve the weights using `get_peft_model`. See https://huggingface.co/docs/peft/v0.6.2/en/quicktour#peftmodel for more information.",
12
+ )
13
+ state_dict = {}
14
+
15
+ for name, module in text_encoder_attn_modules(text_encoder):
16
+ for k, v in module.q_proj.lora_linear_layer.state_dict().items():
17
+ state_dict[f"{name}.q_proj.lora_linear_layer.{k}"] = v
18
+
19
+ for k, v in module.k_proj.lora_linear_layer.state_dict().items():
20
+ state_dict[f"{name}.k_proj.lora_linear_layer.{k}"] = v
21
+
22
+ for k, v in module.v_proj.lora_linear_layer.state_dict().items():
23
+ state_dict[f"{name}.v_proj.lora_linear_layer.{k}"] = v
24
+
25
+ for k, v in module.out_proj.lora_linear_layer.state_dict().items():
26
+ state_dict[f"{name}.out_proj.lora_linear_layer.{k}"] = v
27
+
28
+ return state_dict
29
+
30
+
31
+ if is_transformers_available():
32
+
33
+ def text_encoder_attn_modules(text_encoder):
34
+ deprecate(
35
+ "text_encoder_attn_modules in `models`",
36
+ "0.27.0",
37
+ "`text_encoder_lora_state_dict` is deprecated and will be removed in 0.27.0. Make sure to retrieve the weights using `get_peft_model`. See https://huggingface.co/docs/peft/v0.6.2/en/quicktour#peftmodel for more information.",
38
+ )
39
+ from transformers import CLIPTextModel, CLIPTextModelWithProjection
40
+
41
+ attn_modules = []
42
+
43
+ if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)):
44
+ for i, layer in enumerate(text_encoder.text_model.encoder.layers):
45
+ name = f"text_model.encoder.layers.{i}.self_attn"
46
+ mod = layer.self_attn
47
+ attn_modules.append((name, mod))
48
+ else:
49
+ raise ValueError(f"do not know how to get attention modules for: {text_encoder.__class__.__name__}")
50
+
51
+ return attn_modules
52
+
53
+
54
+ _import_structure = {}
55
+
56
+ if is_torch_available():
57
+ _import_structure["autoencoder"] = ["FromOriginalVAEMixin"]
58
+
59
+ _import_structure["controlnet"] = ["FromOriginalControlNetMixin"]
60
+ _import_structure["unet"] = ["UNet2DConditionLoadersMixin"]
61
+ _import_structure["utils"] = ["AttnProcsLayers"]
62
+ if is_transformers_available():
63
+ _import_structure["single_file"] = ["FromSingleFileMixin"]
64
+ _import_structure["lora"] = ["LoraLoaderMixin", "StableDiffusionXLLoraLoaderMixin"]
65
+ _import_structure["textual_inversion"] = ["TextualInversionLoaderMixin"]
66
+ _import_structure["ip_adapter"] = ["IPAdapterMixin"]
67
+
68
+ _import_structure["peft"] = ["PeftAdapterMixin"]
69
+
70
+
71
+ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
72
+ if is_torch_available():
73
+ from .autoencoder import FromOriginalVAEMixin
74
+ from .controlnet import FromOriginalControlNetMixin
75
+ from .unet import UNet2DConditionLoadersMixin
76
+ from .utils import AttnProcsLayers
77
+
78
+ if is_transformers_available():
79
+ from .ip_adapter import IPAdapterMixin
80
+ from .lora import LoraLoaderMixin, StableDiffusionXLLoraLoaderMixin
81
+ from .single_file import FromSingleFileMixin
82
+ from .textual_inversion import TextualInversionLoaderMixin
83
+
84
+ from .peft import PeftAdapterMixin
85
+ else:
86
+ import sys
87
+
88
+ sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
model/diffusers_c/loaders/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (2.82 kB). View file
 
model/diffusers_c/loaders/__pycache__/__init__.cpython-39.pyc ADDED
Binary file (2.84 kB). View file
 
model/diffusers_c/loaders/__pycache__/autoencoder.cpython-310.pyc ADDED
Binary file (6.4 kB). View file
 
model/diffusers_c/loaders/__pycache__/autoencoder.cpython-39.pyc ADDED
Binary file (6.42 kB). View file
 
model/diffusers_c/loaders/__pycache__/controlnet.cpython-39.pyc ADDED
Binary file (5.83 kB). View file
 
model/diffusers_c/loaders/__pycache__/ip_adapter.cpython-39.pyc ADDED
Binary file (9.86 kB). View file
 
model/diffusers_c/loaders/__pycache__/lora.cpython-39.pyc ADDED
Binary file (44.1 kB). View file
 
model/diffusers_c/loaders/__pycache__/lora_conversion_utils.cpython-39.pyc ADDED
Binary file (7.2 kB). View file
 
model/diffusers_c/loaders/__pycache__/peft.cpython-310.pyc ADDED
Binary file (6.08 kB). View file
 
model/diffusers_c/loaders/__pycache__/peft.cpython-39.pyc ADDED
Binary file (6.18 kB). View file
 
model/diffusers_c/loaders/__pycache__/single_file.cpython-39.pyc ADDED
Binary file (9.77 kB). View file
 
model/diffusers_c/loaders/__pycache__/single_file_utils.cpython-310.pyc ADDED
Binary file (36.4 kB). View file
 
model/diffusers_c/loaders/__pycache__/single_file_utils.cpython-39.pyc ADDED
Binary file (37.5 kB). View file
 
model/diffusers_c/loaders/__pycache__/textual_inversion.cpython-39.pyc ADDED
Binary file (19.3 kB). View file
 
model/diffusers_c/loaders/__pycache__/unet.cpython-310.pyc ADDED
Binary file (27.6 kB). View file
 
model/diffusers_c/loaders/__pycache__/unet.cpython-39.pyc ADDED
Binary file (27.6 kB). View file
 
model/diffusers_c/loaders/__pycache__/utils.cpython-310.pyc ADDED
Binary file (2.08 kB). View file
 
model/diffusers_c/loaders/__pycache__/utils.cpython-39.pyc ADDED
Binary file (2.09 kB). View file
 
model/diffusers_c/loaders/autoencoder.py ADDED
@@ -0,0 +1,146 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from huggingface_hub.utils import validate_hf_hub_args
16
+
17
+ from .single_file_utils import (
18
+ create_diffusers_vae_model_from_ldm,
19
+ fetch_ldm_config_and_checkpoint,
20
+ )
21
+
22
+
23
+ class FromOriginalVAEMixin:
24
+ """
25
+ Load pretrained AutoencoderKL weights saved in the `.ckpt` or `.safetensors` format into a [`AutoencoderKL`].
26
+ """
27
+
28
+ @classmethod
29
+ @validate_hf_hub_args
30
+ def from_single_file(cls, pretrained_model_link_or_path, **kwargs):
31
+ r"""
32
+ Instantiate a [`AutoencoderKL`] from pretrained ControlNet weights saved in the original `.ckpt` or
33
+ `.safetensors` format. The pipeline is set in evaluation mode (`model.eval()`) by default.
34
+
35
+ Parameters:
36
+ pretrained_model_link_or_path (`str` or `os.PathLike`, *optional*):
37
+ Can be either:
38
+ - A link to the `.ckpt` file (for example
39
+ `"https://huggingface.co/<repo_id>/blob/main/<path_to_file>.ckpt"`) on the Hub.
40
+ - A path to a *file* containing all pipeline weights.
41
+ config_file (`str`, *optional*):
42
+ Filepath to the configuration YAML file associated with the model. If not provided it will default to:
43
+ https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
44
+ torch_dtype (`str` or `torch.dtype`, *optional*):
45
+ Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the
46
+ dtype is automatically derived from the model's weights.
47
+ force_download (`bool`, *optional*, defaults to `False`):
48
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
49
+ cached versions if they exist.
50
+ cache_dir (`Union[str, os.PathLike]`, *optional*):
51
+ Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
52
+ is not used.
53
+ resume_download (`bool`, *optional*, defaults to `False`):
54
+ Whether or not to resume downloading the model weights and configuration files. If set to `False`, any
55
+ incompletely downloaded files are deleted.
56
+ proxies (`Dict[str, str]`, *optional*):
57
+ A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
58
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
59
+ local_files_only (`bool`, *optional*, defaults to `False`):
60
+ Whether to only load local model weights and configuration files or not. If set to True, the model
61
+ won't be downloaded from the Hub.
62
+ token (`str` or *bool*, *optional*):
63
+ The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
64
+ `diffusers-cli login` (stored in `~/.huggingface`) is used.
65
+ revision (`str`, *optional*, defaults to `"main"`):
66
+ The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
67
+ allowed by Git.
68
+ image_size (`int`, *optional*, defaults to 512):
69
+ The image size the model was trained on. Use 512 for all Stable Diffusion v1 models and the Stable
70
+ Diffusion v2 base model. Use 768 for Stable Diffusion v2.
71
+ scaling_factor (`float`, *optional*, defaults to 0.18215):
72
+ The component-wise standard deviation of the trained latent space computed using the first batch of the
73
+ training set. This is used to scale the latent space to have unit variance when training the diffusion
74
+ model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
75
+ diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z
76
+ = 1 / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution
77
+ Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
78
+ kwargs (remaining dictionary of keyword arguments, *optional*):
79
+ Can be used to overwrite load and saveable variables (for example the pipeline components of the
80
+ specific pipeline class). The overwritten components are directly passed to the pipelines `__init__`
81
+ method. See example below for more information.
82
+
83
+ <Tip warning={true}>
84
+
85
+ Make sure to pass both `image_size` and `scaling_factor` to `from_single_file()` if you're loading
86
+ a VAE from SDXL or a Stable Diffusion v2 model or higher.
87
+
88
+ </Tip>
89
+
90
+ Examples:
91
+
92
+ ```py
93
+ from diffusers import AutoencoderKL
94
+
95
+ url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors" # can also be local file
96
+ model = AutoencoderKL.from_single_file(url)
97
+ ```
98
+ """
99
+
100
+ original_config_file = kwargs.pop("original_config_file", None)
101
+ config_file = kwargs.pop("config_file", None)
102
+ resume_download = kwargs.pop("resume_download", False)
103
+ force_download = kwargs.pop("force_download", False)
104
+ proxies = kwargs.pop("proxies", None)
105
+ token = kwargs.pop("token", None)
106
+ cache_dir = kwargs.pop("cache_dir", None)
107
+ local_files_only = kwargs.pop("local_files_only", None)
108
+ revision = kwargs.pop("revision", None)
109
+ torch_dtype = kwargs.pop("torch_dtype", None)
110
+
111
+ class_name = cls.__name__
112
+
113
+ if (config_file is not None) and (original_config_file is not None):
114
+ raise ValueError(
115
+ "You cannot pass both `config_file` and `original_config_file` to `from_single_file`. Please use only one of these arguments."
116
+ )
117
+
118
+ original_config_file = original_config_file or config_file
119
+ original_config, checkpoint = fetch_ldm_config_and_checkpoint(
120
+ pretrained_model_link_or_path=pretrained_model_link_or_path,
121
+ class_name=class_name,
122
+ original_config_file=original_config_file,
123
+ resume_download=resume_download,
124
+ force_download=force_download,
125
+ proxies=proxies,
126
+ token=token,
127
+ revision=revision,
128
+ local_files_only=local_files_only,
129
+ cache_dir=cache_dir,
130
+ )
131
+
132
+ image_size = kwargs.pop("image_size", None)
133
+ scaling_factor = kwargs.pop("scaling_factor", None)
134
+ component = create_diffusers_vae_model_from_ldm(
135
+ class_name,
136
+ original_config,
137
+ checkpoint,
138
+ image_size=image_size,
139
+ scaling_factor=scaling_factor,
140
+ torch_dtype=torch_dtype,
141
+ )
142
+ vae = component["vae"]
143
+ if torch_dtype is not None:
144
+ vae = vae.to(torch_dtype)
145
+
146
+ return vae
model/diffusers_c/loaders/controlnet.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from huggingface_hub.utils import validate_hf_hub_args
16
+
17
+ from .single_file_utils import (
18
+ create_diffusers_controlnet_model_from_ldm,
19
+ fetch_ldm_config_and_checkpoint,
20
+ )
21
+
22
+
23
+ class FromOriginalControlNetMixin:
24
+ """
25
+ Load pretrained ControlNet weights saved in the `.ckpt` or `.safetensors` format into a [`ControlNetModel`].
26
+ """
27
+
28
+ @classmethod
29
+ @validate_hf_hub_args
30
+ def from_single_file(cls, pretrained_model_link_or_path, **kwargs):
31
+ r"""
32
+ Instantiate a [`ControlNetModel`] from pretrained ControlNet weights saved in the original `.ckpt` or
33
+ `.safetensors` format. The pipeline is set in evaluation mode (`model.eval()`) by default.
34
+
35
+ Parameters:
36
+ pretrained_model_link_or_path (`str` or `os.PathLike`, *optional*):
37
+ Can be either:
38
+ - A link to the `.ckpt` file (for example
39
+ `"https://huggingface.co/<repo_id>/blob/main/<path_to_file>.ckpt"`) on the Hub.
40
+ - A path to a *file* containing all pipeline weights.
41
+ config_file (`str`, *optional*):
42
+ Filepath to the configuration YAML file associated with the model. If not provided it will default to:
43
+ https://raw.githubusercontent.com/lllyasviel/ControlNet/main/models/cldm_v15.yaml
44
+ torch_dtype (`str` or `torch.dtype`, *optional*):
45
+ Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the
46
+ dtype is automatically derived from the model's weights.
47
+ force_download (`bool`, *optional*, defaults to `False`):
48
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
49
+ cached versions if they exist.
50
+ cache_dir (`Union[str, os.PathLike]`, *optional*):
51
+ Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
52
+ is not used.
53
+ resume_download (`bool`, *optional*, defaults to `False`):
54
+ Whether or not to resume downloading the model weights and configuration files. If set to `False`, any
55
+ incompletely downloaded files are deleted.
56
+ proxies (`Dict[str, str]`, *optional*):
57
+ A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
58
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
59
+ local_files_only (`bool`, *optional*, defaults to `False`):
60
+ Whether to only load local model weights and configuration files or not. If set to True, the model
61
+ won't be downloaded from the Hub.
62
+ token (`str` or *bool*, *optional*):
63
+ The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
64
+ `diffusers-cli login` (stored in `~/.huggingface`) is used.
65
+ revision (`str`, *optional*, defaults to `"main"`):
66
+ The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
67
+ allowed by Git.
68
+ image_size (`int`, *optional*, defaults to 512):
69
+ The image size the model was trained on. Use 512 for all Stable Diffusion v1 models and the Stable
70
+ Diffusion v2 base model. Use 768 for Stable Diffusion v2.
71
+ upcast_attention (`bool`, *optional*, defaults to `None`):
72
+ Whether the attention computation should always be upcasted.
73
+ kwargs (remaining dictionary of keyword arguments, *optional*):
74
+ Can be used to overwrite load and saveable variables (for example the pipeline components of the
75
+ specific pipeline class). The overwritten components are directly passed to the pipelines `__init__`
76
+ method. See example below for more information.
77
+
78
+ Examples:
79
+
80
+ ```py
81
+ from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
82
+
83
+ url = "https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_canny.pth" # can also be a local path
84
+ model = ControlNetModel.from_single_file(url)
85
+
86
+ url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned.safetensors" # can also be a local path
87
+ pipe = StableDiffusionControlNetPipeline.from_single_file(url, controlnet=controlnet)
88
+ ```
89
+ """
90
+ original_config_file = kwargs.pop("original_config_file", None)
91
+ config_file = kwargs.pop("config_file", None)
92
+ resume_download = kwargs.pop("resume_download", False)
93
+ force_download = kwargs.pop("force_download", False)
94
+ proxies = kwargs.pop("proxies", None)
95
+ token = kwargs.pop("token", None)
96
+ cache_dir = kwargs.pop("cache_dir", None)
97
+ local_files_only = kwargs.pop("local_files_only", None)
98
+ revision = kwargs.pop("revision", None)
99
+ torch_dtype = kwargs.pop("torch_dtype", None)
100
+
101
+ class_name = cls.__name__
102
+ if (config_file is not None) and (original_config_file is not None):
103
+ raise ValueError(
104
+ "You cannot pass both `config_file` and `original_config_file` to `from_single_file`. Please use only one of these arguments."
105
+ )
106
+
107
+ original_config_file = config_file or original_config_file
108
+ original_config, checkpoint = fetch_ldm_config_and_checkpoint(
109
+ pretrained_model_link_or_path=pretrained_model_link_or_path,
110
+ class_name=class_name,
111
+ original_config_file=original_config_file,
112
+ resume_download=resume_download,
113
+ force_download=force_download,
114
+ proxies=proxies,
115
+ token=token,
116
+ revision=revision,
117
+ local_files_only=local_files_only,
118
+ cache_dir=cache_dir,
119
+ )
120
+
121
+ upcast_attention = kwargs.pop("upcast_attention", False)
122
+ image_size = kwargs.pop("image_size", None)
123
+
124
+ component = create_diffusers_controlnet_model_from_ldm(
125
+ class_name,
126
+ original_config,
127
+ checkpoint,
128
+ upcast_attention=upcast_attention,
129
+ image_size=image_size,
130
+ torch_dtype=torch_dtype,
131
+ )
132
+ controlnet = component["controlnet"]
133
+ if torch_dtype is not None:
134
+ controlnet = controlnet.to(torch_dtype)
135
+
136
+ return controlnet
model/diffusers_c/loaders/ip_adapter.py ADDED
@@ -0,0 +1,281 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from pathlib import Path
16
+ from typing import Dict, List, Optional, Union
17
+
18
+ import torch
19
+ from huggingface_hub.utils import validate_hf_hub_args
20
+ from safetensors import safe_open
21
+
22
+ from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
23
+ from ..utils import (
24
+ _get_model_file,
25
+ is_accelerate_available,
26
+ is_torch_version,
27
+ is_transformers_available,
28
+ logging,
29
+ )
30
+
31
+
32
+ if is_transformers_available():
33
+ from transformers import (
34
+ CLIPImageProcessor,
35
+ CLIPVisionModelWithProjection,
36
+ )
37
+
38
+ from ..models.attention_processor import (
39
+ IPAdapterAttnProcessor,
40
+ IPAdapterAttnProcessor2_0,
41
+ )
42
+
43
+ logger = logging.get_logger(__name__)
44
+
45
+
46
+ class IPAdapterMixin:
47
+ """Mixin for handling IP Adapters."""
48
+
49
+ @validate_hf_hub_args
50
+ def load_ip_adapter(
51
+ self,
52
+ pretrained_model_name_or_path_or_dict: Union[str, List[str], Dict[str, torch.Tensor]],
53
+ subfolder: Union[str, List[str]],
54
+ weight_name: Union[str, List[str]],
55
+ image_encoder_folder: Optional[str] = "image_encoder",
56
+ **kwargs,
57
+ ):
58
+ """
59
+ Parameters:
60
+ pretrained_model_name_or_path_or_dict (`str` or `List[str]` or `os.PathLike` or `List[os.PathLike]` or `dict` or `List[dict]`):
61
+ Can be either:
62
+
63
+ - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
64
+ the Hub.
65
+ - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
66
+ with [`ModelMixin.save_pretrained`].
67
+ - A [torch state
68
+ dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
69
+ subfolder (`str` or `List[str]`):
70
+ The subfolder location of a model file within a larger model repository on the Hub or locally.
71
+ If a list is passed, it should have the same length as `weight_name`.
72
+ weight_name (`str` or `List[str]`):
73
+ The name of the weight file to load. If a list is passed, it should have the same length as
74
+ `weight_name`.
75
+ image_encoder_folder (`str`, *optional*, defaults to `image_encoder`):
76
+ The subfolder location of the image encoder within a larger model repository on the Hub or locally.
77
+ Pass `None` to not load the image encoder. If the image encoder is located in a folder inside `subfolder`,
78
+ you only need to pass the name of the folder that contains image encoder weights, e.g. `image_encoder_folder="image_encoder"`.
79
+ 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,
80
+ for example, `image_encoder_folder="different_subfolder/image_encoder"`.
81
+ cache_dir (`Union[str, os.PathLike]`, *optional*):
82
+ Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
83
+ is not used.
84
+ force_download (`bool`, *optional*, defaults to `False`):
85
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
86
+ cached versions if they exist.
87
+ resume_download (`bool`, *optional*, defaults to `False`):
88
+ Whether or not to resume downloading the model weights and configuration files. If set to `False`, any
89
+ incompletely downloaded files are deleted.
90
+ proxies (`Dict[str, str]`, *optional*):
91
+ A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
92
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
93
+ local_files_only (`bool`, *optional*, defaults to `False`):
94
+ Whether to only load local model weights and configuration files or not. If set to `True`, the model
95
+ won't be downloaded from the Hub.
96
+ token (`str` or *bool*, *optional*):
97
+ The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
98
+ `diffusers-cli login` (stored in `~/.huggingface`) is used.
99
+ revision (`str`, *optional*, defaults to `"main"`):
100
+ The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
101
+ allowed by Git.
102
+ low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
103
+ Speed up model loading only loading the pretrained weights and not initializing the weights. This also
104
+ tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
105
+ Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
106
+ argument to `True` will raise an error.
107
+ """
108
+
109
+ # handle the list inputs for multiple IP Adapters
110
+ if not isinstance(weight_name, list):
111
+ weight_name = [weight_name]
112
+
113
+ if not isinstance(pretrained_model_name_or_path_or_dict, list):
114
+ pretrained_model_name_or_path_or_dict = [pretrained_model_name_or_path_or_dict]
115
+ if len(pretrained_model_name_or_path_or_dict) == 1:
116
+ pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict * len(weight_name)
117
+
118
+ if not isinstance(subfolder, list):
119
+ subfolder = [subfolder]
120
+ if len(subfolder) == 1:
121
+ subfolder = subfolder * len(weight_name)
122
+
123
+ if len(weight_name) != len(pretrained_model_name_or_path_or_dict):
124
+ raise ValueError("`weight_name` and `pretrained_model_name_or_path_or_dict` must have the same length.")
125
+
126
+ if len(weight_name) != len(subfolder):
127
+ raise ValueError("`weight_name` and `subfolder` must have the same length.")
128
+
129
+ # Load the main state dict first.
130
+ cache_dir = kwargs.pop("cache_dir", None)
131
+ force_download = kwargs.pop("force_download", False)
132
+ resume_download = kwargs.pop("resume_download", False)
133
+ proxies = kwargs.pop("proxies", None)
134
+ local_files_only = kwargs.pop("local_files_only", None)
135
+ token = kwargs.pop("token", None)
136
+ revision = kwargs.pop("revision", None)
137
+ low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
138
+
139
+ if low_cpu_mem_usage and not is_accelerate_available():
140
+ low_cpu_mem_usage = False
141
+ logger.warning(
142
+ "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
143
+ " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
144
+ " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
145
+ " install accelerate\n```\n."
146
+ )
147
+
148
+ if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
149
+ raise NotImplementedError(
150
+ "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"
151
+ " `low_cpu_mem_usage=False`."
152
+ )
153
+
154
+ user_agent = {
155
+ "file_type": "attn_procs_weights",
156
+ "framework": "pytorch",
157
+ }
158
+ state_dicts = []
159
+ for pretrained_model_name_or_path_or_dict, weight_name, subfolder in zip(
160
+ pretrained_model_name_or_path_or_dict, weight_name, subfolder
161
+ ):
162
+ if not isinstance(pretrained_model_name_or_path_or_dict, dict):
163
+ model_file = _get_model_file(
164
+ pretrained_model_name_or_path_or_dict,
165
+ weights_name=weight_name,
166
+ cache_dir=cache_dir,
167
+ force_download=force_download,
168
+ resume_download=resume_download,
169
+ proxies=proxies,
170
+ local_files_only=local_files_only,
171
+ token=token,
172
+ revision=revision,
173
+ subfolder=subfolder,
174
+ user_agent=user_agent,
175
+ )
176
+ if weight_name.endswith(".safetensors"):
177
+ state_dict = {"image_proj": {}, "ip_adapter": {}}
178
+ with safe_open(model_file, framework="pt", device="cpu") as f:
179
+ for key in f.keys():
180
+ if key.startswith("image_proj."):
181
+ state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key)
182
+ elif key.startswith("ip_adapter."):
183
+ state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key)
184
+ else:
185
+ state_dict = torch.load(model_file, map_location="cpu")
186
+ else:
187
+ state_dict = pretrained_model_name_or_path_or_dict
188
+
189
+ keys = list(state_dict.keys())
190
+ if keys != ["image_proj", "ip_adapter"]:
191
+ raise ValueError("Required keys are (`image_proj` and `ip_adapter`) missing from the state dict.")
192
+
193
+ state_dicts.append(state_dict)
194
+
195
+ # load CLIP image encoder here if it has not been registered to the pipeline yet
196
+ if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is None:
197
+ if image_encoder_folder is not None:
198
+ if not isinstance(pretrained_model_name_or_path_or_dict, dict):
199
+ logger.info(f"loading image_encoder from {pretrained_model_name_or_path_or_dict}")
200
+ if image_encoder_folder.count("/") == 0:
201
+ image_encoder_subfolder = Path(subfolder, image_encoder_folder).as_posix()
202
+ else:
203
+ image_encoder_subfolder = Path(image_encoder_folder).as_posix()
204
+
205
+ image_encoder = CLIPVisionModelWithProjection.from_pretrained(
206
+ pretrained_model_name_or_path_or_dict,
207
+ subfolder=image_encoder_subfolder,
208
+ low_cpu_mem_usage=low_cpu_mem_usage,
209
+ ).to(self.device, dtype=self.dtype)
210
+ self.register_modules(image_encoder=image_encoder)
211
+ else:
212
+ raise ValueError(
213
+ "`image_encoder` cannot be loaded because `pretrained_model_name_or_path_or_dict` is a state dict."
214
+ )
215
+ else:
216
+ logger.warning(
217
+ "image_encoder is not loaded since `image_encoder_folder=None` passed. You will not be able to use `ip_adapter_image` when calling the pipeline with IP-Adapter."
218
+ "Use `ip_adapter_image_embeds` to pass pre-generated image embedding instead."
219
+ )
220
+
221
+ # create feature extractor if it has not been registered to the pipeline yet
222
+ if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is None:
223
+ feature_extractor = CLIPImageProcessor()
224
+ self.register_modules(feature_extractor=feature_extractor)
225
+
226
+ # load ip-adapter into unet
227
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
228
+ unet._load_ip_adapter_weights(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage)
229
+
230
+ def set_ip_adapter_scale(self, scale):
231
+ """
232
+ Sets the conditioning scale between text and image.
233
+
234
+ Example:
235
+
236
+ ```py
237
+ pipeline.set_ip_adapter_scale(0.5)
238
+ ```
239
+ """
240
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
241
+ for attn_processor in unet.attn_processors.values():
242
+ if isinstance(attn_processor, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0)):
243
+ if not isinstance(scale, list):
244
+ scale = [scale] * len(attn_processor.scale)
245
+ if len(attn_processor.scale) != len(scale):
246
+ raise ValueError(
247
+ f"`scale` should be a list of same length as the number if ip-adapters "
248
+ f"Expected {len(attn_processor.scale)} but got {len(scale)}."
249
+ )
250
+ attn_processor.scale = scale
251
+
252
+ def unload_ip_adapter(self):
253
+ """
254
+ Unloads the IP Adapter weights
255
+
256
+ Examples:
257
+
258
+ ```python
259
+ >>> # Assuming `pipeline` is already loaded with the IP Adapter weights.
260
+ >>> pipeline.unload_ip_adapter()
261
+ >>> ...
262
+ ```
263
+ """
264
+ # remove CLIP image encoder
265
+ if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is not None:
266
+ self.image_encoder = None
267
+ self.register_to_config(image_encoder=[None, None])
268
+
269
+ # remove feature extractor only when safety_checker is None as safety_checker uses
270
+ # the feature_extractor later
271
+ if not hasattr(self, "safety_checker"):
272
+ if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is not None:
273
+ self.feature_extractor = None
274
+ self.register_to_config(feature_extractor=[None, None])
275
+
276
+ # remove hidden encoder
277
+ self.unet.encoder_hid_proj = None
278
+ self.config.encoder_hid_dim_type = None
279
+
280
+ # restore original Unet attention processors layers
281
+ self.unet.set_default_attn_processor()
model/diffusers_c/loaders/lora.py ADDED
@@ -0,0 +1,1349 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ import inspect
15
+ import os
16
+ from pathlib import Path
17
+ from typing import Callable, Dict, List, Optional, Union
18
+
19
+ import safetensors
20
+ import torch
21
+ from huggingface_hub import model_info
22
+ from huggingface_hub.constants import HF_HUB_OFFLINE
23
+ from huggingface_hub.utils import validate_hf_hub_args
24
+ from packaging import version
25
+ from torch import nn
26
+
27
+ from .. import __version__
28
+ from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
29
+ from ..utils import (
30
+ USE_PEFT_BACKEND,
31
+ _get_model_file,
32
+ convert_state_dict_to_diffusers,
33
+ convert_state_dict_to_peft,
34
+ convert_unet_state_dict_to_peft,
35
+ delete_adapter_layers,
36
+ get_adapter_name,
37
+ get_peft_kwargs,
38
+ is_accelerate_available,
39
+ is_transformers_available,
40
+ logging,
41
+ recurse_remove_peft_layers,
42
+ scale_lora_layers,
43
+ set_adapter_layers,
44
+ set_weights_and_activate_adapters,
45
+ )
46
+ from .lora_conversion_utils import _convert_kohya_lora_to_diffusers, _maybe_map_sgm_blocks_to_diffusers
47
+
48
+
49
+ if is_transformers_available():
50
+ from transformers import PreTrainedModel
51
+
52
+ from ..models.lora import text_encoder_attn_modules, text_encoder_mlp_modules
53
+
54
+ if is_accelerate_available():
55
+ from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module
56
+
57
+ logger = logging.get_logger(__name__)
58
+
59
+ TEXT_ENCODER_NAME = "text_encoder"
60
+ UNET_NAME = "unet"
61
+ TRANSFORMER_NAME = "transformer"
62
+
63
+ LORA_WEIGHT_NAME = "pytorch_lora_weights.bin"
64
+ LORA_WEIGHT_NAME_SAFE = "pytorch_lora_weights.safetensors"
65
+
66
+ LORA_DEPRECATION_MESSAGE = "You are using an old version of LoRA backend. This will be deprecated in the next releases in favor of PEFT make sure to install the latest PEFT and transformers packages in the future."
67
+
68
+
69
+ class LoraLoaderMixin:
70
+ r"""
71
+ Load LoRA layers into [`UNet2DConditionModel`] and
72
+ [`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel).
73
+ """
74
+
75
+ text_encoder_name = TEXT_ENCODER_NAME
76
+ unet_name = UNET_NAME
77
+ transformer_name = TRANSFORMER_NAME
78
+ num_fused_loras = 0
79
+
80
+ def load_lora_weights(
81
+ self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], adapter_name=None, **kwargs
82
+ ):
83
+ """
84
+ Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and
85
+ `self.text_encoder`.
86
+
87
+ All kwargs are forwarded to `self.lora_state_dict`.
88
+
89
+ See [`~loaders.LoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
90
+
91
+ See [`~loaders.LoraLoaderMixin.load_lora_into_unet`] for more details on how the state dict is loaded into
92
+ `self.unet`.
93
+
94
+ See [`~loaders.LoraLoaderMixin.load_lora_into_text_encoder`] for more details on how the state dict is loaded
95
+ into `self.text_encoder`.
96
+
97
+ Parameters:
98
+ pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
99
+ See [`~loaders.LoraLoaderMixin.lora_state_dict`].
100
+ kwargs (`dict`, *optional*):
101
+ See [`~loaders.LoraLoaderMixin.lora_state_dict`].
102
+ adapter_name (`str`, *optional*):
103
+ Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
104
+ `default_{i}` where i is the total number of adapters being loaded.
105
+ """
106
+ if not USE_PEFT_BACKEND:
107
+ raise ValueError("PEFT backend is required for this method.")
108
+
109
+ # if a dict is passed, copy it instead of modifying it inplace
110
+ if isinstance(pretrained_model_name_or_path_or_dict, dict):
111
+ pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
112
+
113
+ # First, ensure that the checkpoint is a compatible one and can be successfully loaded.
114
+ state_dict, network_alphas = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
115
+
116
+ is_correct_format = all("lora" in key for key in state_dict.keys())
117
+ if not is_correct_format:
118
+ raise ValueError("Invalid LoRA checkpoint.")
119
+
120
+ low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
121
+
122
+ self.load_lora_into_unet(
123
+ state_dict,
124
+ network_alphas=network_alphas,
125
+ unet=getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet,
126
+ low_cpu_mem_usage=low_cpu_mem_usage,
127
+ adapter_name=adapter_name,
128
+ _pipeline=self,
129
+ )
130
+ self.load_lora_into_text_encoder(
131
+ state_dict,
132
+ network_alphas=network_alphas,
133
+ text_encoder=getattr(self, self.text_encoder_name)
134
+ if not hasattr(self, "text_encoder")
135
+ else self.text_encoder,
136
+ lora_scale=self.lora_scale,
137
+ low_cpu_mem_usage=low_cpu_mem_usage,
138
+ adapter_name=adapter_name,
139
+ _pipeline=self,
140
+ )
141
+
142
+ @classmethod
143
+ @validate_hf_hub_args
144
+ def lora_state_dict(
145
+ cls,
146
+ pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
147
+ **kwargs,
148
+ ):
149
+ r"""
150
+ Return state dict for lora weights and the network alphas.
151
+
152
+ <Tip warning={true}>
153
+
154
+ We support loading A1111 formatted LoRA checkpoints in a limited capacity.
155
+
156
+ This function is experimental and might change in the future.
157
+
158
+ </Tip>
159
+
160
+ Parameters:
161
+ pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
162
+ Can be either:
163
+
164
+ - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
165
+ the Hub.
166
+ - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
167
+ with [`ModelMixin.save_pretrained`].
168
+ - A [torch state
169
+ dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
170
+
171
+ cache_dir (`Union[str, os.PathLike]`, *optional*):
172
+ Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
173
+ is not used.
174
+ force_download (`bool`, *optional*, defaults to `False`):
175
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
176
+ cached versions if they exist.
177
+ resume_download (`bool`, *optional*, defaults to `False`):
178
+ Whether or not to resume downloading the model weights and configuration files. If set to `False`, any
179
+ incompletely downloaded files are deleted.
180
+ proxies (`Dict[str, str]`, *optional*):
181
+ A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
182
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
183
+ local_files_only (`bool`, *optional*, defaults to `False`):
184
+ Whether to only load local model weights and configuration files or not. If set to `True`, the model
185
+ won't be downloaded from the Hub.
186
+ token (`str` or *bool*, *optional*):
187
+ The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
188
+ `diffusers-cli login` (stored in `~/.huggingface`) is used.
189
+ revision (`str`, *optional*, defaults to `"main"`):
190
+ The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
191
+ allowed by Git.
192
+ subfolder (`str`, *optional*, defaults to `""`):
193
+ The subfolder location of a model file within a larger model repository on the Hub or locally.
194
+ low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
195
+ Speed up model loading only loading the pretrained weights and not initializing the weights. This also
196
+ tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
197
+ Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
198
+ argument to `True` will raise an error.
199
+ mirror (`str`, *optional*):
200
+ Mirror source to resolve accessibility issues if you're downloading a model in China. We do not
201
+ guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
202
+ information.
203
+
204
+ """
205
+ # Load the main state dict first which has the LoRA layers for either of
206
+ # UNet and text encoder or both.
207
+ cache_dir = kwargs.pop("cache_dir", None)
208
+ force_download = kwargs.pop("force_download", False)
209
+ resume_download = kwargs.pop("resume_download", False)
210
+ proxies = kwargs.pop("proxies", None)
211
+ local_files_only = kwargs.pop("local_files_only", None)
212
+ token = kwargs.pop("token", None)
213
+ revision = kwargs.pop("revision", None)
214
+ subfolder = kwargs.pop("subfolder", None)
215
+ weight_name = kwargs.pop("weight_name", None)
216
+ unet_config = kwargs.pop("unet_config", None)
217
+ use_safetensors = kwargs.pop("use_safetensors", None)
218
+
219
+ allow_pickle = False
220
+ if use_safetensors is None:
221
+ use_safetensors = True
222
+ allow_pickle = True
223
+
224
+ user_agent = {
225
+ "file_type": "attn_procs_weights",
226
+ "framework": "pytorch",
227
+ }
228
+
229
+ model_file = None
230
+ if not isinstance(pretrained_model_name_or_path_or_dict, dict):
231
+ # Let's first try to load .safetensors weights
232
+ if (use_safetensors and weight_name is None) or (
233
+ weight_name is not None and weight_name.endswith(".safetensors")
234
+ ):
235
+ try:
236
+ # Here we're relaxing the loading check to enable more Inference API
237
+ # friendliness where sometimes, it's not at all possible to automatically
238
+ # determine `weight_name`.
239
+ if weight_name is None:
240
+ weight_name = cls._best_guess_weight_name(
241
+ pretrained_model_name_or_path_or_dict,
242
+ file_extension=".safetensors",
243
+ local_files_only=local_files_only,
244
+ )
245
+ model_file = _get_model_file(
246
+ pretrained_model_name_or_path_or_dict,
247
+ weights_name=weight_name or LORA_WEIGHT_NAME_SAFE,
248
+ cache_dir=cache_dir,
249
+ force_download=force_download,
250
+ resume_download=resume_download,
251
+ proxies=proxies,
252
+ local_files_only=local_files_only,
253
+ token=token,
254
+ revision=revision,
255
+ subfolder=subfolder,
256
+ user_agent=user_agent,
257
+ )
258
+ state_dict = safetensors.torch.load_file(model_file, device="cpu")
259
+ except (IOError, safetensors.SafetensorError) as e:
260
+ if not allow_pickle:
261
+ raise e
262
+ # try loading non-safetensors weights
263
+ model_file = None
264
+ pass
265
+
266
+ if model_file is None:
267
+ if weight_name is None:
268
+ weight_name = cls._best_guess_weight_name(
269
+ pretrained_model_name_or_path_or_dict, file_extension=".bin", local_files_only=local_files_only
270
+ )
271
+ model_file = _get_model_file(
272
+ pretrained_model_name_or_path_or_dict,
273
+ weights_name=weight_name or LORA_WEIGHT_NAME,
274
+ cache_dir=cache_dir,
275
+ force_download=force_download,
276
+ resume_download=resume_download,
277
+ proxies=proxies,
278
+ local_files_only=local_files_only,
279
+ token=token,
280
+ revision=revision,
281
+ subfolder=subfolder,
282
+ user_agent=user_agent,
283
+ )
284
+ state_dict = torch.load(model_file, map_location="cpu")
285
+ else:
286
+ state_dict = pretrained_model_name_or_path_or_dict
287
+
288
+ network_alphas = None
289
+ # TODO: replace it with a method from `state_dict_utils`
290
+ if all(
291
+ (
292
+ k.startswith("lora_te_")
293
+ or k.startswith("lora_unet_")
294
+ or k.startswith("lora_te1_")
295
+ or k.startswith("lora_te2_")
296
+ )
297
+ for k in state_dict.keys()
298
+ ):
299
+ # Map SDXL blocks correctly.
300
+ if unet_config is not None:
301
+ # use unet config to remap block numbers
302
+ state_dict = _maybe_map_sgm_blocks_to_diffusers(state_dict, unet_config)
303
+ state_dict, network_alphas = _convert_kohya_lora_to_diffusers(state_dict)
304
+
305
+ return state_dict, network_alphas
306
+
307
+ @classmethod
308
+ def _best_guess_weight_name(
309
+ cls, pretrained_model_name_or_path_or_dict, file_extension=".safetensors", local_files_only=False
310
+ ):
311
+ if local_files_only or HF_HUB_OFFLINE:
312
+ raise ValueError("When using the offline mode, you must specify a `weight_name`.")
313
+
314
+ targeted_files = []
315
+
316
+ if os.path.isfile(pretrained_model_name_or_path_or_dict):
317
+ return
318
+ elif os.path.isdir(pretrained_model_name_or_path_or_dict):
319
+ targeted_files = [
320
+ f for f in os.listdir(pretrained_model_name_or_path_or_dict) if f.endswith(file_extension)
321
+ ]
322
+ else:
323
+ files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings
324
+ targeted_files = [f.rfilename for f in files_in_repo if f.rfilename.endswith(file_extension)]
325
+ if len(targeted_files) == 0:
326
+ return
327
+
328
+ # "scheduler" does not correspond to a LoRA checkpoint.
329
+ # "optimizer" does not correspond to a LoRA checkpoint
330
+ # only top-level checkpoints are considered and not the other ones, hence "checkpoint".
331
+ unallowed_substrings = {"scheduler", "optimizer", "checkpoint"}
332
+ targeted_files = list(
333
+ filter(lambda x: all(substring not in x for substring in unallowed_substrings), targeted_files)
334
+ )
335
+
336
+ if any(f.endswith(LORA_WEIGHT_NAME) for f in targeted_files):
337
+ targeted_files = list(filter(lambda x: x.endswith(LORA_WEIGHT_NAME), targeted_files))
338
+ elif any(f.endswith(LORA_WEIGHT_NAME_SAFE) for f in targeted_files):
339
+ targeted_files = list(filter(lambda x: x.endswith(LORA_WEIGHT_NAME_SAFE), targeted_files))
340
+
341
+ if len(targeted_files) > 1:
342
+ raise ValueError(
343
+ f"Provided path contains more than one weights file in the {file_extension} format. Either specify `weight_name` in `load_lora_weights` or make sure there's only one `.safetensors` or `.bin` file in {pretrained_model_name_or_path_or_dict}."
344
+ )
345
+ weight_name = targeted_files[0]
346
+ return weight_name
347
+
348
+ @classmethod
349
+ def _optionally_disable_offloading(cls, _pipeline):
350
+ """
351
+ Optionally removes offloading in case the pipeline has been already sequentially offloaded to CPU.
352
+
353
+ Args:
354
+ _pipeline (`DiffusionPipeline`):
355
+ The pipeline to disable offloading for.
356
+
357
+ Returns:
358
+ tuple:
359
+ A tuple indicating if `is_model_cpu_offload` or `is_sequential_cpu_offload` is True.
360
+ """
361
+ is_model_cpu_offload = False
362
+ is_sequential_cpu_offload = False
363
+
364
+ if _pipeline is not None:
365
+ for _, component in _pipeline.components.items():
366
+ if isinstance(component, nn.Module) and hasattr(component, "_hf_hook"):
367
+ if not is_model_cpu_offload:
368
+ is_model_cpu_offload = isinstance(component._hf_hook, CpuOffload)
369
+ if not is_sequential_cpu_offload:
370
+ is_sequential_cpu_offload = isinstance(component._hf_hook, AlignDevicesHook)
371
+
372
+ logger.info(
373
+ "Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again."
374
+ )
375
+ remove_hook_from_module(component, recurse=is_sequential_cpu_offload)
376
+
377
+ return (is_model_cpu_offload, is_sequential_cpu_offload)
378
+
379
+ @classmethod
380
+ def load_lora_into_unet(
381
+ cls, state_dict, network_alphas, unet, low_cpu_mem_usage=None, adapter_name=None, _pipeline=None
382
+ ):
383
+ """
384
+ This will load the LoRA layers specified in `state_dict` into `unet`.
385
+
386
+ Parameters:
387
+ state_dict (`dict`):
388
+ A standard state dict containing the lora layer parameters. The keys can either be indexed directly
389
+ into the unet or prefixed with an additional `unet` which can be used to distinguish between text
390
+ encoder lora layers.
391
+ network_alphas (`Dict[str, float]`):
392
+ See `LoRALinearLayer` for more details.
393
+ unet (`UNet2DConditionModel`):
394
+ The UNet model to load the LoRA layers into.
395
+ low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
396
+ Speed up model loading only loading the pretrained weights and not initializing the weights. This also
397
+ tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
398
+ Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
399
+ argument to `True` will raise an error.
400
+ adapter_name (`str`, *optional*):
401
+ Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
402
+ `default_{i}` where i is the total number of adapters being loaded.
403
+ """
404
+ if not USE_PEFT_BACKEND:
405
+ raise ValueError("PEFT backend is required for this method.")
406
+
407
+ from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
408
+
409
+ low_cpu_mem_usage = low_cpu_mem_usage if low_cpu_mem_usage is not None else _LOW_CPU_MEM_USAGE_DEFAULT
410
+ # If the serialization format is new (introduced in https://github.com/huggingface/diffusers/pull/2918),
411
+ # then the `state_dict` keys should have `cls.unet_name` and/or `cls.text_encoder_name` as
412
+ # their prefixes.
413
+ keys = list(state_dict.keys())
414
+
415
+ if all(key.startswith(cls.unet_name) or key.startswith(cls.text_encoder_name) for key in keys):
416
+ # Load the layers corresponding to UNet.
417
+ logger.info(f"Loading {cls.unet_name}.")
418
+
419
+ unet_keys = [k for k in keys if k.startswith(cls.unet_name)]
420
+ state_dict = {k.replace(f"{cls.unet_name}.", ""): v for k, v in state_dict.items() if k in unet_keys}
421
+
422
+ if network_alphas is not None:
423
+ alpha_keys = [k for k in network_alphas.keys() if k.startswith(cls.unet_name)]
424
+ network_alphas = {
425
+ k.replace(f"{cls.unet_name}.", ""): v for k, v in network_alphas.items() if k in alpha_keys
426
+ }
427
+
428
+ else:
429
+ # Otherwise, we're dealing with the old format. This means the `state_dict` should only
430
+ # contain the module names of the `unet` as its keys WITHOUT any prefix.
431
+ if not USE_PEFT_BACKEND:
432
+ warn_message = "You have saved the LoRA weights using the old format. To convert the old LoRA weights to the new format, you can first load them in a dictionary and then create a new dictionary like the following: `new_state_dict = {f'unet.{module_name}': params for module_name, params in old_state_dict.items()}`."
433
+ logger.warn(warn_message)
434
+
435
+ if len(state_dict.keys()) > 0:
436
+ if adapter_name in getattr(unet, "peft_config", {}):
437
+ raise ValueError(
438
+ f"Adapter name {adapter_name} already in use in the Unet - please select a new adapter name."
439
+ )
440
+
441
+ state_dict = convert_unet_state_dict_to_peft(state_dict)
442
+
443
+ if network_alphas is not None:
444
+ # The alphas state dict have the same structure as Unet, thus we convert it to peft format using
445
+ # `convert_unet_state_dict_to_peft` method.
446
+ network_alphas = convert_unet_state_dict_to_peft(network_alphas)
447
+
448
+ rank = {}
449
+ for key, val in state_dict.items():
450
+ if "lora_B" in key:
451
+ rank[key] = val.shape[1]
452
+
453
+ lora_config_kwargs = get_peft_kwargs(rank, network_alphas, state_dict, is_unet=True)
454
+ lora_config = LoraConfig(**lora_config_kwargs)
455
+
456
+ # adapter_name
457
+ if adapter_name is None:
458
+ adapter_name = get_adapter_name(unet)
459
+
460
+ # In case the pipeline has been already offloaded to CPU - temporarily remove the hooks
461
+ # otherwise loading LoRA weights will lead to an error
462
+ is_model_cpu_offload, is_sequential_cpu_offload = cls._optionally_disable_offloading(_pipeline)
463
+
464
+ inject_adapter_in_model(lora_config, unet, adapter_name=adapter_name)
465
+ incompatible_keys = set_peft_model_state_dict(unet, state_dict, adapter_name)
466
+
467
+ if incompatible_keys is not None:
468
+ # check only for unexpected keys
469
+ unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
470
+ if unexpected_keys:
471
+ logger.warning(
472
+ f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
473
+ f" {unexpected_keys}. "
474
+ )
475
+
476
+ # Offload back.
477
+ if is_model_cpu_offload:
478
+ _pipeline.enable_model_cpu_offload()
479
+ elif is_sequential_cpu_offload:
480
+ _pipeline.enable_sequential_cpu_offload()
481
+ # Unsafe code />
482
+
483
+ unet.load_attn_procs(
484
+ state_dict, network_alphas=network_alphas, low_cpu_mem_usage=low_cpu_mem_usage, _pipeline=_pipeline
485
+ )
486
+
487
+ @classmethod
488
+ def load_lora_into_text_encoder(
489
+ cls,
490
+ state_dict,
491
+ network_alphas,
492
+ text_encoder,
493
+ prefix=None,
494
+ lora_scale=1.0,
495
+ low_cpu_mem_usage=None,
496
+ adapter_name=None,
497
+ _pipeline=None,
498
+ ):
499
+ """
500
+ This will load the LoRA layers specified in `state_dict` into `text_encoder`
501
+
502
+ Parameters:
503
+ state_dict (`dict`):
504
+ A standard state dict containing the lora layer parameters. The key should be prefixed with an
505
+ additional `text_encoder` to distinguish between unet lora layers.
506
+ network_alphas (`Dict[str, float]`):
507
+ See `LoRALinearLayer` for more details.
508
+ text_encoder (`CLIPTextModel`):
509
+ The text encoder model to load the LoRA layers into.
510
+ prefix (`str`):
511
+ Expected prefix of the `text_encoder` in the `state_dict`.
512
+ lora_scale (`float`):
513
+ How much to scale the output of the lora linear layer before it is added with the output of the regular
514
+ lora layer.
515
+ low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
516
+ Speed up model loading only loading the pretrained weights and not initializing the weights. This also
517
+ tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
518
+ Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
519
+ argument to `True` will raise an error.
520
+ adapter_name (`str`, *optional*):
521
+ Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
522
+ `default_{i}` where i is the total number of adapters being loaded.
523
+ """
524
+ if not USE_PEFT_BACKEND:
525
+ raise ValueError("PEFT backend is required for this method.")
526
+
527
+ from peft import LoraConfig
528
+
529
+ low_cpu_mem_usage = low_cpu_mem_usage if low_cpu_mem_usage is not None else _LOW_CPU_MEM_USAGE_DEFAULT
530
+
531
+ # If the serialization format is new (introduced in https://github.com/huggingface/diffusers/pull/2918),
532
+ # then the `state_dict` keys should have `self.unet_name` and/or `self.text_encoder_name` as
533
+ # their prefixes.
534
+ keys = list(state_dict.keys())
535
+ prefix = cls.text_encoder_name if prefix is None else prefix
536
+
537
+ # Safe prefix to check with.
538
+ if any(cls.text_encoder_name in key for key in keys):
539
+ # Load the layers corresponding to text encoder and make necessary adjustments.
540
+ text_encoder_keys = [k for k in keys if k.startswith(prefix) and k.split(".")[0] == prefix]
541
+ text_encoder_lora_state_dict = {
542
+ k.replace(f"{prefix}.", ""): v for k, v in state_dict.items() if k in text_encoder_keys
543
+ }
544
+
545
+ if len(text_encoder_lora_state_dict) > 0:
546
+ logger.info(f"Loading {prefix}.")
547
+ rank = {}
548
+ text_encoder_lora_state_dict = convert_state_dict_to_diffusers(text_encoder_lora_state_dict)
549
+
550
+ # convert state dict
551
+ text_encoder_lora_state_dict = convert_state_dict_to_peft(text_encoder_lora_state_dict)
552
+
553
+ for name, _ in text_encoder_attn_modules(text_encoder):
554
+ rank_key = f"{name}.out_proj.lora_B.weight"
555
+ rank[rank_key] = text_encoder_lora_state_dict[rank_key].shape[1]
556
+
557
+ patch_mlp = any(".mlp." in key for key in text_encoder_lora_state_dict.keys())
558
+ if patch_mlp:
559
+ for name, _ in text_encoder_mlp_modules(text_encoder):
560
+ rank_key_fc1 = f"{name}.fc1.lora_B.weight"
561
+ rank_key_fc2 = f"{name}.fc2.lora_B.weight"
562
+
563
+ rank[rank_key_fc1] = text_encoder_lora_state_dict[rank_key_fc1].shape[1]
564
+ rank[rank_key_fc2] = text_encoder_lora_state_dict[rank_key_fc2].shape[1]
565
+
566
+ if network_alphas is not None:
567
+ alpha_keys = [
568
+ k for k in network_alphas.keys() if k.startswith(prefix) and k.split(".")[0] == prefix
569
+ ]
570
+ network_alphas = {
571
+ k.replace(f"{prefix}.", ""): v for k, v in network_alphas.items() if k in alpha_keys
572
+ }
573
+
574
+ lora_config_kwargs = get_peft_kwargs(rank, network_alphas, text_encoder_lora_state_dict, is_unet=False)
575
+ lora_config = LoraConfig(**lora_config_kwargs)
576
+
577
+ # adapter_name
578
+ if adapter_name is None:
579
+ adapter_name = get_adapter_name(text_encoder)
580
+
581
+ is_model_cpu_offload, is_sequential_cpu_offload = cls._optionally_disable_offloading(_pipeline)
582
+
583
+ # inject LoRA layers and load the state dict
584
+ # in transformers we automatically check whether the adapter name is already in use or not
585
+ text_encoder.load_adapter(
586
+ adapter_name=adapter_name,
587
+ adapter_state_dict=text_encoder_lora_state_dict,
588
+ peft_config=lora_config,
589
+ )
590
+
591
+ # scale LoRA layers with `lora_scale`
592
+ scale_lora_layers(text_encoder, weight=lora_scale)
593
+
594
+ text_encoder.to(device=text_encoder.device, dtype=text_encoder.dtype)
595
+
596
+ # Offload back.
597
+ if is_model_cpu_offload:
598
+ _pipeline.enable_model_cpu_offload()
599
+ elif is_sequential_cpu_offload:
600
+ _pipeline.enable_sequential_cpu_offload()
601
+ # Unsafe code />
602
+
603
+ @classmethod
604
+ def load_lora_into_transformer(
605
+ cls, state_dict, network_alphas, transformer, low_cpu_mem_usage=None, adapter_name=None, _pipeline=None
606
+ ):
607
+ """
608
+ This will load the LoRA layers specified in `state_dict` into `transformer`.
609
+
610
+ Parameters:
611
+ state_dict (`dict`):
612
+ A standard state dict containing the lora layer parameters. The keys can either be indexed directly
613
+ into the unet or prefixed with an additional `unet` which can be used to distinguish between text
614
+ encoder lora layers.
615
+ network_alphas (`Dict[str, float]`):
616
+ See `LoRALinearLayer` for more details.
617
+ unet (`UNet2DConditionModel`):
618
+ The UNet model to load the LoRA layers into.
619
+ low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
620
+ Speed up model loading only loading the pretrained weights and not initializing the weights. This also
621
+ tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
622
+ Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
623
+ argument to `True` will raise an error.
624
+ adapter_name (`str`, *optional*):
625
+ Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
626
+ `default_{i}` where i is the total number of adapters being loaded.
627
+ """
628
+ from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
629
+
630
+ low_cpu_mem_usage = low_cpu_mem_usage if low_cpu_mem_usage is not None else _LOW_CPU_MEM_USAGE_DEFAULT
631
+
632
+ keys = list(state_dict.keys())
633
+
634
+ transformer_keys = [k for k in keys if k.startswith(cls.transformer_name)]
635
+ state_dict = {
636
+ k.replace(f"{cls.transformer_name}.", ""): v for k, v in state_dict.items() if k in transformer_keys
637
+ }
638
+
639
+ if network_alphas is not None:
640
+ alpha_keys = [k for k in network_alphas.keys() if k.startswith(cls.transformer_name)]
641
+ network_alphas = {
642
+ k.replace(f"{cls.transformer_name}.", ""): v for k, v in network_alphas.items() if k in alpha_keys
643
+ }
644
+
645
+ if len(state_dict.keys()) > 0:
646
+ if adapter_name in getattr(transformer, "peft_config", {}):
647
+ raise ValueError(
648
+ f"Adapter name {adapter_name} already in use in the transformer - please select a new adapter name."
649
+ )
650
+
651
+ rank = {}
652
+ for key, val in state_dict.items():
653
+ if "lora_B" in key:
654
+ rank[key] = val.shape[1]
655
+
656
+ lora_config_kwargs = get_peft_kwargs(rank, network_alphas, state_dict)
657
+ lora_config = LoraConfig(**lora_config_kwargs)
658
+
659
+ # adapter_name
660
+ if adapter_name is None:
661
+ adapter_name = get_adapter_name(transformer)
662
+
663
+ # In case the pipeline has been already offloaded to CPU - temporarily remove the hooks
664
+ # otherwise loading LoRA weights will lead to an error
665
+ is_model_cpu_offload, is_sequential_cpu_offload = cls._optionally_disable_offloading(_pipeline)
666
+
667
+ inject_adapter_in_model(lora_config, transformer, adapter_name=adapter_name)
668
+ incompatible_keys = set_peft_model_state_dict(transformer, state_dict, adapter_name)
669
+
670
+ if incompatible_keys is not None:
671
+ # check only for unexpected keys
672
+ unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
673
+ if unexpected_keys:
674
+ logger.warning(
675
+ f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
676
+ f" {unexpected_keys}. "
677
+ )
678
+
679
+ # Offload back.
680
+ if is_model_cpu_offload:
681
+ _pipeline.enable_model_cpu_offload()
682
+ elif is_sequential_cpu_offload:
683
+ _pipeline.enable_sequential_cpu_offload()
684
+ # Unsafe code />
685
+
686
+ @property
687
+ def lora_scale(self) -> float:
688
+ # property function that returns the lora scale which can be set at run time by the pipeline.
689
+ # if _lora_scale has not been set, return 1
690
+ return self._lora_scale if hasattr(self, "_lora_scale") else 1.0
691
+
692
+ def _remove_text_encoder_monkey_patch(self):
693
+ remove_method = recurse_remove_peft_layers
694
+ if hasattr(self, "text_encoder"):
695
+ remove_method(self.text_encoder)
696
+ # In case text encoder have no Lora attached
697
+ if getattr(self.text_encoder, "peft_config", None) is not None:
698
+ del self.text_encoder.peft_config
699
+ self.text_encoder._hf_peft_config_loaded = None
700
+
701
+ if hasattr(self, "text_encoder_2"):
702
+ remove_method(self.text_encoder_2)
703
+ if getattr(self.text_encoder_2, "peft_config", None) is not None:
704
+ del self.text_encoder_2.peft_config
705
+ self.text_encoder_2._hf_peft_config_loaded = None
706
+
707
+ @classmethod
708
+ def save_lora_weights(
709
+ cls,
710
+ save_directory: Union[str, os.PathLike],
711
+ unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
712
+ text_encoder_lora_layers: Dict[str, torch.nn.Module] = None,
713
+ transformer_lora_layers: Dict[str, torch.nn.Module] = None,
714
+ is_main_process: bool = True,
715
+ weight_name: str = None,
716
+ save_function: Callable = None,
717
+ safe_serialization: bool = True,
718
+ ):
719
+ r"""
720
+ Save the LoRA parameters corresponding to the UNet and text encoder.
721
+
722
+ Arguments:
723
+ save_directory (`str` or `os.PathLike`):
724
+ Directory to save LoRA parameters to. Will be created if it doesn't exist.
725
+ unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
726
+ State dict of the LoRA layers corresponding to the `unet`.
727
+ text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
728
+ State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
729
+ encoder LoRA state dict because it comes from 🤗 Transformers.
730
+ is_main_process (`bool`, *optional*, defaults to `True`):
731
+ Whether the process calling this is the main process or not. Useful during distributed training and you
732
+ need to call this function on all processes. In this case, set `is_main_process=True` only on the main
733
+ process to avoid race conditions.
734
+ save_function (`Callable`):
735
+ The function to use to save the state dictionary. Useful during distributed training when you need to
736
+ replace `torch.save` with another method. Can be configured with the environment variable
737
+ `DIFFUSERS_SAVE_MODE`.
738
+ safe_serialization (`bool`, *optional*, defaults to `True`):
739
+ Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
740
+ """
741
+ state_dict = {}
742
+
743
+ def pack_weights(layers, prefix):
744
+ layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers
745
+ layers_state_dict = {f"{prefix}.{module_name}": param for module_name, param in layers_weights.items()}
746
+ return layers_state_dict
747
+
748
+ if not (unet_lora_layers or text_encoder_lora_layers or transformer_lora_layers):
749
+ raise ValueError(
750
+ "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers`, or `transformer_lora_layers`."
751
+ )
752
+
753
+ if unet_lora_layers:
754
+ state_dict.update(pack_weights(unet_lora_layers, cls.unet_name))
755
+
756
+ if text_encoder_lora_layers:
757
+ state_dict.update(pack_weights(text_encoder_lora_layers, cls.text_encoder_name))
758
+
759
+ if transformer_lora_layers:
760
+ state_dict.update(pack_weights(transformer_lora_layers, "transformer"))
761
+
762
+ # Save the model
763
+ cls.write_lora_layers(
764
+ state_dict=state_dict,
765
+ save_directory=save_directory,
766
+ is_main_process=is_main_process,
767
+ weight_name=weight_name,
768
+ save_function=save_function,
769
+ safe_serialization=safe_serialization,
770
+ )
771
+
772
+ @staticmethod
773
+ def write_lora_layers(
774
+ state_dict: Dict[str, torch.Tensor],
775
+ save_directory: str,
776
+ is_main_process: bool,
777
+ weight_name: str,
778
+ save_function: Callable,
779
+ safe_serialization: bool,
780
+ ):
781
+ if os.path.isfile(save_directory):
782
+ logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
783
+ return
784
+
785
+ if save_function is None:
786
+ if safe_serialization:
787
+
788
+ def save_function(weights, filename):
789
+ return safetensors.torch.save_file(weights, filename, metadata={"format": "pt"})
790
+
791
+ else:
792
+ save_function = torch.save
793
+
794
+ os.makedirs(save_directory, exist_ok=True)
795
+
796
+ if weight_name is None:
797
+ if safe_serialization:
798
+ weight_name = LORA_WEIGHT_NAME_SAFE
799
+ else:
800
+ weight_name = LORA_WEIGHT_NAME
801
+
802
+ save_path = Path(save_directory, weight_name).as_posix()
803
+ save_function(state_dict, save_path)
804
+ logger.info(f"Model weights saved in {save_path}")
805
+
806
+ def unload_lora_weights(self):
807
+ """
808
+ Unloads the LoRA parameters.
809
+
810
+ Examples:
811
+
812
+ ```python
813
+ >>> # Assuming `pipeline` is already loaded with the LoRA parameters.
814
+ >>> pipeline.unload_lora_weights()
815
+ >>> ...
816
+ ```
817
+ """
818
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
819
+
820
+ if not USE_PEFT_BACKEND:
821
+ if version.parse(__version__) > version.parse("0.23"):
822
+ logger.warning(
823
+ "You are using `unload_lora_weights` to disable and unload lora weights. If you want to iteratively enable and disable adapter weights,"
824
+ "you can use `pipe.enable_lora()` or `pipe.disable_lora()`. After installing the latest version of PEFT."
825
+ )
826
+
827
+ for _, module in unet.named_modules():
828
+ if hasattr(module, "set_lora_layer"):
829
+ module.set_lora_layer(None)
830
+ else:
831
+ recurse_remove_peft_layers(unet)
832
+ if hasattr(unet, "peft_config"):
833
+ del unet.peft_config
834
+
835
+ # Safe to call the following regardless of LoRA.
836
+ self._remove_text_encoder_monkey_patch()
837
+
838
+ def fuse_lora(
839
+ self,
840
+ fuse_unet: bool = True,
841
+ fuse_text_encoder: bool = True,
842
+ lora_scale: float = 1.0,
843
+ safe_fusing: bool = False,
844
+ adapter_names: Optional[List[str]] = None,
845
+ ):
846
+ r"""
847
+ Fuses the LoRA parameters into the original parameters of the corresponding blocks.
848
+
849
+ <Tip warning={true}>
850
+
851
+ This is an experimental API.
852
+
853
+ </Tip>
854
+
855
+ Args:
856
+ fuse_unet (`bool`, defaults to `True`): Whether to fuse the UNet LoRA parameters.
857
+ fuse_text_encoder (`bool`, defaults to `True`):
858
+ Whether to fuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
859
+ LoRA parameters then it won't have any effect.
860
+ lora_scale (`float`, defaults to 1.0):
861
+ Controls how much to influence the outputs with the LoRA parameters.
862
+ safe_fusing (`bool`, defaults to `False`):
863
+ Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
864
+ adapter_names (`List[str]`, *optional*):
865
+ Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
866
+
867
+ Example:
868
+
869
+ ```py
870
+ from diffusers import DiffusionPipeline
871
+ import torch
872
+
873
+ pipeline = DiffusionPipeline.from_pretrained(
874
+ "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
875
+ ).to("cuda")
876
+ pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
877
+ pipeline.fuse_lora(lora_scale=0.7)
878
+ ```
879
+ """
880
+ from peft.tuners.tuners_utils import BaseTunerLayer
881
+
882
+ if fuse_unet or fuse_text_encoder:
883
+ self.num_fused_loras += 1
884
+ if self.num_fused_loras > 1:
885
+ logger.warn(
886
+ "The current API is supported for operating with a single LoRA file. You are trying to load and fuse more than one LoRA which is not well-supported.",
887
+ )
888
+
889
+ if fuse_unet:
890
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
891
+ unet.fuse_lora(lora_scale, safe_fusing=safe_fusing, adapter_names=adapter_names)
892
+
893
+ def fuse_text_encoder_lora(text_encoder, lora_scale=1.0, safe_fusing=False, adapter_names=None):
894
+ merge_kwargs = {"safe_merge": safe_fusing}
895
+
896
+ for module in text_encoder.modules():
897
+ if isinstance(module, BaseTunerLayer):
898
+ if lora_scale != 1.0:
899
+ module.scale_layer(lora_scale)
900
+
901
+ # For BC with previous PEFT versions, we need to check the signature
902
+ # of the `merge` method to see if it supports the `adapter_names` argument.
903
+ supported_merge_kwargs = list(inspect.signature(module.merge).parameters)
904
+ if "adapter_names" in supported_merge_kwargs:
905
+ merge_kwargs["adapter_names"] = adapter_names
906
+ elif "adapter_names" not in supported_merge_kwargs and adapter_names is not None:
907
+ raise ValueError(
908
+ "The `adapter_names` argument is not supported with your PEFT version. "
909
+ "Please upgrade to the latest version of PEFT. `pip install -U peft`"
910
+ )
911
+
912
+ module.merge(**merge_kwargs)
913
+
914
+ if fuse_text_encoder:
915
+ if hasattr(self, "text_encoder"):
916
+ fuse_text_encoder_lora(self.text_encoder, lora_scale, safe_fusing, adapter_names=adapter_names)
917
+ if hasattr(self, "text_encoder_2"):
918
+ fuse_text_encoder_lora(self.text_encoder_2, lora_scale, safe_fusing, adapter_names=adapter_names)
919
+
920
+ def unfuse_lora(self, unfuse_unet: bool = True, unfuse_text_encoder: bool = True):
921
+ r"""
922
+ Reverses the effect of
923
+ [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraLoaderMixin.fuse_lora).
924
+
925
+ <Tip warning={true}>
926
+
927
+ This is an experimental API.
928
+
929
+ </Tip>
930
+
931
+ Args:
932
+ unfuse_unet (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
933
+ unfuse_text_encoder (`bool`, defaults to `True`):
934
+ Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
935
+ LoRA parameters then it won't have any effect.
936
+ """
937
+ from peft.tuners.tuners_utils import BaseTunerLayer
938
+
939
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
940
+ if unfuse_unet:
941
+ for module in unet.modules():
942
+ if isinstance(module, BaseTunerLayer):
943
+ module.unmerge()
944
+
945
+ def unfuse_text_encoder_lora(text_encoder):
946
+ for module in text_encoder.modules():
947
+ if isinstance(module, BaseTunerLayer):
948
+ module.unmerge()
949
+
950
+ if unfuse_text_encoder:
951
+ if hasattr(self, "text_encoder"):
952
+ unfuse_text_encoder_lora(self.text_encoder)
953
+ if hasattr(self, "text_encoder_2"):
954
+ unfuse_text_encoder_lora(self.text_encoder_2)
955
+
956
+ self.num_fused_loras -= 1
957
+
958
+ def set_adapters_for_text_encoder(
959
+ self,
960
+ adapter_names: Union[List[str], str],
961
+ text_encoder: Optional["PreTrainedModel"] = None, # noqa: F821
962
+ text_encoder_weights: List[float] = None,
963
+ ):
964
+ """
965
+ Sets the adapter layers for the text encoder.
966
+
967
+ Args:
968
+ adapter_names (`List[str]` or `str`):
969
+ The names of the adapters to use.
970
+ text_encoder (`torch.nn.Module`, *optional*):
971
+ The text encoder module to set the adapter layers for. If `None`, it will try to get the `text_encoder`
972
+ attribute.
973
+ text_encoder_weights (`List[float]`, *optional*):
974
+ The weights to use for the text encoder. If `None`, the weights are set to `1.0` for all the adapters.
975
+ """
976
+ if not USE_PEFT_BACKEND:
977
+ raise ValueError("PEFT backend is required for this method.")
978
+
979
+ def process_weights(adapter_names, weights):
980
+ if weights is None:
981
+ weights = [1.0] * len(adapter_names)
982
+ elif isinstance(weights, float):
983
+ weights = [weights]
984
+
985
+ if len(adapter_names) != len(weights):
986
+ raise ValueError(
987
+ f"Length of adapter names {len(adapter_names)} is not equal to the length of the weights {len(weights)}"
988
+ )
989
+ return weights
990
+
991
+ adapter_names = [adapter_names] if isinstance(adapter_names, str) else adapter_names
992
+ text_encoder_weights = process_weights(adapter_names, text_encoder_weights)
993
+ text_encoder = text_encoder or getattr(self, "text_encoder", None)
994
+ if text_encoder is None:
995
+ raise ValueError(
996
+ "The pipeline does not have a default `pipe.text_encoder` class. Please make sure to pass a `text_encoder` instead."
997
+ )
998
+ set_weights_and_activate_adapters(text_encoder, adapter_names, text_encoder_weights)
999
+
1000
+ def disable_lora_for_text_encoder(self, text_encoder: Optional["PreTrainedModel"] = None):
1001
+ """
1002
+ Disables the LoRA layers for the text encoder.
1003
+
1004
+ Args:
1005
+ text_encoder (`torch.nn.Module`, *optional*):
1006
+ The text encoder module to disable the LoRA layers for. If `None`, it will try to get the
1007
+ `text_encoder` attribute.
1008
+ """
1009
+ if not USE_PEFT_BACKEND:
1010
+ raise ValueError("PEFT backend is required for this method.")
1011
+
1012
+ text_encoder = text_encoder or getattr(self, "text_encoder", None)
1013
+ if text_encoder is None:
1014
+ raise ValueError("Text Encoder not found.")
1015
+ set_adapter_layers(text_encoder, enabled=False)
1016
+
1017
+ def enable_lora_for_text_encoder(self, text_encoder: Optional["PreTrainedModel"] = None):
1018
+ """
1019
+ Enables the LoRA layers for the text encoder.
1020
+
1021
+ Args:
1022
+ text_encoder (`torch.nn.Module`, *optional*):
1023
+ The text encoder module to enable the LoRA layers for. If `None`, it will try to get the `text_encoder`
1024
+ attribute.
1025
+ """
1026
+ if not USE_PEFT_BACKEND:
1027
+ raise ValueError("PEFT backend is required for this method.")
1028
+ text_encoder = text_encoder or getattr(self, "text_encoder", None)
1029
+ if text_encoder is None:
1030
+ raise ValueError("Text Encoder not found.")
1031
+ set_adapter_layers(self.text_encoder, enabled=True)
1032
+
1033
+ def set_adapters(
1034
+ self,
1035
+ adapter_names: Union[List[str], str],
1036
+ adapter_weights: Optional[List[float]] = None,
1037
+ ):
1038
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
1039
+ # Handle the UNET
1040
+ unet.set_adapters(adapter_names, adapter_weights)
1041
+
1042
+ # Handle the Text Encoder
1043
+ if hasattr(self, "text_encoder"):
1044
+ self.set_adapters_for_text_encoder(adapter_names, self.text_encoder, adapter_weights)
1045
+ if hasattr(self, "text_encoder_2"):
1046
+ self.set_adapters_for_text_encoder(adapter_names, self.text_encoder_2, adapter_weights)
1047
+
1048
+ def disable_lora(self):
1049
+ if not USE_PEFT_BACKEND:
1050
+ raise ValueError("PEFT backend is required for this method.")
1051
+
1052
+ # Disable unet adapters
1053
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
1054
+ unet.disable_lora()
1055
+
1056
+ # Disable text encoder adapters
1057
+ if hasattr(self, "text_encoder"):
1058
+ self.disable_lora_for_text_encoder(self.text_encoder)
1059
+ if hasattr(self, "text_encoder_2"):
1060
+ self.disable_lora_for_text_encoder(self.text_encoder_2)
1061
+
1062
+ def enable_lora(self):
1063
+ if not USE_PEFT_BACKEND:
1064
+ raise ValueError("PEFT backend is required for this method.")
1065
+
1066
+ # Enable unet adapters
1067
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
1068
+ unet.enable_lora()
1069
+
1070
+ # Enable text encoder adapters
1071
+ if hasattr(self, "text_encoder"):
1072
+ self.enable_lora_for_text_encoder(self.text_encoder)
1073
+ if hasattr(self, "text_encoder_2"):
1074
+ self.enable_lora_for_text_encoder(self.text_encoder_2)
1075
+
1076
+ def delete_adapters(self, adapter_names: Union[List[str], str]):
1077
+ """
1078
+ Args:
1079
+ Deletes the LoRA layers of `adapter_name` for the unet and text-encoder(s).
1080
+ adapter_names (`Union[List[str], str]`):
1081
+ The names of the adapter to delete. Can be a single string or a list of strings
1082
+ """
1083
+ if not USE_PEFT_BACKEND:
1084
+ raise ValueError("PEFT backend is required for this method.")
1085
+
1086
+ if isinstance(adapter_names, str):
1087
+ adapter_names = [adapter_names]
1088
+
1089
+ # Delete unet adapters
1090
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
1091
+ unet.delete_adapters(adapter_names)
1092
+
1093
+ for adapter_name in adapter_names:
1094
+ # Delete text encoder adapters
1095
+ if hasattr(self, "text_encoder"):
1096
+ delete_adapter_layers(self.text_encoder, adapter_name)
1097
+ if hasattr(self, "text_encoder_2"):
1098
+ delete_adapter_layers(self.text_encoder_2, adapter_name)
1099
+
1100
+ def get_active_adapters(self) -> List[str]:
1101
+ """
1102
+ Gets the list of the current active adapters.
1103
+
1104
+ Example:
1105
+
1106
+ ```python
1107
+ from diffusers import DiffusionPipeline
1108
+
1109
+ pipeline = DiffusionPipeline.from_pretrained(
1110
+ "stabilityai/stable-diffusion-xl-base-1.0",
1111
+ ).to("cuda")
1112
+ pipeline.load_lora_weights("CiroN2022/toy-face", weight_name="toy_face_sdxl.safetensors", adapter_name="toy")
1113
+ pipeline.get_active_adapters()
1114
+ ```
1115
+ """
1116
+ if not USE_PEFT_BACKEND:
1117
+ raise ValueError(
1118
+ "PEFT backend is required for this method. Please install the latest version of PEFT `pip install -U peft`"
1119
+ )
1120
+
1121
+ from peft.tuners.tuners_utils import BaseTunerLayer
1122
+
1123
+ active_adapters = []
1124
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
1125
+ for module in unet.modules():
1126
+ if isinstance(module, BaseTunerLayer):
1127
+ active_adapters = module.active_adapters
1128
+ break
1129
+
1130
+ return active_adapters
1131
+
1132
+ def get_list_adapters(self) -> Dict[str, List[str]]:
1133
+ """
1134
+ Gets the current list of all available adapters in the pipeline.
1135
+ """
1136
+ if not USE_PEFT_BACKEND:
1137
+ raise ValueError(
1138
+ "PEFT backend is required for this method. Please install the latest version of PEFT `pip install -U peft`"
1139
+ )
1140
+
1141
+ set_adapters = {}
1142
+
1143
+ if hasattr(self, "text_encoder") and hasattr(self.text_encoder, "peft_config"):
1144
+ set_adapters["text_encoder"] = list(self.text_encoder.peft_config.keys())
1145
+
1146
+ if hasattr(self, "text_encoder_2") and hasattr(self.text_encoder_2, "peft_config"):
1147
+ set_adapters["text_encoder_2"] = list(self.text_encoder_2.peft_config.keys())
1148
+
1149
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
1150
+ if hasattr(self, self.unet_name) and hasattr(unet, "peft_config"):
1151
+ set_adapters[self.unet_name] = list(self.unet.peft_config.keys())
1152
+
1153
+ return set_adapters
1154
+
1155
+ def set_lora_device(self, adapter_names: List[str], device: Union[torch.device, str, int]) -> None:
1156
+ """
1157
+ Moves the LoRAs listed in `adapter_names` to a target device. Useful for offloading the LoRA to the CPU in case
1158
+ you want to load multiple adapters and free some GPU memory.
1159
+
1160
+ Args:
1161
+ adapter_names (`List[str]`):
1162
+ List of adapters to send device to.
1163
+ device (`Union[torch.device, str, int]`):
1164
+ Device to send the adapters to. Can be either a torch device, a str or an integer.
1165
+ """
1166
+ if not USE_PEFT_BACKEND:
1167
+ raise ValueError("PEFT backend is required for this method.")
1168
+
1169
+ from peft.tuners.tuners_utils import BaseTunerLayer
1170
+
1171
+ # Handle the UNET
1172
+ unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet
1173
+ for unet_module in unet.modules():
1174
+ if isinstance(unet_module, BaseTunerLayer):
1175
+ for adapter_name in adapter_names:
1176
+ unet_module.lora_A[adapter_name].to(device)
1177
+ unet_module.lora_B[adapter_name].to(device)
1178
+
1179
+ # Handle the text encoder
1180
+ modules_to_process = []
1181
+ if hasattr(self, "text_encoder"):
1182
+ modules_to_process.append(self.text_encoder)
1183
+
1184
+ if hasattr(self, "text_encoder_2"):
1185
+ modules_to_process.append(self.text_encoder_2)
1186
+
1187
+ for text_encoder in modules_to_process:
1188
+ # loop over submodules
1189
+ for text_encoder_module in text_encoder.modules():
1190
+ if isinstance(text_encoder_module, BaseTunerLayer):
1191
+ for adapter_name in adapter_names:
1192
+ text_encoder_module.lora_A[adapter_name].to(device)
1193
+ text_encoder_module.lora_B[adapter_name].to(device)
1194
+
1195
+
1196
+ class StableDiffusionXLLoraLoaderMixin(LoraLoaderMixin):
1197
+ """This class overrides `LoraLoaderMixin` with LoRA loading/saving code that's specific to SDXL"""
1198
+
1199
+ # Override to properly handle the loading and unloading of the additional text encoder.
1200
+ def load_lora_weights(
1201
+ self,
1202
+ pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
1203
+ adapter_name: Optional[str] = None,
1204
+ **kwargs,
1205
+ ):
1206
+ """
1207
+ Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and
1208
+ `self.text_encoder`.
1209
+
1210
+ All kwargs are forwarded to `self.lora_state_dict`.
1211
+
1212
+ See [`~loaders.LoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
1213
+
1214
+ See [`~loaders.LoraLoaderMixin.load_lora_into_unet`] for more details on how the state dict is loaded into
1215
+ `self.unet`.
1216
+
1217
+ See [`~loaders.LoraLoaderMixin.load_lora_into_text_encoder`] for more details on how the state dict is loaded
1218
+ into `self.text_encoder`.
1219
+
1220
+ Parameters:
1221
+ pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
1222
+ See [`~loaders.LoraLoaderMixin.lora_state_dict`].
1223
+ adapter_name (`str`, *optional*):
1224
+ Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
1225
+ `default_{i}` where i is the total number of adapters being loaded.
1226
+ kwargs (`dict`, *optional*):
1227
+ See [`~loaders.LoraLoaderMixin.lora_state_dict`].
1228
+ """
1229
+ if not USE_PEFT_BACKEND:
1230
+ raise ValueError("PEFT backend is required for this method.")
1231
+
1232
+ # We could have accessed the unet config from `lora_state_dict()` too. We pass
1233
+ # it here explicitly to be able to tell that it's coming from an SDXL
1234
+ # pipeline.
1235
+
1236
+ # if a dict is passed, copy it instead of modifying it inplace
1237
+ if isinstance(pretrained_model_name_or_path_or_dict, dict):
1238
+ pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
1239
+
1240
+ # First, ensure that the checkpoint is a compatible one and can be successfully loaded.
1241
+ state_dict, network_alphas = self.lora_state_dict(
1242
+ pretrained_model_name_or_path_or_dict,
1243
+ unet_config=self.unet.config,
1244
+ **kwargs,
1245
+ )
1246
+ is_correct_format = all("lora" in key for key in state_dict.keys())
1247
+ if not is_correct_format:
1248
+ raise ValueError("Invalid LoRA checkpoint.")
1249
+
1250
+ self.load_lora_into_unet(
1251
+ state_dict, network_alphas=network_alphas, unet=self.unet, adapter_name=adapter_name, _pipeline=self
1252
+ )
1253
+ text_encoder_state_dict = {k: v for k, v in state_dict.items() if "text_encoder." in k}
1254
+ if len(text_encoder_state_dict) > 0:
1255
+ self.load_lora_into_text_encoder(
1256
+ text_encoder_state_dict,
1257
+ network_alphas=network_alphas,
1258
+ text_encoder=self.text_encoder,
1259
+ prefix="text_encoder",
1260
+ lora_scale=self.lora_scale,
1261
+ adapter_name=adapter_name,
1262
+ _pipeline=self,
1263
+ )
1264
+
1265
+ text_encoder_2_state_dict = {k: v for k, v in state_dict.items() if "text_encoder_2." in k}
1266
+ if len(text_encoder_2_state_dict) > 0:
1267
+ self.load_lora_into_text_encoder(
1268
+ text_encoder_2_state_dict,
1269
+ network_alphas=network_alphas,
1270
+ text_encoder=self.text_encoder_2,
1271
+ prefix="text_encoder_2",
1272
+ lora_scale=self.lora_scale,
1273
+ adapter_name=adapter_name,
1274
+ _pipeline=self,
1275
+ )
1276
+
1277
+ @classmethod
1278
+ def save_lora_weights(
1279
+ cls,
1280
+ save_directory: Union[str, os.PathLike],
1281
+ unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
1282
+ text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
1283
+ text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
1284
+ is_main_process: bool = True,
1285
+ weight_name: str = None,
1286
+ save_function: Callable = None,
1287
+ safe_serialization: bool = True,
1288
+ ):
1289
+ r"""
1290
+ Save the LoRA parameters corresponding to the UNet and text encoder.
1291
+
1292
+ Arguments:
1293
+ save_directory (`str` or `os.PathLike`):
1294
+ Directory to save LoRA parameters to. Will be created if it doesn't exist.
1295
+ unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
1296
+ State dict of the LoRA layers corresponding to the `unet`.
1297
+ text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
1298
+ State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
1299
+ encoder LoRA state dict because it comes from 🤗 Transformers.
1300
+ is_main_process (`bool`, *optional*, defaults to `True`):
1301
+ Whether the process calling this is the main process or not. Useful during distributed training and you
1302
+ need to call this function on all processes. In this case, set `is_main_process=True` only on the main
1303
+ process to avoid race conditions.
1304
+ save_function (`Callable`):
1305
+ The function to use to save the state dictionary. Useful during distributed training when you need to
1306
+ replace `torch.save` with another method. Can be configured with the environment variable
1307
+ `DIFFUSERS_SAVE_MODE`.
1308
+ safe_serialization (`bool`, *optional*, defaults to `True`):
1309
+ Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
1310
+ """
1311
+ state_dict = {}
1312
+
1313
+ def pack_weights(layers, prefix):
1314
+ layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers
1315
+ layers_state_dict = {f"{prefix}.{module_name}": param for module_name, param in layers_weights.items()}
1316
+ return layers_state_dict
1317
+
1318
+ if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers):
1319
+ raise ValueError(
1320
+ "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers` or `text_encoder_2_lora_layers`."
1321
+ )
1322
+
1323
+ if unet_lora_layers:
1324
+ state_dict.update(pack_weights(unet_lora_layers, "unet"))
1325
+
1326
+ if text_encoder_lora_layers and text_encoder_2_lora_layers:
1327
+ state_dict.update(pack_weights(text_encoder_lora_layers, "text_encoder"))
1328
+ state_dict.update(pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
1329
+
1330
+ cls.write_lora_layers(
1331
+ state_dict=state_dict,
1332
+ save_directory=save_directory,
1333
+ is_main_process=is_main_process,
1334
+ weight_name=weight_name,
1335
+ save_function=save_function,
1336
+ safe_serialization=safe_serialization,
1337
+ )
1338
+
1339
+ def _remove_text_encoder_monkey_patch(self):
1340
+ recurse_remove_peft_layers(self.text_encoder)
1341
+ # TODO: @younesbelkada handle this in transformers side
1342
+ if getattr(self.text_encoder, "peft_config", None) is not None:
1343
+ del self.text_encoder.peft_config
1344
+ self.text_encoder._hf_peft_config_loaded = None
1345
+
1346
+ recurse_remove_peft_layers(self.text_encoder_2)
1347
+ if getattr(self.text_encoder_2, "peft_config", None) is not None:
1348
+ del self.text_encoder_2.peft_config
1349
+ self.text_encoder_2._hf_peft_config_loaded = None