shgao commited on
Commit
61bbe53
β€’
1 Parent(s): 294c39a

update demo

Browse files
LICENSE ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright [yyyy] [name of copyright owner]
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
README.md CHANGED
@@ -10,3 +10,6 @@ pinned: false
10
  ---
11
 
12
  Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
10
  ---
11
 
12
  Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
13
+
14
+ # Edit Anything by Segment-Anything
15
+
annotator/util.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import cv2
3
+ import os
4
+
5
+
6
+ annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts')
7
+
8
+
9
+ def HWC3(x):
10
+ assert x.dtype == np.uint8
11
+ if x.ndim == 2:
12
+ x = x[:, :, None]
13
+ assert x.ndim == 3
14
+ H, W, C = x.shape
15
+ assert C == 1 or C == 3 or C == 4
16
+ if C == 3:
17
+ return x
18
+ if C == 1:
19
+ return np.concatenate([x, x, x], axis=2)
20
+ if C == 4:
21
+ color = x[:, :, 0:3].astype(np.float32)
22
+ alpha = x[:, :, 3:4].astype(np.float32) / 255.0
23
+ y = color * alpha + 255.0 * (1.0 - alpha)
24
+ y = y.clip(0, 255).astype(np.uint8)
25
+ return y
26
+
27
+
28
+ def resize_image(input_image, resolution):
29
+ H, W, C = input_image.shape
30
+ H = float(H)
31
+ W = float(W)
32
+ k = float(resolution) / min(H, W)
33
+ H *= k
34
+ W *= k
35
+ H = int(np.round(H / 64.0)) * 64
36
+ W = int(np.round(W / 64.0)) * 64
37
+ img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
38
+ return img
app.py ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+
3
+
4
+ from sam2edit import create_demo as create_demo_edit_anything
5
+ from sam2image import create_demo as create_demo_generate_anything
6
+
7
+
8
+ DESCRIPTION = '# [Edit Anything](https://github.com/sail-sg/EditAnything)'
9
+
10
+ with gr.Blocks() as demo:
11
+ gr.Markdown(DESCRIPTION)
12
+ with gr.Tabs():
13
+ with gr.TabItem('Edit Anything'):
14
+ create_demo_edit_anything()
15
+ with gr.TabItem('Generate Anything'):
16
+ create_demo_generate_anything()
17
+
18
+
19
+ demo.queue(api_open=False).launch()
requirements.txt ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ cudatoolkit==11.3
2
+ pytorch==1.13.1
3
+ torchvision==0.14.1
4
+ numpy==1.23.1
5
+ gradio==3.16.2
6
+ albumentations==1.3.0
7
+ opencv-contrib-python==4.3.0.36
8
+ imageio==2.9.0
9
+ imageio-ffmpeg==0.4.2
10
+ pytorch-lightning==1.5.0
11
+ omegaconf==2.1.1
12
+ test-tube>=0.7.5
13
+ streamlit==1.12.1
14
+ einops==0.3.0
15
+ webdataset==0.2.5
16
+ kornia==0.6
17
+ open_clip_torch==2.0.2
18
+ invisible-watermark>=0.1.5
19
+ streamlit-drawable-canvas==0.8.0
20
+ torchmetrics==0.6.0
21
+ timm==0.6.12
22
+ addict==2.4.0
23
+ yapf==0.32.0
24
+ prettytable==3.6.0
25
+ safetensors==0.2.7
26
+ basicsr==1.4.2
27
+ diffusers==0.14.0
28
+ accelerate==0.14.0
29
+ transformers==4.27.4
sam2edit.py ADDED
@@ -0,0 +1,291 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Edit Anything trained with Stable Diffusion + ControlNet + SAM + BLIP2
2
+ from torchvision.utils import save_image
3
+ from PIL import Image
4
+ from pytorch_lightning import seed_everything
5
+ import subprocess
6
+ from collections import OrderedDict
7
+
8
+ import cv2
9
+ import einops
10
+ import gradio as gr
11
+ import numpy as np
12
+ import torch
13
+ import random
14
+ import os
15
+ import requests
16
+ from io import BytesIO
17
+ from annotator.util import resize_image, HWC3
18
+
19
+ def create_demo():
20
+ device = "cuda" if torch.cuda.is_available() else "cpu"
21
+ use_blip = True
22
+ use_gradio = True
23
+
24
+ # Diffusion init using diffusers.
25
+
26
+ # diffusers==0.14.0 required.
27
+ from diffusers import ControlNetModel, UniPCMultistepScheduler
28
+ from utils.stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
29
+ from diffusers.utils import load_image
30
+
31
+ base_model_path = "stabilityai/stable-diffusion-2-inpainting"
32
+ config_dict = OrderedDict([('SAM Pretrained(v0-1)', 'shgao/edit-anything-v0-1-1'),
33
+ ('LAION Pretrained(v0-3)', 'shgao/edit-anything-v0-3'),
34
+ ])
35
+ def obtain_generation_model(controlnet_path):
36
+ controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
37
+ pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
38
+ base_model_path, controlnet=controlnet, torch_dtype=torch.float16
39
+ )
40
+ # speed up diffusion process with faster scheduler and memory optimization
41
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
42
+ # remove following line if xformers is not installed
43
+ pipe.enable_xformers_memory_efficient_attention()
44
+
45
+ # pipe.enable_model_cpu_offload() # disable for now because of unknow bug in accelerate
46
+ pipe.to(device)
47
+ return pipe
48
+ global default_controlnet_path
49
+ default_controlnet_path = config_dict['LAION Pretrained(v0-3)']
50
+ pipe = obtain_generation_model(default_controlnet_path)
51
+
52
+ # Segment-Anything init.
53
+ # pip install git+https://github.com/facebookresearch/segment-anything.git
54
+ from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
55
+
56
+ try:
57
+ from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
58
+ except ImportError:
59
+ print('segment_anything not installed')
60
+ result = subprocess.run(['pip', 'install', 'git+https://github.com/facebookresearch/segment-anything.git'], check=True)
61
+ print(f'Install segment_anything {result}')
62
+ if not os.path.exists('./models/sam_vit_h_4b8939.pth'):
63
+ result = subprocess.run(['wget', 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth', '-P', 'models'], check=True)
64
+ print(f'Download sam_vit_h_4b8939.pth {result}')
65
+ sam_checkpoint = "models/sam_vit_h_4b8939.pth"
66
+ model_type = "default"
67
+ sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
68
+ sam.to(device=device)
69
+ mask_generator = SamAutomaticMaskGenerator(sam)
70
+
71
+
72
+ # BLIP2 init.
73
+ if use_blip:
74
+ # need the latest transformers
75
+ # pip install git+https://github.com/huggingface/transformers.git
76
+ from transformers import AutoProcessor, Blip2ForConditionalGeneration
77
+
78
+ processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
79
+ blip_model = Blip2ForConditionalGeneration.from_pretrained(
80
+ "Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16)
81
+ blip_model.to(device)
82
+ blip_model.to(device)
83
+
84
+
85
+ def get_blip2_text(image):
86
+ inputs = processor(image, return_tensors="pt").to(device, torch.float16)
87
+ generated_ids = blip_model.generate(**inputs, max_new_tokens=50)
88
+ generated_text = processor.batch_decode(
89
+ generated_ids, skip_special_tokens=True)[0].strip()
90
+ return generated_text
91
+
92
+
93
+ def show_anns(anns):
94
+ if len(anns) == 0:
95
+ return
96
+ sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
97
+ full_img = None
98
+
99
+ # for ann in sorted_anns:
100
+ for i in range(len(sorted_anns)):
101
+ ann = anns[i]
102
+ m = ann['segmentation']
103
+ if full_img is None:
104
+ full_img = np.zeros((m.shape[0], m.shape[1], 3))
105
+ map = np.zeros((m.shape[0], m.shape[1]), dtype=np.uint16)
106
+ map[m != 0] = i + 1
107
+ color_mask = np.random.random((1, 3)).tolist()[0]
108
+ full_img[m != 0] = color_mask
109
+ full_img = full_img*255
110
+ # anno encoding from https://github.com/LUSSeg/ImageNet-S
111
+ res = np.zeros((map.shape[0], map.shape[1], 3))
112
+ res[:, :, 0] = map % 256
113
+ res[:, :, 1] = map // 256
114
+ res.astype(np.float32)
115
+ full_img = Image.fromarray(np.uint8(full_img))
116
+ return full_img, res
117
+
118
+
119
+ def get_sam_control(image):
120
+ masks = mask_generator.generate(image)
121
+ full_img, res = show_anns(masks)
122
+ return full_img, res
123
+
124
+
125
+ def process(condition_model, source_image, mask_image, enable_auto_prompt, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta):
126
+
127
+ input_image = source_image["image"]
128
+ if mask_image is None:
129
+ mask_image = source_image["mask"]
130
+ global default_controlnet_path
131
+ global pipe
132
+ print("To Use:", config_dict[condition_model], "Current:", default_controlnet_path)
133
+ if default_controlnet_path!=config_dict[condition_model]:
134
+ print("Change condition model to:", config_dict[condition_model])
135
+ pipe = obtain_generation_model(config_dict[condition_model])
136
+ default_controlnet_path = config_dict[condition_model]
137
+
138
+ with torch.no_grad():
139
+ if use_blip and (enable_auto_prompt or len(prompt) == 0):
140
+ print("Generating text:")
141
+ blip2_prompt = get_blip2_text(input_image)
142
+ print("Generated text:", blip2_prompt)
143
+ if len(prompt)>0:
144
+ prompt = blip2_prompt + ',' + prompt
145
+ else:
146
+ prompt = blip2_prompt
147
+ print("All text:", prompt)
148
+
149
+ input_image = HWC3(input_image)
150
+
151
+ img = resize_image(input_image, image_resolution)
152
+ H, W, C = img.shape
153
+
154
+ print("Generating SAM seg:")
155
+ # the default SAM model is trained with 1024 size.
156
+ full_segmask, detected_map = get_sam_control(
157
+ resize_image(input_image, detect_resolution))
158
+
159
+ detected_map = HWC3(detected_map.astype(np.uint8))
160
+ detected_map = cv2.resize(
161
+ detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
162
+
163
+ control = torch.from_numpy(
164
+ detected_map.copy()).float().cuda()
165
+ control = torch.stack([control for _ in range(num_samples)], dim=0)
166
+ control = einops.rearrange(control, 'b h w c -> b c h w').clone()
167
+
168
+ mask_image = HWC3(mask_image.astype(np.uint8))
169
+ mask_image = cv2.resize(
170
+ mask_image, (W, H), interpolation=cv2.INTER_LINEAR)
171
+ mask_image = Image.fromarray(mask_image)
172
+
173
+
174
+ if seed == -1:
175
+ seed = random.randint(0, 65535)
176
+ seed_everything(seed)
177
+ generator = torch.manual_seed(seed)
178
+ x_samples = pipe(
179
+ image=img,
180
+ mask_image=mask_image,
181
+ prompt=[prompt + ', ' + a_prompt] * num_samples,
182
+ negative_prompt=[n_prompt] * num_samples,
183
+ num_images_per_prompt=num_samples,
184
+ num_inference_steps=ddim_steps,
185
+ generator=generator,
186
+ controlnet_conditioning_image=control.type(torch.float16),
187
+ height=H,
188
+ width=W,
189
+ ).images
190
+
191
+
192
+ results = [x_samples[i] for i in range(num_samples)]
193
+ return [full_segmask, mask_image] + results, prompt
194
+
195
+
196
+ def download_image(url):
197
+ response = requests.get(url)
198
+ return Image.open(BytesIO(response.content)).convert("RGB")
199
+
200
+ # disable gradio when not using GUI.
201
+ if not use_gradio:
202
+ # This part is not updated, it's just a example to use it without GUI.
203
+ image_path = "../data/samples/sa_223750.jpg"
204
+ mask_path = "../data/samples/sa_223750inpaint.png"
205
+ input_image = Image.open(image_path)
206
+ mask_image = Image.open(mask_path)
207
+ enable_auto_prompt = True
208
+ input_image = np.array(input_image, dtype=np.uint8)
209
+ mask_image = np.array(mask_image, dtype=np.uint8)
210
+ prompt = "esplendent sunset sky, red brick wall"
211
+ a_prompt = 'best quality, extremely detailed'
212
+ n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
213
+ num_samples = 3
214
+ image_resolution = 512
215
+ detect_resolution = 512
216
+ ddim_steps = 30
217
+ guess_mode = False
218
+ strength = 1.0
219
+ scale = 9.0
220
+ seed = -1
221
+ eta = 0.0
222
+
223
+ outputs = process(condition_model, input_image, mask_image, enable_auto_prompt, prompt, a_prompt, n_prompt, num_samples, image_resolution,
224
+ detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta)
225
+
226
+ image_list = []
227
+ input_image = resize_image(input_image, 512)
228
+ image_list.append(torch.tensor(input_image))
229
+ for i in range(len(outputs)):
230
+ each = outputs[i]
231
+ if type(each) is not np.ndarray:
232
+ each = np.array(each, dtype=np.uint8)
233
+ each = resize_image(each, 512)
234
+ print(i, each.shape)
235
+ image_list.append(torch.tensor(each))
236
+
237
+ image_list = torch.stack(image_list).permute(0, 3, 1, 2)
238
+
239
+ save_image(image_list, "sample.jpg", nrow=3,
240
+ normalize=True, value_range=(0, 255))
241
+ else:
242
+ print("The GUI is not fully tested yet. Please open an issue if you find bugs.")
243
+ block = gr.Blocks()
244
+ with block as demo:
245
+ with gr.Row():
246
+ gr.Markdown(
247
+ "## Edit Anything")
248
+ with gr.Row():
249
+ with gr.Column():
250
+ source_image = gr.Image(source='upload',label="Image (support sketch)", type="numpy", tool="sketch")
251
+ mask_image = gr.Image(source='upload', label="Edit region (Optional)", type="numpy", value=None)
252
+ prompt = gr.Textbox(label="Prompt")
253
+ enable_auto_prompt = gr.Checkbox(label='Auto generated BLIP2 prompt', value=True)
254
+ run_button = gr.Button(label="Run")
255
+ condition_model = gr.Dropdown(choices=list(config_dict.keys()),
256
+ value=list(config_dict.keys())[0],
257
+ label='Model',
258
+ multiselect=False)
259
+ num_samples = gr.Slider(
260
+ label="Images", minimum=1, maximum=12, value=1, step=1)
261
+ with gr.Accordion("Advanced options", open=False):
262
+ image_resolution = gr.Slider(
263
+ label="Image Resolution", minimum=256, maximum=768, value=512, step=64)
264
+ strength = gr.Slider(
265
+ label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
266
+ guess_mode = gr.Checkbox(label='Guess Mode', value=False)
267
+ detect_resolution = gr.Slider(
268
+ label="SAM Resolution", minimum=128, maximum=2048, value=1024, step=1)
269
+ ddim_steps = gr.Slider(
270
+ label="Steps", minimum=1, maximum=100, value=20, step=1)
271
+ scale = gr.Slider(
272
+ label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
273
+ seed = gr.Slider(label="Seed", minimum=-1,
274
+ maximum=2147483647, step=1, randomize=True)
275
+ eta = gr.Number(label="eta (DDIM)", value=0.0)
276
+ a_prompt = gr.Textbox(
277
+ label="Added Prompt", value='best quality, extremely detailed')
278
+ n_prompt = gr.Textbox(label="Negative Prompt",
279
+ value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
280
+ with gr.Column():
281
+ result_gallery = gr.Gallery(
282
+ label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
283
+ result_text = gr.Text(label='BLIP2+Human Prompt Text')
284
+ ips = [condition_model, source_image, mask_image, enable_auto_prompt, prompt, a_prompt, n_prompt, num_samples, image_resolution,
285
+ detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta]
286
+ run_button.click(fn=process, inputs=ips, outputs=[result_gallery, result_text])
287
+ return demo
288
+
289
+ if __name__ == '__main__':
290
+ demo = create_demo()
291
+ demo.queue().launch(server_name='0.0.0.0')
sam2image.py ADDED
@@ -0,0 +1,273 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Edit Anything trained with Stable Diffusion + ControlNet + SAM + BLIP2
2
+ from diffusers.utils import load_image
3
+ from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
4
+ from torchvision.utils import save_image
5
+ from PIL import Image
6
+ from pytorch_lightning import seed_everything
7
+ import subprocess
8
+ from collections import OrderedDict
9
+
10
+ import cv2
11
+ import einops
12
+ import gradio as gr
13
+ import numpy as np
14
+ import torch
15
+ import random
16
+ import os
17
+ from annotator.util import resize_image, HWC3
18
+
19
+ def create_demo():
20
+ device = "cuda" if torch.cuda.is_available() else "cpu"
21
+ use_blip = True
22
+ use_gradio = True
23
+
24
+ # Diffusion init using diffusers.
25
+ # diffusers==0.14.0 required.
26
+
27
+ base_model_path = "stabilityai/stable-diffusion-2-1"
28
+
29
+
30
+ config_dict = OrderedDict([('SAM Pretrained(v0-1)', 'shgao/edit-anything-v0-1-1'),
31
+ ('LAION Pretrained(v0-3)', 'shgao/edit-anything-v0-3'),
32
+ ])
33
+
34
+
35
+ def obtain_generation_model(controlnet_path):
36
+ controlnet = ControlNetModel.from_pretrained(
37
+ controlnet_path, torch_dtype=torch.float16)
38
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
39
+ base_model_path, controlnet=controlnet, torch_dtype=torch.float16
40
+ )
41
+ # speed up diffusion process with faster scheduler and memory optimization
42
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
43
+ # remove following line if xformers is not installed
44
+ pipe.enable_xformers_memory_efficient_attention()
45
+ # pipe.enable_model_cpu_offload() # disable for now because of unknow bug in accelerate
46
+ pipe.to(device)
47
+ return pipe
48
+
49
+ global default_controlnet_path
50
+ default_controlnet_path = config_dict['LAION Pretrained(v0-3)']
51
+ pipe = obtain_generation_model(default_controlnet_path)
52
+
53
+
54
+ # Segment-Anything init.
55
+ # pip install git+https://github.com/facebookresearch/segment-anything.git
56
+ try:
57
+ from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
58
+ except ImportError:
59
+ print('segment_anything not installed')
60
+ result = subprocess.run(['pip', 'install', 'git+https://github.com/facebookresearch/segment-anything.git'], check=True)
61
+ print(f'Install segment_anything {result}')
62
+ if not os.path.exists('./models/sam_vit_h_4b8939.pth'):
63
+ result = subprocess.run(['wget', 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth', '-P', 'models'], check=True)
64
+ print(f'Download sam_vit_h_4b8939.pth {result}')
65
+ sam_checkpoint = "models/sam_vit_h_4b8939.pth"
66
+ model_type = "default"
67
+ sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
68
+ sam.to(device=device)
69
+ mask_generator = SamAutomaticMaskGenerator(sam)
70
+
71
+
72
+ # BLIP2 init.
73
+ if use_blip:
74
+ # need the latest transformers
75
+ # pip install git+https://github.com/huggingface/transformers.git
76
+ from transformers import AutoProcessor, Blip2ForConditionalGeneration
77
+
78
+ processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
79
+ blip_model = Blip2ForConditionalGeneration.from_pretrained(
80
+ "Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16)
81
+ blip_model.to(device)
82
+ blip_model.to(device)
83
+
84
+
85
+ def get_blip2_text(image):
86
+ inputs = processor(image, return_tensors="pt").to(device, torch.float16)
87
+ generated_ids = blip_model.generate(**inputs, max_new_tokens=50)
88
+ generated_text = processor.batch_decode(
89
+ generated_ids, skip_special_tokens=True)[0].strip()
90
+ return generated_text
91
+
92
+
93
+ def show_anns(anns):
94
+ if len(anns) == 0:
95
+ return
96
+ sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
97
+ full_img = None
98
+
99
+ # for ann in sorted_anns:
100
+ for i in range(len(sorted_anns)):
101
+ ann = anns[i]
102
+ m = ann['segmentation']
103
+ if full_img is None:
104
+ full_img = np.zeros((m.shape[0], m.shape[1], 3))
105
+ map = np.zeros((m.shape[0], m.shape[1]), dtype=np.uint16)
106
+ map[m != 0] = i + 1
107
+ color_mask = np.random.random((1, 3)).tolist()[0]
108
+ full_img[m != 0] = color_mask
109
+ full_img = full_img*255
110
+ # anno encoding from https://github.com/LUSSeg/ImageNet-S
111
+ res = np.zeros((map.shape[0], map.shape[1], 3))
112
+ res[:, :, 0] = map % 256
113
+ res[:, :, 1] = map // 256
114
+ res.astype(np.float32)
115
+ full_img = Image.fromarray(np.uint8(full_img))
116
+ return full_img, res
117
+
118
+
119
+ def get_sam_control(image):
120
+ masks = mask_generator.generate(image)
121
+ full_img, res = show_anns(masks)
122
+ return full_img, res
123
+
124
+
125
+ def process(condition_model, input_image, enable_auto_prompt, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta):
126
+
127
+ global default_controlnet_path
128
+ global pipe
129
+ print("To Use:", config_dict[condition_model], "Current:", default_controlnet_path)
130
+ if default_controlnet_path!=config_dict[condition_model]:
131
+ print("Change condition model to:", config_dict[condition_model])
132
+ pipe = obtain_generation_model(config_dict[condition_model])
133
+ default_controlnet_path = config_dict[condition_model]
134
+
135
+ with torch.no_grad():
136
+ if use_blip and (enable_auto_prompt or len(prompt) == 0):
137
+ print("Generating text:")
138
+ blip2_prompt = get_blip2_text(input_image)
139
+ print("Generated text:", blip2_prompt)
140
+ if len(prompt) > 0:
141
+ prompt = blip2_prompt + ',' + prompt
142
+ else:
143
+ prompt = blip2_prompt
144
+ print("All text:", prompt)
145
+
146
+ input_image = HWC3(input_image)
147
+
148
+ img = resize_image(input_image, image_resolution)
149
+ H, W, C = img.shape
150
+
151
+ print("Generating SAM seg:")
152
+ # the default SAM model is trained with 1024 size.
153
+ full_segmask, detected_map = get_sam_control(
154
+ resize_image(input_image, detect_resolution))
155
+
156
+ detected_map = HWC3(detected_map.astype(np.uint8))
157
+ detected_map = cv2.resize(
158
+ detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
159
+
160
+ control = torch.from_numpy(
161
+ detected_map.copy()).float().cuda()
162
+ control = torch.stack([control for _ in range(num_samples)], dim=0)
163
+ control = einops.rearrange(control, 'b h w c -> b c h w').clone()
164
+
165
+ if seed == -1:
166
+ seed = random.randint(0, 65535)
167
+ seed_everything(seed)
168
+ print("control.shape", control.shape)
169
+ generator = torch.manual_seed(seed)
170
+ x_samples = pipe(
171
+ prompt=[prompt + ', ' + a_prompt] * num_samples,
172
+ negative_prompt=[n_prompt] * num_samples,
173
+ num_images_per_prompt=num_samples,
174
+ num_inference_steps=ddim_steps,
175
+ generator=generator,
176
+ height=H,
177
+ width=W,
178
+ image=control.type(torch.float16),
179
+ ).images
180
+
181
+ results = [x_samples[i] for i in range(num_samples)]
182
+ return [full_segmask] + results, prompt
183
+
184
+
185
+ # disable gradio when not using GUI.
186
+ if not use_gradio:
187
+ # This part is not updated, it's just a example to use it without GUI.
188
+ condition_model = 'shgao/edit-anything-v0-1-1'
189
+ image_path = "images/sa_309398.jpg"
190
+ input_image = Image.open(image_path)
191
+ input_image = np.array(input_image, dtype=np.uint8)
192
+ prompt = ""
193
+ a_prompt = 'best quality, extremely detailed'
194
+ n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
195
+ num_samples = 4
196
+ image_resolution = 512
197
+ detect_resolution = 512
198
+ ddim_steps = 100
199
+ guess_mode = False
200
+ strength = 1.0
201
+ scale = 9.0
202
+ seed = 10086
203
+ eta = 0.0
204
+
205
+ outputs, full_text = process(condition_model, input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution,
206
+ detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta)
207
+
208
+ image_list = []
209
+ input_image = resize_image(input_image, 512)
210
+ image_list.append(torch.tensor(input_image))
211
+ for i in range(len(outputs)):
212
+ each = outputs[i]
213
+ if type(each) is not np.ndarray:
214
+ each = np.array(each, dtype=np.uint8)
215
+ each = resize_image(each, 512)
216
+ print(i, each.shape)
217
+ image_list.append(torch.tensor(each))
218
+
219
+ image_list = torch.stack(image_list).permute(0, 3, 1, 2)
220
+
221
+ save_image(image_list, "sample.jpg", nrow=3,
222
+ normalize=True, value_range=(0, 255))
223
+ else:
224
+ block = gr.Blocks()
225
+ with block as demo:
226
+ with gr.Row():
227
+ gr.Markdown(
228
+ "## Generate Anything")
229
+ with gr.Row():
230
+ with gr.Column():
231
+ input_image = gr.Image(source='upload', type="numpy")
232
+ prompt = gr.Textbox(label="Prompt (Optional)")
233
+ run_button = gr.Button(label="Run")
234
+ condition_model = gr.Dropdown(choices=list(config_dict.keys()),
235
+ value=list(config_dict.keys())[0],
236
+ label='Model',
237
+ multiselect=False)
238
+ num_samples = gr.Slider(
239
+ label="Images", minimum=1, maximum=12, value=1, step=1)
240
+
241
+ enable_auto_prompt = gr.Checkbox(label='Auto generated BLIP2 prompt', value=True)
242
+ with gr.Accordion("Advanced options", open=False):
243
+ image_resolution = gr.Slider(
244
+ label="Image Resolution", minimum=256, maximum=768, value=512, step=64)
245
+ strength = gr.Slider(
246
+ label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
247
+ guess_mode = gr.Checkbox(label='Guess Mode', value=False)
248
+ detect_resolution = gr.Slider(
249
+ label="SAM Resolution", minimum=128, maximum=2048, value=1024, step=1)
250
+ ddim_steps = gr.Slider(
251
+ label="Steps", minimum=1, maximum=100, value=20, step=1)
252
+ scale = gr.Slider(
253
+ label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
254
+ seed = gr.Slider(label="Seed", minimum=-1,
255
+ maximum=2147483647, step=1, randomize=True)
256
+ eta = gr.Number(label="eta (DDIM)", value=0.0)
257
+ a_prompt = gr.Textbox(
258
+ label="Added Prompt", value='best quality, extremely detailed')
259
+ n_prompt = gr.Textbox(label="Negative Prompt",
260
+ value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
261
+
262
+ with gr.Column():
263
+ result_gallery = gr.Gallery(
264
+ label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
265
+ result_text = gr.Text(label='BLIP2+Human Prompt Text')
266
+ ips = [condition_model, input_image, enable_auto_prompt, prompt, a_prompt, n_prompt, num_samples, image_resolution,
267
+ detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta]
268
+ run_button.click(fn=process, inputs=ips, outputs=[result_gallery, result_text])
269
+ return demo
270
+
271
+ if __name__ == '__main__':
272
+ demo = create_demo()
273
+ demo.queue().launch(server_name='0.0.0.0')
stable_diffusion_controlnet_inpaint.py ADDED
@@ -0,0 +1,1076 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Inspired by: https://github.com/haofanwang/ControlNet-for-Diffusers/
2
+ # From https://raw.githubusercontent.com/huggingface/diffusers/53377ef83c6446033f3ee506e3ef718db817b293/examples/community/stable_diffusion_controlnet_inpaint.py
3
+ import inspect
4
+ from typing import Any, Callable, Dict, List, Optional, Union
5
+
6
+ import numpy as np
7
+ import PIL.Image
8
+ import torch
9
+ import torch.nn.functional as F
10
+ from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
11
+
12
+ from diffusers import AutoencoderKL, ControlNetModel, DiffusionPipeline, UNet2DConditionModel, logging
13
+ from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker
14
+ from diffusers.schedulers import KarrasDiffusionSchedulers
15
+ from diffusers.utils import (
16
+ PIL_INTERPOLATION,
17
+ is_accelerate_available,
18
+ is_accelerate_version,
19
+ randn_tensor,
20
+ replace_example_docstring,
21
+ )
22
+
23
+
24
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
25
+
26
+ EXAMPLE_DOC_STRING = """
27
+ Examples:
28
+ ```py
29
+ >>> import numpy as np
30
+ >>> import torch
31
+ >>> from PIL import Image
32
+ >>> from stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
33
+
34
+ >>> from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
35
+ >>> from diffusers import ControlNetModel, UniPCMultistepScheduler
36
+ >>> from diffusers.utils import load_image
37
+
38
+ >>> def ade_palette():
39
+ return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
40
+ [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
41
+ [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
42
+ [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
43
+ [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
44
+ [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
45
+ [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
46
+ [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
47
+ [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
48
+ [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
49
+ [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
50
+ [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
51
+ [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
52
+ [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
53
+ [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
54
+ [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
55
+ [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
56
+ [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
57
+ [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
58
+ [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
59
+ [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
60
+ [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
61
+ [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
62
+ [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
63
+ [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
64
+ [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
65
+ [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
66
+ [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
67
+ [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
68
+ [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
69
+ [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
70
+ [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
71
+ [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
72
+ [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
73
+ [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
74
+ [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
75
+ [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
76
+ [102, 255, 0], [92, 0, 255]]
77
+
78
+ >>> image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
79
+ >>> image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
80
+
81
+ >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-seg", torch_dtype=torch.float16)
82
+
83
+ >>> pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
84
+ "runwayml/stable-diffusion-inpainting", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
85
+ )
86
+
87
+ >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
88
+ >>> pipe.enable_xformers_memory_efficient_attention()
89
+ >>> pipe.enable_model_cpu_offload()
90
+
91
+ >>> def image_to_seg(image):
92
+ pixel_values = image_processor(image, return_tensors="pt").pixel_values
93
+ with torch.no_grad():
94
+ outputs = image_segmentor(pixel_values)
95
+ seg = image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
96
+ color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
97
+ palette = np.array(ade_palette())
98
+ for label, color in enumerate(palette):
99
+ color_seg[seg == label, :] = color
100
+ color_seg = color_seg.astype(np.uint8)
101
+ seg_image = Image.fromarray(color_seg)
102
+ return seg_image
103
+
104
+ >>> image = load_image(
105
+ "https://github.com/CompVis/latent-diffusion/raw/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
106
+ )
107
+
108
+ >>> mask_image = load_image(
109
+ "https://github.com/CompVis/latent-diffusion/raw/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
110
+ )
111
+
112
+ >>> controlnet_conditioning_image = image_to_seg(image)
113
+
114
+ >>> image = pipe(
115
+ "Face of a yellow cat, high resolution, sitting on a park bench",
116
+ image,
117
+ mask_image,
118
+ controlnet_conditioning_image,
119
+ num_inference_steps=20,
120
+ ).images[0]
121
+
122
+ >>> image.save("out.png")
123
+ ```
124
+ """
125
+
126
+
127
+ def prepare_image(image):
128
+ if isinstance(image, torch.Tensor):
129
+ # Batch single image
130
+ if image.ndim == 3:
131
+ image = image.unsqueeze(0)
132
+
133
+ image = image.to(dtype=torch.float32)
134
+ else:
135
+ # preprocess image
136
+ if isinstance(image, (PIL.Image.Image, np.ndarray)):
137
+ image = [image]
138
+
139
+ if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):
140
+ image = [np.array(i.convert("RGB"))[None, :] for i in image]
141
+ image = np.concatenate(image, axis=0)
142
+ elif isinstance(image, list) and isinstance(image[0], np.ndarray):
143
+ image = np.concatenate([i[None, :] for i in image], axis=0)
144
+
145
+ image = image.transpose(0, 3, 1, 2)
146
+ image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0
147
+
148
+ return image
149
+
150
+
151
+ def prepare_mask_image(mask_image):
152
+ if isinstance(mask_image, torch.Tensor):
153
+ if mask_image.ndim == 2:
154
+ # Batch and add channel dim for single mask
155
+ mask_image = mask_image.unsqueeze(0).unsqueeze(0)
156
+ elif mask_image.ndim == 3 and mask_image.shape[0] == 1:
157
+ # Single mask, the 0'th dimension is considered to be
158
+ # the existing batch size of 1
159
+ mask_image = mask_image.unsqueeze(0)
160
+ elif mask_image.ndim == 3 and mask_image.shape[0] != 1:
161
+ # Batch of mask, the 0'th dimension is considered to be
162
+ # the batching dimension
163
+ mask_image = mask_image.unsqueeze(1)
164
+
165
+ # Binarize mask
166
+ mask_image[mask_image < 0.5] = 0
167
+ mask_image[mask_image >= 0.5] = 1
168
+ else:
169
+ # preprocess mask
170
+ if isinstance(mask_image, (PIL.Image.Image, np.ndarray)):
171
+ mask_image = [mask_image]
172
+
173
+ if isinstance(mask_image, list) and isinstance(mask_image[0], PIL.Image.Image):
174
+ mask_image = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask_image], axis=0)
175
+ mask_image = mask_image.astype(np.float32) / 255.0
176
+ elif isinstance(mask_image, list) and isinstance(mask_image[0], np.ndarray):
177
+ mask_image = np.concatenate([m[None, None, :] for m in mask_image], axis=0)
178
+
179
+ mask_image[mask_image < 0.5] = 0
180
+ mask_image[mask_image >= 0.5] = 1
181
+ mask_image = torch.from_numpy(mask_image)
182
+
183
+ return mask_image
184
+
185
+
186
+ def prepare_controlnet_conditioning_image(
187
+ controlnet_conditioning_image, width, height, batch_size, num_images_per_prompt, device, dtype
188
+ ):
189
+ if not isinstance(controlnet_conditioning_image, torch.Tensor):
190
+ if isinstance(controlnet_conditioning_image, PIL.Image.Image):
191
+ controlnet_conditioning_image = [controlnet_conditioning_image]
192
+
193
+ if isinstance(controlnet_conditioning_image[0], PIL.Image.Image):
194
+ controlnet_conditioning_image = [
195
+ np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"]))[None, :]
196
+ for i in controlnet_conditioning_image
197
+ ]
198
+ controlnet_conditioning_image = np.concatenate(controlnet_conditioning_image, axis=0)
199
+ controlnet_conditioning_image = np.array(controlnet_conditioning_image).astype(np.float32) / 255.0
200
+ controlnet_conditioning_image = controlnet_conditioning_image.transpose(0, 3, 1, 2)
201
+ controlnet_conditioning_image = torch.from_numpy(controlnet_conditioning_image)
202
+ elif isinstance(controlnet_conditioning_image[0], torch.Tensor):
203
+ controlnet_conditioning_image = torch.cat(controlnet_conditioning_image, dim=0)
204
+
205
+ image_batch_size = controlnet_conditioning_image.shape[0]
206
+
207
+ if image_batch_size == 1:
208
+ repeat_by = batch_size
209
+ else:
210
+ # image batch size is the same as prompt batch size
211
+ repeat_by = num_images_per_prompt
212
+
213
+ controlnet_conditioning_image = controlnet_conditioning_image.repeat_interleave(repeat_by, dim=0)
214
+
215
+ controlnet_conditioning_image = controlnet_conditioning_image.to(device=device, dtype=dtype)
216
+
217
+ return controlnet_conditioning_image
218
+
219
+
220
+ class StableDiffusionControlNetInpaintPipeline(DiffusionPipeline):
221
+ """
222
+ Inspired by: https://github.com/haofanwang/ControlNet-for-Diffusers/
223
+ """
224
+
225
+ _optional_components = ["safety_checker", "feature_extractor"]
226
+
227
+ def __init__(
228
+ self,
229
+ vae: AutoencoderKL,
230
+ text_encoder: CLIPTextModel,
231
+ tokenizer: CLIPTokenizer,
232
+ unet: UNet2DConditionModel,
233
+ controlnet: ControlNetModel,
234
+ scheduler: KarrasDiffusionSchedulers,
235
+ safety_checker: StableDiffusionSafetyChecker,
236
+ feature_extractor: CLIPImageProcessor,
237
+ requires_safety_checker: bool = True,
238
+ ):
239
+ super().__init__()
240
+
241
+ if safety_checker is None and requires_safety_checker:
242
+ logger.warning(
243
+ f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
244
+ " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
245
+ " results in services or applications open to the public. Both the diffusers team and Hugging Face"
246
+ " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
247
+ " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
248
+ " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
249
+ )
250
+
251
+ if safety_checker is not None and feature_extractor is None:
252
+ raise ValueError(
253
+ "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
254
+ " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
255
+ )
256
+
257
+ self.register_modules(
258
+ vae=vae,
259
+ text_encoder=text_encoder,
260
+ tokenizer=tokenizer,
261
+ unet=unet,
262
+ controlnet=controlnet,
263
+ scheduler=scheduler,
264
+ safety_checker=safety_checker,
265
+ feature_extractor=feature_extractor,
266
+ )
267
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
268
+ self.register_to_config(requires_safety_checker=requires_safety_checker)
269
+
270
+ def enable_vae_slicing(self):
271
+ r"""
272
+ Enable sliced VAE decoding.
273
+
274
+ When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
275
+ steps. This is useful to save some memory and allow larger batch sizes.
276
+ """
277
+ self.vae.enable_slicing()
278
+
279
+ def disable_vae_slicing(self):
280
+ r"""
281
+ Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
282
+ computing decoding in one step.
283
+ """
284
+ self.vae.disable_slicing()
285
+
286
+ def enable_sequential_cpu_offload(self, gpu_id=0):
287
+ r"""
288
+ Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
289
+ text_encoder, vae, controlnet, and safety checker have their state dicts saved to CPU and then are moved to a
290
+ `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
291
+ Note that offloading happens on a submodule basis. Memory savings are higher than with
292
+ `enable_model_cpu_offload`, but performance is lower.
293
+ """
294
+ if is_accelerate_available():
295
+ from accelerate import cpu_offload
296
+ else:
297
+ raise ImportError("Please install accelerate via `pip install accelerate`")
298
+
299
+ device = torch.device(f"cuda:{gpu_id}")
300
+
301
+ for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.controlnet]:
302
+ cpu_offload(cpu_offloaded_model, device)
303
+
304
+ if self.safety_checker is not None:
305
+ cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True)
306
+
307
+ def enable_model_cpu_offload(self, gpu_id=0):
308
+ r"""
309
+ Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
310
+ to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
311
+ method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
312
+ `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
313
+ """
314
+ if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
315
+ from accelerate import cpu_offload_with_hook
316
+ else:
317
+ raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")
318
+
319
+ device = torch.device(f"cuda:{gpu_id}")
320
+
321
+ hook = None
322
+ for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
323
+ _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
324
+
325
+ if self.safety_checker is not None:
326
+ # the safety checker can offload the vae again
327
+ _, hook = cpu_offload_with_hook(self.safety_checker, device, prev_module_hook=hook)
328
+
329
+ # control net hook has be manually offloaded as it alternates with unet
330
+ cpu_offload_with_hook(self.controlnet, device)
331
+
332
+ # We'll offload the last model manually.
333
+ self.final_offload_hook = hook
334
+
335
+ @property
336
+ def _execution_device(self):
337
+ r"""
338
+ Returns the device on which the pipeline's models will be executed. After calling
339
+ `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
340
+ hooks.
341
+ """
342
+ if not hasattr(self.unet, "_hf_hook"):
343
+ return self.device
344
+ for module in self.unet.modules():
345
+ if (
346
+ hasattr(module, "_hf_hook")
347
+ and hasattr(module._hf_hook, "execution_device")
348
+ and module._hf_hook.execution_device is not None
349
+ ):
350
+ return torch.device(module._hf_hook.execution_device)
351
+ return self.device
352
+
353
+ def _encode_prompt(
354
+ self,
355
+ prompt,
356
+ device,
357
+ num_images_per_prompt,
358
+ do_classifier_free_guidance,
359
+ negative_prompt=None,
360
+ prompt_embeds: Optional[torch.FloatTensor] = None,
361
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
362
+ ):
363
+ r"""
364
+ Encodes the prompt into text encoder hidden states.
365
+
366
+ Args:
367
+ prompt (`str` or `List[str]`, *optional*):
368
+ prompt to be encoded
369
+ device: (`torch.device`):
370
+ torch device
371
+ num_images_per_prompt (`int`):
372
+ number of images that should be generated per prompt
373
+ do_classifier_free_guidance (`bool`):
374
+ whether to use classifier free guidance or not
375
+ negative_prompt (`str` or `List[str]`, *optional*):
376
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead.
377
+ Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
378
+ prompt_embeds (`torch.FloatTensor`, *optional*):
379
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
380
+ provided, text embeddings will be generated from `prompt` input argument.
381
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
382
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
383
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
384
+ argument.
385
+ """
386
+ if prompt is not None and isinstance(prompt, str):
387
+ batch_size = 1
388
+ elif prompt is not None and isinstance(prompt, list):
389
+ batch_size = len(prompt)
390
+ else:
391
+ batch_size = prompt_embeds.shape[0]
392
+
393
+ if prompt_embeds is None:
394
+ text_inputs = self.tokenizer(
395
+ prompt,
396
+ padding="max_length",
397
+ max_length=self.tokenizer.model_max_length,
398
+ truncation=True,
399
+ return_tensors="pt",
400
+ )
401
+ text_input_ids = text_inputs.input_ids
402
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
403
+
404
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
405
+ text_input_ids, untruncated_ids
406
+ ):
407
+ removed_text = self.tokenizer.batch_decode(
408
+ untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
409
+ )
410
+ logger.warning(
411
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
412
+ f" {self.tokenizer.model_max_length} tokens: {removed_text}"
413
+ )
414
+
415
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
416
+ attention_mask = text_inputs.attention_mask.to(device)
417
+ else:
418
+ attention_mask = None
419
+
420
+ prompt_embeds = self.text_encoder(
421
+ text_input_ids.to(device),
422
+ attention_mask=attention_mask,
423
+ )
424
+ prompt_embeds = prompt_embeds[0]
425
+
426
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
427
+
428
+ bs_embed, seq_len, _ = prompt_embeds.shape
429
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
430
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
431
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
432
+
433
+ # get unconditional embeddings for classifier free guidance
434
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
435
+ uncond_tokens: List[str]
436
+ if negative_prompt is None:
437
+ uncond_tokens = [""] * batch_size
438
+ elif type(prompt) is not type(negative_prompt):
439
+ raise TypeError(
440
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
441
+ f" {type(prompt)}."
442
+ )
443
+ elif isinstance(negative_prompt, str):
444
+ uncond_tokens = [negative_prompt]
445
+ elif batch_size != len(negative_prompt):
446
+ raise ValueError(
447
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
448
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
449
+ " the batch size of `prompt`."
450
+ )
451
+ else:
452
+ uncond_tokens = negative_prompt
453
+
454
+ max_length = prompt_embeds.shape[1]
455
+ uncond_input = self.tokenizer(
456
+ uncond_tokens,
457
+ padding="max_length",
458
+ max_length=max_length,
459
+ truncation=True,
460
+ return_tensors="pt",
461
+ )
462
+
463
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
464
+ attention_mask = uncond_input.attention_mask.to(device)
465
+ else:
466
+ attention_mask = None
467
+
468
+ negative_prompt_embeds = self.text_encoder(
469
+ uncond_input.input_ids.to(device),
470
+ attention_mask=attention_mask,
471
+ )
472
+ negative_prompt_embeds = negative_prompt_embeds[0]
473
+
474
+ if do_classifier_free_guidance:
475
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
476
+ seq_len = negative_prompt_embeds.shape[1]
477
+
478
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
479
+
480
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
481
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
482
+
483
+ # For classifier free guidance, we need to do two forward passes.
484
+ # Here we concatenate the unconditional and text embeddings into a single batch
485
+ # to avoid doing two forward passes
486
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
487
+
488
+ return prompt_embeds
489
+
490
+ def run_safety_checker(self, image, device, dtype):
491
+ if self.safety_checker is not None:
492
+ safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)
493
+ image, has_nsfw_concept = self.safety_checker(
494
+ images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
495
+ )
496
+ else:
497
+ has_nsfw_concept = None
498
+ return image, has_nsfw_concept
499
+
500
+ def decode_latents(self, latents):
501
+ latents = 1 / self.vae.config.scaling_factor * latents
502
+ image = self.vae.decode(latents).sample
503
+ image = (image / 2 + 0.5).clamp(0, 1)
504
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
505
+ image = image.cpu().permute(0, 2, 3, 1).float().numpy()
506
+ return image
507
+
508
+ def prepare_extra_step_kwargs(self, generator, eta):
509
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
510
+ # eta (Ξ·) is only used with the DDIMScheduler, it will be ignored for other schedulers.
511
+ # eta corresponds to Ξ· in DDIM paper: https://arxiv.org/abs/2010.02502
512
+ # and should be between [0, 1]
513
+
514
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
515
+ extra_step_kwargs = {}
516
+ if accepts_eta:
517
+ extra_step_kwargs["eta"] = eta
518
+
519
+ # check if the scheduler accepts generator
520
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
521
+ if accepts_generator:
522
+ extra_step_kwargs["generator"] = generator
523
+ return extra_step_kwargs
524
+
525
+ def check_inputs(
526
+ self,
527
+ prompt,
528
+ image,
529
+ mask_image,
530
+ controlnet_conditioning_image,
531
+ height,
532
+ width,
533
+ callback_steps,
534
+ negative_prompt=None,
535
+ prompt_embeds=None,
536
+ negative_prompt_embeds=None,
537
+ ):
538
+ if height % 8 != 0 or width % 8 != 0:
539
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
540
+
541
+ if (callback_steps is None) or (
542
+ callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
543
+ ):
544
+ raise ValueError(
545
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
546
+ f" {type(callback_steps)}."
547
+ )
548
+
549
+ if prompt is not None and prompt_embeds is not None:
550
+ raise ValueError(
551
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
552
+ " only forward one of the two."
553
+ )
554
+ elif prompt is None and prompt_embeds is None:
555
+ raise ValueError(
556
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
557
+ )
558
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
559
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
560
+
561
+ if negative_prompt is not None and negative_prompt_embeds is not None:
562
+ raise ValueError(
563
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
564
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
565
+ )
566
+
567
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
568
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
569
+ raise ValueError(
570
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
571
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
572
+ f" {negative_prompt_embeds.shape}."
573
+ )
574
+
575
+ controlnet_cond_image_is_pil = isinstance(controlnet_conditioning_image, PIL.Image.Image)
576
+ controlnet_cond_image_is_tensor = isinstance(controlnet_conditioning_image, torch.Tensor)
577
+ controlnet_cond_image_is_pil_list = isinstance(controlnet_conditioning_image, list) and isinstance(
578
+ controlnet_conditioning_image[0], PIL.Image.Image
579
+ )
580
+ controlnet_cond_image_is_tensor_list = isinstance(controlnet_conditioning_image, list) and isinstance(
581
+ controlnet_conditioning_image[0], torch.Tensor
582
+ )
583
+
584
+ if (
585
+ not controlnet_cond_image_is_pil
586
+ and not controlnet_cond_image_is_tensor
587
+ and not controlnet_cond_image_is_pil_list
588
+ and not controlnet_cond_image_is_tensor_list
589
+ ):
590
+ raise TypeError(
591
+ "image must be passed and be one of PIL image, torch tensor, list of PIL images, or list of torch tensors"
592
+ )
593
+
594
+ if controlnet_cond_image_is_pil:
595
+ controlnet_cond_image_batch_size = 1
596
+ elif controlnet_cond_image_is_tensor:
597
+ controlnet_cond_image_batch_size = controlnet_conditioning_image.shape[0]
598
+ elif controlnet_cond_image_is_pil_list:
599
+ controlnet_cond_image_batch_size = len(controlnet_conditioning_image)
600
+ elif controlnet_cond_image_is_tensor_list:
601
+ controlnet_cond_image_batch_size = len(controlnet_conditioning_image)
602
+
603
+ if prompt is not None and isinstance(prompt, str):
604
+ prompt_batch_size = 1
605
+ elif prompt is not None and isinstance(prompt, list):
606
+ prompt_batch_size = len(prompt)
607
+ elif prompt_embeds is not None:
608
+ prompt_batch_size = prompt_embeds.shape[0]
609
+
610
+ if controlnet_cond_image_batch_size != 1 and controlnet_cond_image_batch_size != prompt_batch_size:
611
+ raise ValueError(
612
+ f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {controlnet_cond_image_batch_size}, prompt batch size: {prompt_batch_size}"
613
+ )
614
+
615
+ if isinstance(image, torch.Tensor) and not isinstance(mask_image, torch.Tensor):
616
+ raise TypeError("if `image` is a tensor, `mask_image` must also be a tensor")
617
+
618
+ if isinstance(image, PIL.Image.Image) and not isinstance(mask_image, PIL.Image.Image):
619
+ raise TypeError("if `image` is a PIL image, `mask_image` must also be a PIL image")
620
+
621
+ if isinstance(image, torch.Tensor):
622
+ if image.ndim != 3 and image.ndim != 4:
623
+ raise ValueError("`image` must have 3 or 4 dimensions")
624
+
625
+ if mask_image.ndim != 2 and mask_image.ndim != 3 and mask_image.ndim != 4:
626
+ raise ValueError("`mask_image` must have 2, 3, or 4 dimensions")
627
+
628
+ if image.ndim == 3:
629
+ image_batch_size = 1
630
+ image_channels, image_height, image_width = image.shape
631
+ elif image.ndim == 4:
632
+ image_batch_size, image_channels, image_height, image_width = image.shape
633
+
634
+ if mask_image.ndim == 2:
635
+ mask_image_batch_size = 1
636
+ mask_image_channels = 1
637
+ mask_image_height, mask_image_width = mask_image.shape
638
+ elif mask_image.ndim == 3:
639
+ mask_image_channels = 1
640
+ mask_image_batch_size, mask_image_height, mask_image_width = mask_image.shape
641
+ elif mask_image.ndim == 4:
642
+ mask_image_batch_size, mask_image_channels, mask_image_height, mask_image_width = mask_image.shape
643
+
644
+ if image_channels != 3:
645
+ raise ValueError("`image` must have 3 channels")
646
+
647
+ if mask_image_channels != 1:
648
+ raise ValueError("`mask_image` must have 1 channel")
649
+
650
+ if image_batch_size != mask_image_batch_size:
651
+ raise ValueError("`image` and `mask_image` mush have the same batch sizes")
652
+
653
+ if image_height != mask_image_height or image_width != mask_image_width:
654
+ raise ValueError("`image` and `mask_image` must have the same height and width dimensions")
655
+
656
+ if image.min() < -1 or image.max() > 1:
657
+ raise ValueError("`image` should be in range [-1, 1]")
658
+
659
+ if mask_image.min() < 0 or mask_image.max() > 1:
660
+ raise ValueError("`mask_image` should be in range [0, 1]")
661
+ else:
662
+ mask_image_channels = 1
663
+ image_channels = 3
664
+
665
+ single_image_latent_channels = self.vae.config.latent_channels
666
+
667
+ total_latent_channels = single_image_latent_channels * 2 + mask_image_channels
668
+
669
+ if total_latent_channels != self.unet.config.in_channels:
670
+ raise ValueError(
671
+ f"The config of `pipeline.unet` expects {self.unet.config.in_channels} but received"
672
+ f" non inpainting latent channels: {single_image_latent_channels},"
673
+ f" mask channels: {mask_image_channels}, and masked image channels: {single_image_latent_channels}."
674
+ f" Please verify the config of `pipeline.unet` and the `mask_image` and `image` inputs."
675
+ )
676
+
677
+ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
678
+ shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
679
+ if isinstance(generator, list) and len(generator) != batch_size:
680
+ raise ValueError(
681
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
682
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
683
+ )
684
+
685
+ if latents is None:
686
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
687
+ else:
688
+ latents = latents.to(device)
689
+
690
+ # scale the initial noise by the standard deviation required by the scheduler
691
+ latents = latents * self.scheduler.init_noise_sigma
692
+
693
+ return latents
694
+
695
+ def prepare_mask_latents(self, mask_image, batch_size, height, width, dtype, device, do_classifier_free_guidance):
696
+ # resize the mask to latents shape as we concatenate the mask to the latents
697
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
698
+ # and half precision
699
+ mask_image = F.interpolate(mask_image, size=(height // self.vae_scale_factor, width // self.vae_scale_factor))
700
+ mask_image = mask_image.to(device=device, dtype=dtype)
701
+
702
+ # duplicate mask for each generation per prompt, using mps friendly method
703
+ if mask_image.shape[0] < batch_size:
704
+ if not batch_size % mask_image.shape[0] == 0:
705
+ raise ValueError(
706
+ "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
707
+ f" a total batch size of {batch_size}, but {mask_image.shape[0]} masks were passed. Make sure the number"
708
+ " of masks that you pass is divisible by the total requested batch size."
709
+ )
710
+ mask_image = mask_image.repeat(batch_size // mask_image.shape[0], 1, 1, 1)
711
+
712
+ mask_image = torch.cat([mask_image] * 2) if do_classifier_free_guidance else mask_image
713
+
714
+ mask_image_latents = mask_image
715
+
716
+ return mask_image_latents
717
+
718
+ def prepare_masked_image_latents(
719
+ self, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance
720
+ ):
721
+ masked_image = masked_image.to(device=device, dtype=dtype)
722
+
723
+ # encode the mask image into latents space so we can concatenate it to the latents
724
+ if isinstance(generator, list):
725
+ masked_image_latents = [
726
+ self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(generator=generator[i])
727
+ for i in range(batch_size)
728
+ ]
729
+ masked_image_latents = torch.cat(masked_image_latents, dim=0)
730
+ else:
731
+ masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator)
732
+ masked_image_latents = self.vae.config.scaling_factor * masked_image_latents
733
+
734
+ # duplicate masked_image_latents for each generation per prompt, using mps friendly method
735
+ if masked_image_latents.shape[0] < batch_size:
736
+ if not batch_size % masked_image_latents.shape[0] == 0:
737
+ raise ValueError(
738
+ "The passed images and the required batch size don't match. Images are supposed to be duplicated"
739
+ f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
740
+ " Make sure the number of images that you pass is divisible by the total requested batch size."
741
+ )
742
+ masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)
743
+
744
+ masked_image_latents = (
745
+ torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents
746
+ )
747
+
748
+ # aligning device to prevent device errors when concating it with the latent model input
749
+ masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
750
+ return masked_image_latents
751
+
752
+ def _default_height_width(self, height, width, image):
753
+ if isinstance(image, list):
754
+ image = image[0]
755
+
756
+ if height is None:
757
+ if isinstance(image, PIL.Image.Image):
758
+ height = image.height
759
+ elif isinstance(image, torch.Tensor):
760
+ height = image.shape[3]
761
+
762
+ height = (height // 8) * 8 # round down to nearest multiple of 8
763
+
764
+ if width is None:
765
+ if isinstance(image, PIL.Image.Image):
766
+ width = image.width
767
+ elif isinstance(image, torch.Tensor):
768
+ width = image.shape[2]
769
+
770
+ width = (width // 8) * 8 # round down to nearest multiple of 8
771
+
772
+ return height, width
773
+
774
+ @torch.no_grad()
775
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
776
+ def __call__(
777
+ self,
778
+ prompt: Union[str, List[str]] = None,
779
+ image: Union[torch.Tensor, PIL.Image.Image] = None,
780
+ mask_image: Union[torch.Tensor, PIL.Image.Image] = None,
781
+ controlnet_conditioning_image: Union[
782
+ torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]
783
+ ] = None,
784
+ height: Optional[int] = None,
785
+ width: Optional[int] = None,
786
+ num_inference_steps: int = 50,
787
+ guidance_scale: float = 7.5,
788
+ negative_prompt: Optional[Union[str, List[str]]] = None,
789
+ num_images_per_prompt: Optional[int] = 1,
790
+ eta: float = 0.0,
791
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
792
+ latents: Optional[torch.FloatTensor] = None,
793
+ prompt_embeds: Optional[torch.FloatTensor] = None,
794
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
795
+ output_type: Optional[str] = "pil",
796
+ return_dict: bool = True,
797
+ callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
798
+ callback_steps: int = 1,
799
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
800
+ controlnet_conditioning_scale: float = 1.0,
801
+ ):
802
+ r"""
803
+ Function invoked when calling the pipeline for generation.
804
+
805
+ Args:
806
+ prompt (`str` or `List[str]`, *optional*):
807
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
808
+ instead.
809
+ image (`torch.Tensor` or `PIL.Image.Image`):
810
+ `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will
811
+ be masked out with `mask_image` and repainted according to `prompt`.
812
+ mask_image (`torch.Tensor` or `PIL.Image.Image`):
813
+ `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be
814
+ repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted
815
+ to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)
816
+ instead of 3, so the expected shape would be `(B, H, W, 1)`.
817
+ controlnet_conditioning_image (`torch.FloatTensor`, `PIL.Image.Image`, `List[torch.FloatTensor]` or `List[PIL.Image.Image]`):
818
+ The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If
819
+ the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. PIL.Image.Image` can
820
+ also be accepted as an image. The control image is automatically resized to fit the output image.
821
+ height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
822
+ The height in pixels of the generated image.
823
+ width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
824
+ The width in pixels of the generated image.
825
+ num_inference_steps (`int`, *optional*, defaults to 50):
826
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
827
+ expense of slower inference.
828
+ guidance_scale (`float`, *optional*, defaults to 7.5):
829
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
830
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
831
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
832
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
833
+ usually at the expense of lower image quality.
834
+ negative_prompt (`str` or `List[str]`, *optional*):
835
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead.
836
+ Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
837
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
838
+ The number of images to generate per prompt.
839
+ eta (`float`, *optional*, defaults to 0.0):
840
+ Corresponds to parameter eta (Ξ·) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
841
+ [`schedulers.DDIMScheduler`], will be ignored for others.
842
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
843
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
844
+ to make generation deterministic.
845
+ latents (`torch.FloatTensor`, *optional*):
846
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
847
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
848
+ tensor will ge generated by sampling using the supplied random `generator`.
849
+ prompt_embeds (`torch.FloatTensor`, *optional*):
850
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
851
+ provided, text embeddings will be generated from `prompt` input argument.
852
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
853
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
854
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
855
+ argument.
856
+ output_type (`str`, *optional*, defaults to `"pil"`):
857
+ The output format of the generate image. Choose between
858
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
859
+ return_dict (`bool`, *optional*, defaults to `True`):
860
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
861
+ plain tuple.
862
+ callback (`Callable`, *optional*):
863
+ A function that will be called every `callback_steps` steps during inference. The function will be
864
+ called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
865
+ callback_steps (`int`, *optional*, defaults to 1):
866
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
867
+ called at every step.
868
+ cross_attention_kwargs (`dict`, *optional*):
869
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
870
+ `self.processor` in
871
+ [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
872
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
873
+ The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
874
+ to the residual in the original unet.
875
+
876
+ Examples:
877
+
878
+ Returns:
879
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
880
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
881
+ When returning a tuple, the first element is a list with the generated images, and the second element is a
882
+ list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
883
+ (nsfw) content, according to the `safety_checker`.
884
+ """
885
+ # 0. Default height and width to unet
886
+ height, width = self._default_height_width(height, width, controlnet_conditioning_image)
887
+
888
+ # 1. Check inputs. Raise error if not correct
889
+ self.check_inputs(
890
+ prompt,
891
+ image,
892
+ mask_image,
893
+ controlnet_conditioning_image,
894
+ height,
895
+ width,
896
+ callback_steps,
897
+ negative_prompt,
898
+ prompt_embeds,
899
+ negative_prompt_embeds,
900
+ )
901
+
902
+ # 2. Define call parameters
903
+ if prompt is not None and isinstance(prompt, str):
904
+ batch_size = 1
905
+ elif prompt is not None and isinstance(prompt, list):
906
+ batch_size = len(prompt)
907
+ else:
908
+ batch_size = prompt_embeds.shape[0]
909
+
910
+ device = self._execution_device
911
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
912
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
913
+ # corresponds to doing no classifier free guidance.
914
+ do_classifier_free_guidance = guidance_scale > 1.0
915
+
916
+ # 3. Encode input prompt
917
+ prompt_embeds = self._encode_prompt(
918
+ prompt,
919
+ device,
920
+ num_images_per_prompt,
921
+ do_classifier_free_guidance,
922
+ negative_prompt,
923
+ prompt_embeds=prompt_embeds,
924
+ negative_prompt_embeds=negative_prompt_embeds,
925
+ )
926
+
927
+ # 4. Prepare mask, image, and controlnet_conditioning_image
928
+ image = prepare_image(image)
929
+
930
+ mask_image = prepare_mask_image(mask_image)
931
+
932
+ controlnet_conditioning_image = prepare_controlnet_conditioning_image(
933
+ controlnet_conditioning_image,
934
+ width,
935
+ height,
936
+ batch_size * num_images_per_prompt,
937
+ num_images_per_prompt,
938
+ device,
939
+ self.controlnet.dtype,
940
+ )
941
+
942
+ masked_image = image * (mask_image < 0.5)
943
+
944
+ # 5. Prepare timesteps
945
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
946
+ timesteps = self.scheduler.timesteps
947
+
948
+ # 6. Prepare latent variables
949
+ num_channels_latents = self.vae.config.latent_channels
950
+ latents = self.prepare_latents(
951
+ batch_size * num_images_per_prompt,
952
+ num_channels_latents,
953
+ height,
954
+ width,
955
+ prompt_embeds.dtype,
956
+ device,
957
+ generator,
958
+ latents,
959
+ )
960
+
961
+ mask_image_latents = self.prepare_mask_latents(
962
+ mask_image,
963
+ batch_size * num_images_per_prompt,
964
+ height,
965
+ width,
966
+ prompt_embeds.dtype,
967
+ device,
968
+ do_classifier_free_guidance,
969
+ )
970
+
971
+ masked_image_latents = self.prepare_masked_image_latents(
972
+ masked_image,
973
+ batch_size * num_images_per_prompt,
974
+ height,
975
+ width,
976
+ prompt_embeds.dtype,
977
+ device,
978
+ generator,
979
+ do_classifier_free_guidance,
980
+ )
981
+
982
+ if do_classifier_free_guidance:
983
+ controlnet_conditioning_image = torch.cat([controlnet_conditioning_image] * 2)
984
+
985
+ # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
986
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
987
+
988
+ # 8. Denoising loop
989
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
990
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
991
+ for i, t in enumerate(timesteps):
992
+ # expand the latents if we are doing classifier free guidance
993
+ non_inpainting_latent_model_input = (
994
+ torch.cat([latents] * 2) if do_classifier_free_guidance else latents
995
+ )
996
+
997
+ non_inpainting_latent_model_input = self.scheduler.scale_model_input(
998
+ non_inpainting_latent_model_input, t
999
+ )
1000
+
1001
+ inpainting_latent_model_input = torch.cat(
1002
+ [non_inpainting_latent_model_input, mask_image_latents, masked_image_latents], dim=1
1003
+ )
1004
+
1005
+ down_block_res_samples, mid_block_res_sample = self.controlnet(
1006
+ non_inpainting_latent_model_input,
1007
+ t,
1008
+ encoder_hidden_states=prompt_embeds,
1009
+ controlnet_cond=controlnet_conditioning_image,
1010
+ return_dict=False,
1011
+ )
1012
+
1013
+ down_block_res_samples = [
1014
+ down_block_res_sample * controlnet_conditioning_scale
1015
+ for down_block_res_sample in down_block_res_samples
1016
+ ]
1017
+ mid_block_res_sample *= controlnet_conditioning_scale
1018
+
1019
+ # predict the noise residual
1020
+ noise_pred = self.unet(
1021
+ inpainting_latent_model_input,
1022
+ t,
1023
+ encoder_hidden_states=prompt_embeds,
1024
+ cross_attention_kwargs=cross_attention_kwargs,
1025
+ down_block_additional_residuals=down_block_res_samples,
1026
+ mid_block_additional_residual=mid_block_res_sample,
1027
+ ).sample
1028
+
1029
+ # perform guidance
1030
+ if do_classifier_free_guidance:
1031
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
1032
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
1033
+
1034
+ # compute the previous noisy sample x_t -> x_t-1
1035
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
1036
+
1037
+ # call the callback, if provided
1038
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
1039
+ progress_bar.update()
1040
+ if callback is not None and i % callback_steps == 0:
1041
+ callback(i, t, latents)
1042
+
1043
+ # If we do sequential model offloading, let's offload unet and controlnet
1044
+ # manually for max memory savings
1045
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
1046
+ self.unet.to("cpu")
1047
+ self.controlnet.to("cpu")
1048
+ torch.cuda.empty_cache()
1049
+
1050
+ if output_type == "latent":
1051
+ image = latents
1052
+ has_nsfw_concept = None
1053
+ elif output_type == "pil":
1054
+ # 8. Post-processing
1055
+ image = self.decode_latents(latents)
1056
+
1057
+ # 9. Run safety checker
1058
+ image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
1059
+
1060
+ # 10. Convert to PIL
1061
+ image = self.numpy_to_pil(image)
1062
+ else:
1063
+ # 8. Post-processing
1064
+ image = self.decode_latents(latents)
1065
+
1066
+ # 9. Run safety checker
1067
+ image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
1068
+
1069
+ # Offload last model to CPU
1070
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
1071
+ self.final_offload_hook.offload()
1072
+
1073
+ if not return_dict:
1074
+ return (image, has_nsfw_concept)
1075
+
1076
+ return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
utils/stable_diffusion_controlnet_inpaint.py ADDED
@@ -0,0 +1,1076 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Inspired by: https://github.com/haofanwang/ControlNet-for-Diffusers/
2
+ # From https://raw.githubusercontent.com/huggingface/diffusers/53377ef83c6446033f3ee506e3ef718db817b293/examples/community/stable_diffusion_controlnet_inpaint.py
3
+ import inspect
4
+ from typing import Any, Callable, Dict, List, Optional, Union
5
+
6
+ import numpy as np
7
+ import PIL.Image
8
+ import torch
9
+ import torch.nn.functional as F
10
+ from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
11
+
12
+ from diffusers import AutoencoderKL, ControlNetModel, DiffusionPipeline, UNet2DConditionModel, logging
13
+ from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker
14
+ from diffusers.schedulers import KarrasDiffusionSchedulers
15
+ from diffusers.utils import (
16
+ PIL_INTERPOLATION,
17
+ is_accelerate_available,
18
+ is_accelerate_version,
19
+ randn_tensor,
20
+ replace_example_docstring,
21
+ )
22
+
23
+
24
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
25
+
26
+ EXAMPLE_DOC_STRING = """
27
+ Examples:
28
+ ```py
29
+ >>> import numpy as np
30
+ >>> import torch
31
+ >>> from PIL import Image
32
+ >>> from stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
33
+
34
+ >>> from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
35
+ >>> from diffusers import ControlNetModel, UniPCMultistepScheduler
36
+ >>> from diffusers.utils import load_image
37
+
38
+ >>> def ade_palette():
39
+ return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
40
+ [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
41
+ [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
42
+ [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
43
+ [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
44
+ [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
45
+ [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
46
+ [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
47
+ [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
48
+ [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
49
+ [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
50
+ [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
51
+ [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
52
+ [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
53
+ [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
54
+ [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
55
+ [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
56
+ [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
57
+ [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
58
+ [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
59
+ [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
60
+ [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
61
+ [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
62
+ [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
63
+ [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
64
+ [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
65
+ [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
66
+ [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
67
+ [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
68
+ [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
69
+ [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
70
+ [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
71
+ [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
72
+ [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
73
+ [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
74
+ [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
75
+ [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
76
+ [102, 255, 0], [92, 0, 255]]
77
+
78
+ >>> image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
79
+ >>> image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
80
+
81
+ >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-seg", torch_dtype=torch.float16)
82
+
83
+ >>> pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
84
+ "runwayml/stable-diffusion-inpainting", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
85
+ )
86
+
87
+ >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
88
+ >>> pipe.enable_xformers_memory_efficient_attention()
89
+ >>> pipe.enable_model_cpu_offload()
90
+
91
+ >>> def image_to_seg(image):
92
+ pixel_values = image_processor(image, return_tensors="pt").pixel_values
93
+ with torch.no_grad():
94
+ outputs = image_segmentor(pixel_values)
95
+ seg = image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
96
+ color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
97
+ palette = np.array(ade_palette())
98
+ for label, color in enumerate(palette):
99
+ color_seg[seg == label, :] = color
100
+ color_seg = color_seg.astype(np.uint8)
101
+ seg_image = Image.fromarray(color_seg)
102
+ return seg_image
103
+
104
+ >>> image = load_image(
105
+ "https://github.com/CompVis/latent-diffusion/raw/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
106
+ )
107
+
108
+ >>> mask_image = load_image(
109
+ "https://github.com/CompVis/latent-diffusion/raw/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
110
+ )
111
+
112
+ >>> controlnet_conditioning_image = image_to_seg(image)
113
+
114
+ >>> image = pipe(
115
+ "Face of a yellow cat, high resolution, sitting on a park bench",
116
+ image,
117
+ mask_image,
118
+ controlnet_conditioning_image,
119
+ num_inference_steps=20,
120
+ ).images[0]
121
+
122
+ >>> image.save("out.png")
123
+ ```
124
+ """
125
+
126
+
127
+ def prepare_image(image):
128
+ if isinstance(image, torch.Tensor):
129
+ # Batch single image
130
+ if image.ndim == 3:
131
+ image = image.unsqueeze(0)
132
+
133
+ image = image.to(dtype=torch.float32)
134
+ else:
135
+ # preprocess image
136
+ if isinstance(image, (PIL.Image.Image, np.ndarray)):
137
+ image = [image]
138
+
139
+ if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):
140
+ image = [np.array(i.convert("RGB"))[None, :] for i in image]
141
+ image = np.concatenate(image, axis=0)
142
+ elif isinstance(image, list) and isinstance(image[0], np.ndarray):
143
+ image = np.concatenate([i[None, :] for i in image], axis=0)
144
+
145
+ image = image.transpose(0, 3, 1, 2)
146
+ image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0
147
+
148
+ return image
149
+
150
+
151
+ def prepare_mask_image(mask_image):
152
+ if isinstance(mask_image, torch.Tensor):
153
+ if mask_image.ndim == 2:
154
+ # Batch and add channel dim for single mask
155
+ mask_image = mask_image.unsqueeze(0).unsqueeze(0)
156
+ elif mask_image.ndim == 3 and mask_image.shape[0] == 1:
157
+ # Single mask, the 0'th dimension is considered to be
158
+ # the existing batch size of 1
159
+ mask_image = mask_image.unsqueeze(0)
160
+ elif mask_image.ndim == 3 and mask_image.shape[0] != 1:
161
+ # Batch of mask, the 0'th dimension is considered to be
162
+ # the batching dimension
163
+ mask_image = mask_image.unsqueeze(1)
164
+
165
+ # Binarize mask
166
+ mask_image[mask_image < 0.5] = 0
167
+ mask_image[mask_image >= 0.5] = 1
168
+ else:
169
+ # preprocess mask
170
+ if isinstance(mask_image, (PIL.Image.Image, np.ndarray)):
171
+ mask_image = [mask_image]
172
+
173
+ if isinstance(mask_image, list) and isinstance(mask_image[0], PIL.Image.Image):
174
+ mask_image = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask_image], axis=0)
175
+ mask_image = mask_image.astype(np.float32) / 255.0
176
+ elif isinstance(mask_image, list) and isinstance(mask_image[0], np.ndarray):
177
+ mask_image = np.concatenate([m[None, None, :] for m in mask_image], axis=0)
178
+
179
+ mask_image[mask_image < 0.5] = 0
180
+ mask_image[mask_image >= 0.5] = 1
181
+ mask_image = torch.from_numpy(mask_image)
182
+
183
+ return mask_image
184
+
185
+
186
+ def prepare_controlnet_conditioning_image(
187
+ controlnet_conditioning_image, width, height, batch_size, num_images_per_prompt, device, dtype
188
+ ):
189
+ if not isinstance(controlnet_conditioning_image, torch.Tensor):
190
+ if isinstance(controlnet_conditioning_image, PIL.Image.Image):
191
+ controlnet_conditioning_image = [controlnet_conditioning_image]
192
+
193
+ if isinstance(controlnet_conditioning_image[0], PIL.Image.Image):
194
+ controlnet_conditioning_image = [
195
+ np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"]))[None, :]
196
+ for i in controlnet_conditioning_image
197
+ ]
198
+ controlnet_conditioning_image = np.concatenate(controlnet_conditioning_image, axis=0)
199
+ controlnet_conditioning_image = np.array(controlnet_conditioning_image).astype(np.float32) / 255.0
200
+ controlnet_conditioning_image = controlnet_conditioning_image.transpose(0, 3, 1, 2)
201
+ controlnet_conditioning_image = torch.from_numpy(controlnet_conditioning_image)
202
+ elif isinstance(controlnet_conditioning_image[0], torch.Tensor):
203
+ controlnet_conditioning_image = torch.cat(controlnet_conditioning_image, dim=0)
204
+
205
+ image_batch_size = controlnet_conditioning_image.shape[0]
206
+
207
+ if image_batch_size == 1:
208
+ repeat_by = batch_size
209
+ else:
210
+ # image batch size is the same as prompt batch size
211
+ repeat_by = num_images_per_prompt
212
+
213
+ controlnet_conditioning_image = controlnet_conditioning_image.repeat_interleave(repeat_by, dim=0)
214
+
215
+ controlnet_conditioning_image = controlnet_conditioning_image.to(device=device, dtype=dtype)
216
+
217
+ return controlnet_conditioning_image
218
+
219
+
220
+ class StableDiffusionControlNetInpaintPipeline(DiffusionPipeline):
221
+ """
222
+ Inspired by: https://github.com/haofanwang/ControlNet-for-Diffusers/
223
+ """
224
+
225
+ _optional_components = ["safety_checker", "feature_extractor"]
226
+
227
+ def __init__(
228
+ self,
229
+ vae: AutoencoderKL,
230
+ text_encoder: CLIPTextModel,
231
+ tokenizer: CLIPTokenizer,
232
+ unet: UNet2DConditionModel,
233
+ controlnet: ControlNetModel,
234
+ scheduler: KarrasDiffusionSchedulers,
235
+ safety_checker: StableDiffusionSafetyChecker,
236
+ feature_extractor: CLIPImageProcessor,
237
+ requires_safety_checker: bool = True,
238
+ ):
239
+ super().__init__()
240
+
241
+ if safety_checker is None and requires_safety_checker:
242
+ logger.warning(
243
+ f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
244
+ " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
245
+ " results in services or applications open to the public. Both the diffusers team and Hugging Face"
246
+ " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
247
+ " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
248
+ " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
249
+ )
250
+
251
+ if safety_checker is not None and feature_extractor is None:
252
+ raise ValueError(
253
+ "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
254
+ " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
255
+ )
256
+
257
+ self.register_modules(
258
+ vae=vae,
259
+ text_encoder=text_encoder,
260
+ tokenizer=tokenizer,
261
+ unet=unet,
262
+ controlnet=controlnet,
263
+ scheduler=scheduler,
264
+ safety_checker=safety_checker,
265
+ feature_extractor=feature_extractor,
266
+ )
267
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
268
+ self.register_to_config(requires_safety_checker=requires_safety_checker)
269
+
270
+ def enable_vae_slicing(self):
271
+ r"""
272
+ Enable sliced VAE decoding.
273
+
274
+ When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
275
+ steps. This is useful to save some memory and allow larger batch sizes.
276
+ """
277
+ self.vae.enable_slicing()
278
+
279
+ def disable_vae_slicing(self):
280
+ r"""
281
+ Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
282
+ computing decoding in one step.
283
+ """
284
+ self.vae.disable_slicing()
285
+
286
+ def enable_sequential_cpu_offload(self, gpu_id=0):
287
+ r"""
288
+ Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
289
+ text_encoder, vae, controlnet, and safety checker have their state dicts saved to CPU and then are moved to a
290
+ `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
291
+ Note that offloading happens on a submodule basis. Memory savings are higher than with
292
+ `enable_model_cpu_offload`, but performance is lower.
293
+ """
294
+ if is_accelerate_available():
295
+ from accelerate import cpu_offload
296
+ else:
297
+ raise ImportError("Please install accelerate via `pip install accelerate`")
298
+
299
+ device = torch.device(f"cuda:{gpu_id}")
300
+
301
+ for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.controlnet]:
302
+ cpu_offload(cpu_offloaded_model, device)
303
+
304
+ if self.safety_checker is not None:
305
+ cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True)
306
+
307
+ def enable_model_cpu_offload(self, gpu_id=0):
308
+ r"""
309
+ Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
310
+ to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
311
+ method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
312
+ `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
313
+ """
314
+ if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
315
+ from accelerate import cpu_offload_with_hook
316
+ else:
317
+ raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")
318
+
319
+ device = torch.device(f"cuda:{gpu_id}")
320
+
321
+ hook = None
322
+ for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
323
+ _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
324
+
325
+ if self.safety_checker is not None:
326
+ # the safety checker can offload the vae again
327
+ _, hook = cpu_offload_with_hook(self.safety_checker, device, prev_module_hook=hook)
328
+
329
+ # control net hook has be manually offloaded as it alternates with unet
330
+ cpu_offload_with_hook(self.controlnet, device)
331
+
332
+ # We'll offload the last model manually.
333
+ self.final_offload_hook = hook
334
+
335
+ @property
336
+ def _execution_device(self):
337
+ r"""
338
+ Returns the device on which the pipeline's models will be executed. After calling
339
+ `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
340
+ hooks.
341
+ """
342
+ if not hasattr(self.unet, "_hf_hook"):
343
+ return self.device
344
+ for module in self.unet.modules():
345
+ if (
346
+ hasattr(module, "_hf_hook")
347
+ and hasattr(module._hf_hook, "execution_device")
348
+ and module._hf_hook.execution_device is not None
349
+ ):
350
+ return torch.device(module._hf_hook.execution_device)
351
+ return self.device
352
+
353
+ def _encode_prompt(
354
+ self,
355
+ prompt,
356
+ device,
357
+ num_images_per_prompt,
358
+ do_classifier_free_guidance,
359
+ negative_prompt=None,
360
+ prompt_embeds: Optional[torch.FloatTensor] = None,
361
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
362
+ ):
363
+ r"""
364
+ Encodes the prompt into text encoder hidden states.
365
+
366
+ Args:
367
+ prompt (`str` or `List[str]`, *optional*):
368
+ prompt to be encoded
369
+ device: (`torch.device`):
370
+ torch device
371
+ num_images_per_prompt (`int`):
372
+ number of images that should be generated per prompt
373
+ do_classifier_free_guidance (`bool`):
374
+ whether to use classifier free guidance or not
375
+ negative_prompt (`str` or `List[str]`, *optional*):
376
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead.
377
+ Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
378
+ prompt_embeds (`torch.FloatTensor`, *optional*):
379
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
380
+ provided, text embeddings will be generated from `prompt` input argument.
381
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
382
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
383
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
384
+ argument.
385
+ """
386
+ if prompt is not None and isinstance(prompt, str):
387
+ batch_size = 1
388
+ elif prompt is not None and isinstance(prompt, list):
389
+ batch_size = len(prompt)
390
+ else:
391
+ batch_size = prompt_embeds.shape[0]
392
+
393
+ if prompt_embeds is None:
394
+ text_inputs = self.tokenizer(
395
+ prompt,
396
+ padding="max_length",
397
+ max_length=self.tokenizer.model_max_length,
398
+ truncation=True,
399
+ return_tensors="pt",
400
+ )
401
+ text_input_ids = text_inputs.input_ids
402
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
403
+
404
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
405
+ text_input_ids, untruncated_ids
406
+ ):
407
+ removed_text = self.tokenizer.batch_decode(
408
+ untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
409
+ )
410
+ logger.warning(
411
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
412
+ f" {self.tokenizer.model_max_length} tokens: {removed_text}"
413
+ )
414
+
415
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
416
+ attention_mask = text_inputs.attention_mask.to(device)
417
+ else:
418
+ attention_mask = None
419
+
420
+ prompt_embeds = self.text_encoder(
421
+ text_input_ids.to(device),
422
+ attention_mask=attention_mask,
423
+ )
424
+ prompt_embeds = prompt_embeds[0]
425
+
426
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
427
+
428
+ bs_embed, seq_len, _ = prompt_embeds.shape
429
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
430
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
431
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
432
+
433
+ # get unconditional embeddings for classifier free guidance
434
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
435
+ uncond_tokens: List[str]
436
+ if negative_prompt is None:
437
+ uncond_tokens = [""] * batch_size
438
+ elif type(prompt) is not type(negative_prompt):
439
+ raise TypeError(
440
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
441
+ f" {type(prompt)}."
442
+ )
443
+ elif isinstance(negative_prompt, str):
444
+ uncond_tokens = [negative_prompt]
445
+ elif batch_size != len(negative_prompt):
446
+ raise ValueError(
447
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
448
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
449
+ " the batch size of `prompt`."
450
+ )
451
+ else:
452
+ uncond_tokens = negative_prompt
453
+
454
+ max_length = prompt_embeds.shape[1]
455
+ uncond_input = self.tokenizer(
456
+ uncond_tokens,
457
+ padding="max_length",
458
+ max_length=max_length,
459
+ truncation=True,
460
+ return_tensors="pt",
461
+ )
462
+
463
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
464
+ attention_mask = uncond_input.attention_mask.to(device)
465
+ else:
466
+ attention_mask = None
467
+
468
+ negative_prompt_embeds = self.text_encoder(
469
+ uncond_input.input_ids.to(device),
470
+ attention_mask=attention_mask,
471
+ )
472
+ negative_prompt_embeds = negative_prompt_embeds[0]
473
+
474
+ if do_classifier_free_guidance:
475
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
476
+ seq_len = negative_prompt_embeds.shape[1]
477
+
478
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
479
+
480
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
481
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
482
+
483
+ # For classifier free guidance, we need to do two forward passes.
484
+ # Here we concatenate the unconditional and text embeddings into a single batch
485
+ # to avoid doing two forward passes
486
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
487
+
488
+ return prompt_embeds
489
+
490
+ def run_safety_checker(self, image, device, dtype):
491
+ if self.safety_checker is not None:
492
+ safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)
493
+ image, has_nsfw_concept = self.safety_checker(
494
+ images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
495
+ )
496
+ else:
497
+ has_nsfw_concept = None
498
+ return image, has_nsfw_concept
499
+
500
+ def decode_latents(self, latents):
501
+ latents = 1 / self.vae.config.scaling_factor * latents
502
+ image = self.vae.decode(latents).sample
503
+ image = (image / 2 + 0.5).clamp(0, 1)
504
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
505
+ image = image.cpu().permute(0, 2, 3, 1).float().numpy()
506
+ return image
507
+
508
+ def prepare_extra_step_kwargs(self, generator, eta):
509
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
510
+ # eta (Ξ·) is only used with the DDIMScheduler, it will be ignored for other schedulers.
511
+ # eta corresponds to Ξ· in DDIM paper: https://arxiv.org/abs/2010.02502
512
+ # and should be between [0, 1]
513
+
514
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
515
+ extra_step_kwargs = {}
516
+ if accepts_eta:
517
+ extra_step_kwargs["eta"] = eta
518
+
519
+ # check if the scheduler accepts generator
520
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
521
+ if accepts_generator:
522
+ extra_step_kwargs["generator"] = generator
523
+ return extra_step_kwargs
524
+
525
+ def check_inputs(
526
+ self,
527
+ prompt,
528
+ image,
529
+ mask_image,
530
+ controlnet_conditioning_image,
531
+ height,
532
+ width,
533
+ callback_steps,
534
+ negative_prompt=None,
535
+ prompt_embeds=None,
536
+ negative_prompt_embeds=None,
537
+ ):
538
+ if height % 8 != 0 or width % 8 != 0:
539
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
540
+
541
+ if (callback_steps is None) or (
542
+ callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
543
+ ):
544
+ raise ValueError(
545
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
546
+ f" {type(callback_steps)}."
547
+ )
548
+
549
+ if prompt is not None and prompt_embeds is not None:
550
+ raise ValueError(
551
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
552
+ " only forward one of the two."
553
+ )
554
+ elif prompt is None and prompt_embeds is None:
555
+ raise ValueError(
556
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
557
+ )
558
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
559
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
560
+
561
+ if negative_prompt is not None and negative_prompt_embeds is not None:
562
+ raise ValueError(
563
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
564
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
565
+ )
566
+
567
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
568
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
569
+ raise ValueError(
570
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
571
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
572
+ f" {negative_prompt_embeds.shape}."
573
+ )
574
+
575
+ controlnet_cond_image_is_pil = isinstance(controlnet_conditioning_image, PIL.Image.Image)
576
+ controlnet_cond_image_is_tensor = isinstance(controlnet_conditioning_image, torch.Tensor)
577
+ controlnet_cond_image_is_pil_list = isinstance(controlnet_conditioning_image, list) and isinstance(
578
+ controlnet_conditioning_image[0], PIL.Image.Image
579
+ )
580
+ controlnet_cond_image_is_tensor_list = isinstance(controlnet_conditioning_image, list) and isinstance(
581
+ controlnet_conditioning_image[0], torch.Tensor
582
+ )
583
+
584
+ if (
585
+ not controlnet_cond_image_is_pil
586
+ and not controlnet_cond_image_is_tensor
587
+ and not controlnet_cond_image_is_pil_list
588
+ and not controlnet_cond_image_is_tensor_list
589
+ ):
590
+ raise TypeError(
591
+ "image must be passed and be one of PIL image, torch tensor, list of PIL images, or list of torch tensors"
592
+ )
593
+
594
+ if controlnet_cond_image_is_pil:
595
+ controlnet_cond_image_batch_size = 1
596
+ elif controlnet_cond_image_is_tensor:
597
+ controlnet_cond_image_batch_size = controlnet_conditioning_image.shape[0]
598
+ elif controlnet_cond_image_is_pil_list:
599
+ controlnet_cond_image_batch_size = len(controlnet_conditioning_image)
600
+ elif controlnet_cond_image_is_tensor_list:
601
+ controlnet_cond_image_batch_size = len(controlnet_conditioning_image)
602
+
603
+ if prompt is not None and isinstance(prompt, str):
604
+ prompt_batch_size = 1
605
+ elif prompt is not None and isinstance(prompt, list):
606
+ prompt_batch_size = len(prompt)
607
+ elif prompt_embeds is not None:
608
+ prompt_batch_size = prompt_embeds.shape[0]
609
+
610
+ if controlnet_cond_image_batch_size != 1 and controlnet_cond_image_batch_size != prompt_batch_size:
611
+ raise ValueError(
612
+ f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {controlnet_cond_image_batch_size}, prompt batch size: {prompt_batch_size}"
613
+ )
614
+
615
+ if isinstance(image, torch.Tensor) and not isinstance(mask_image, torch.Tensor):
616
+ raise TypeError("if `image` is a tensor, `mask_image` must also be a tensor")
617
+
618
+ if isinstance(image, PIL.Image.Image) and not isinstance(mask_image, PIL.Image.Image):
619
+ raise TypeError("if `image` is a PIL image, `mask_image` must also be a PIL image")
620
+
621
+ if isinstance(image, torch.Tensor):
622
+ if image.ndim != 3 and image.ndim != 4:
623
+ raise ValueError("`image` must have 3 or 4 dimensions")
624
+
625
+ if mask_image.ndim != 2 and mask_image.ndim != 3 and mask_image.ndim != 4:
626
+ raise ValueError("`mask_image` must have 2, 3, or 4 dimensions")
627
+
628
+ if image.ndim == 3:
629
+ image_batch_size = 1
630
+ image_channels, image_height, image_width = image.shape
631
+ elif image.ndim == 4:
632
+ image_batch_size, image_channels, image_height, image_width = image.shape
633
+
634
+ if mask_image.ndim == 2:
635
+ mask_image_batch_size = 1
636
+ mask_image_channels = 1
637
+ mask_image_height, mask_image_width = mask_image.shape
638
+ elif mask_image.ndim == 3:
639
+ mask_image_channels = 1
640
+ mask_image_batch_size, mask_image_height, mask_image_width = mask_image.shape
641
+ elif mask_image.ndim == 4:
642
+ mask_image_batch_size, mask_image_channels, mask_image_height, mask_image_width = mask_image.shape
643
+
644
+ if image_channels != 3:
645
+ raise ValueError("`image` must have 3 channels")
646
+
647
+ if mask_image_channels != 1:
648
+ raise ValueError("`mask_image` must have 1 channel")
649
+
650
+ if image_batch_size != mask_image_batch_size:
651
+ raise ValueError("`image` and `mask_image` mush have the same batch sizes")
652
+
653
+ if image_height != mask_image_height or image_width != mask_image_width:
654
+ raise ValueError("`image` and `mask_image` must have the same height and width dimensions")
655
+
656
+ if image.min() < -1 or image.max() > 1:
657
+ raise ValueError("`image` should be in range [-1, 1]")
658
+
659
+ if mask_image.min() < 0 or mask_image.max() > 1:
660
+ raise ValueError("`mask_image` should be in range [0, 1]")
661
+ else:
662
+ mask_image_channels = 1
663
+ image_channels = 3
664
+
665
+ single_image_latent_channels = self.vae.config.latent_channels
666
+
667
+ total_latent_channels = single_image_latent_channels * 2 + mask_image_channels
668
+
669
+ if total_latent_channels != self.unet.config.in_channels:
670
+ raise ValueError(
671
+ f"The config of `pipeline.unet` expects {self.unet.config.in_channels} but received"
672
+ f" non inpainting latent channels: {single_image_latent_channels},"
673
+ f" mask channels: {mask_image_channels}, and masked image channels: {single_image_latent_channels}."
674
+ f" Please verify the config of `pipeline.unet` and the `mask_image` and `image` inputs."
675
+ )
676
+
677
+ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
678
+ shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
679
+ if isinstance(generator, list) and len(generator) != batch_size:
680
+ raise ValueError(
681
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
682
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
683
+ )
684
+
685
+ if latents is None:
686
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
687
+ else:
688
+ latents = latents.to(device)
689
+
690
+ # scale the initial noise by the standard deviation required by the scheduler
691
+ latents = latents * self.scheduler.init_noise_sigma
692
+
693
+ return latents
694
+
695
+ def prepare_mask_latents(self, mask_image, batch_size, height, width, dtype, device, do_classifier_free_guidance):
696
+ # resize the mask to latents shape as we concatenate the mask to the latents
697
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
698
+ # and half precision
699
+ mask_image = F.interpolate(mask_image, size=(height // self.vae_scale_factor, width // self.vae_scale_factor))
700
+ mask_image = mask_image.to(device=device, dtype=dtype)
701
+
702
+ # duplicate mask for each generation per prompt, using mps friendly method
703
+ if mask_image.shape[0] < batch_size:
704
+ if not batch_size % mask_image.shape[0] == 0:
705
+ raise ValueError(
706
+ "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
707
+ f" a total batch size of {batch_size}, but {mask_image.shape[0]} masks were passed. Make sure the number"
708
+ " of masks that you pass is divisible by the total requested batch size."
709
+ )
710
+ mask_image = mask_image.repeat(batch_size // mask_image.shape[0], 1, 1, 1)
711
+
712
+ mask_image = torch.cat([mask_image] * 2) if do_classifier_free_guidance else mask_image
713
+
714
+ mask_image_latents = mask_image
715
+
716
+ return mask_image_latents
717
+
718
+ def prepare_masked_image_latents(
719
+ self, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance
720
+ ):
721
+ masked_image = masked_image.to(device=device, dtype=dtype)
722
+
723
+ # encode the mask image into latents space so we can concatenate it to the latents
724
+ if isinstance(generator, list):
725
+ masked_image_latents = [
726
+ self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(generator=generator[i])
727
+ for i in range(batch_size)
728
+ ]
729
+ masked_image_latents = torch.cat(masked_image_latents, dim=0)
730
+ else:
731
+ masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator)
732
+ masked_image_latents = self.vae.config.scaling_factor * masked_image_latents
733
+
734
+ # duplicate masked_image_latents for each generation per prompt, using mps friendly method
735
+ if masked_image_latents.shape[0] < batch_size:
736
+ if not batch_size % masked_image_latents.shape[0] == 0:
737
+ raise ValueError(
738
+ "The passed images and the required batch size don't match. Images are supposed to be duplicated"
739
+ f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
740
+ " Make sure the number of images that you pass is divisible by the total requested batch size."
741
+ )
742
+ masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)
743
+
744
+ masked_image_latents = (
745
+ torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents
746
+ )
747
+
748
+ # aligning device to prevent device errors when concating it with the latent model input
749
+ masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
750
+ return masked_image_latents
751
+
752
+ def _default_height_width(self, height, width, image):
753
+ if isinstance(image, list):
754
+ image = image[0]
755
+
756
+ if height is None:
757
+ if isinstance(image, PIL.Image.Image):
758
+ height = image.height
759
+ elif isinstance(image, torch.Tensor):
760
+ height = image.shape[3]
761
+
762
+ height = (height // 8) * 8 # round down to nearest multiple of 8
763
+
764
+ if width is None:
765
+ if isinstance(image, PIL.Image.Image):
766
+ width = image.width
767
+ elif isinstance(image, torch.Tensor):
768
+ width = image.shape[2]
769
+
770
+ width = (width // 8) * 8 # round down to nearest multiple of 8
771
+
772
+ return height, width
773
+
774
+ @torch.no_grad()
775
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
776
+ def __call__(
777
+ self,
778
+ prompt: Union[str, List[str]] = None,
779
+ image: Union[torch.Tensor, PIL.Image.Image] = None,
780
+ mask_image: Union[torch.Tensor, PIL.Image.Image] = None,
781
+ controlnet_conditioning_image: Union[
782
+ torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]
783
+ ] = None,
784
+ height: Optional[int] = None,
785
+ width: Optional[int] = None,
786
+ num_inference_steps: int = 50,
787
+ guidance_scale: float = 7.5,
788
+ negative_prompt: Optional[Union[str, List[str]]] = None,
789
+ num_images_per_prompt: Optional[int] = 1,
790
+ eta: float = 0.0,
791
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
792
+ latents: Optional[torch.FloatTensor] = None,
793
+ prompt_embeds: Optional[torch.FloatTensor] = None,
794
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
795
+ output_type: Optional[str] = "pil",
796
+ return_dict: bool = True,
797
+ callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
798
+ callback_steps: int = 1,
799
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
800
+ controlnet_conditioning_scale: float = 1.0,
801
+ ):
802
+ r"""
803
+ Function invoked when calling the pipeline for generation.
804
+
805
+ Args:
806
+ prompt (`str` or `List[str]`, *optional*):
807
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
808
+ instead.
809
+ image (`torch.Tensor` or `PIL.Image.Image`):
810
+ `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will
811
+ be masked out with `mask_image` and repainted according to `prompt`.
812
+ mask_image (`torch.Tensor` or `PIL.Image.Image`):
813
+ `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be
814
+ repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted
815
+ to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)
816
+ instead of 3, so the expected shape would be `(B, H, W, 1)`.
817
+ controlnet_conditioning_image (`torch.FloatTensor`, `PIL.Image.Image`, `List[torch.FloatTensor]` or `List[PIL.Image.Image]`):
818
+ The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If
819
+ the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. PIL.Image.Image` can
820
+ also be accepted as an image. The control image is automatically resized to fit the output image.
821
+ height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
822
+ The height in pixels of the generated image.
823
+ width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
824
+ The width in pixels of the generated image.
825
+ num_inference_steps (`int`, *optional*, defaults to 50):
826
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
827
+ expense of slower inference.
828
+ guidance_scale (`float`, *optional*, defaults to 7.5):
829
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
830
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
831
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
832
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
833
+ usually at the expense of lower image quality.
834
+ negative_prompt (`str` or `List[str]`, *optional*):
835
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead.
836
+ Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
837
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
838
+ The number of images to generate per prompt.
839
+ eta (`float`, *optional*, defaults to 0.0):
840
+ Corresponds to parameter eta (Ξ·) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
841
+ [`schedulers.DDIMScheduler`], will be ignored for others.
842
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
843
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
844
+ to make generation deterministic.
845
+ latents (`torch.FloatTensor`, *optional*):
846
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
847
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
848
+ tensor will ge generated by sampling using the supplied random `generator`.
849
+ prompt_embeds (`torch.FloatTensor`, *optional*):
850
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
851
+ provided, text embeddings will be generated from `prompt` input argument.
852
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
853
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
854
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
855
+ argument.
856
+ output_type (`str`, *optional*, defaults to `"pil"`):
857
+ The output format of the generate image. Choose between
858
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
859
+ return_dict (`bool`, *optional*, defaults to `True`):
860
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
861
+ plain tuple.
862
+ callback (`Callable`, *optional*):
863
+ A function that will be called every `callback_steps` steps during inference. The function will be
864
+ called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
865
+ callback_steps (`int`, *optional*, defaults to 1):
866
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
867
+ called at every step.
868
+ cross_attention_kwargs (`dict`, *optional*):
869
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
870
+ `self.processor` in
871
+ [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
872
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
873
+ The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
874
+ to the residual in the original unet.
875
+
876
+ Examples:
877
+
878
+ Returns:
879
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
880
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
881
+ When returning a tuple, the first element is a list with the generated images, and the second element is a
882
+ list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
883
+ (nsfw) content, according to the `safety_checker`.
884
+ """
885
+ # 0. Default height and width to unet
886
+ height, width = self._default_height_width(height, width, controlnet_conditioning_image)
887
+
888
+ # 1. Check inputs. Raise error if not correct
889
+ self.check_inputs(
890
+ prompt,
891
+ image,
892
+ mask_image,
893
+ controlnet_conditioning_image,
894
+ height,
895
+ width,
896
+ callback_steps,
897
+ negative_prompt,
898
+ prompt_embeds,
899
+ negative_prompt_embeds,
900
+ )
901
+
902
+ # 2. Define call parameters
903
+ if prompt is not None and isinstance(prompt, str):
904
+ batch_size = 1
905
+ elif prompt is not None and isinstance(prompt, list):
906
+ batch_size = len(prompt)
907
+ else:
908
+ batch_size = prompt_embeds.shape[0]
909
+
910
+ device = self._execution_device
911
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
912
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
913
+ # corresponds to doing no classifier free guidance.
914
+ do_classifier_free_guidance = guidance_scale > 1.0
915
+
916
+ # 3. Encode input prompt
917
+ prompt_embeds = self._encode_prompt(
918
+ prompt,
919
+ device,
920
+ num_images_per_prompt,
921
+ do_classifier_free_guidance,
922
+ negative_prompt,
923
+ prompt_embeds=prompt_embeds,
924
+ negative_prompt_embeds=negative_prompt_embeds,
925
+ )
926
+
927
+ # 4. Prepare mask, image, and controlnet_conditioning_image
928
+ image = prepare_image(image)
929
+
930
+ mask_image = prepare_mask_image(mask_image)
931
+
932
+ controlnet_conditioning_image = prepare_controlnet_conditioning_image(
933
+ controlnet_conditioning_image,
934
+ width,
935
+ height,
936
+ batch_size * num_images_per_prompt,
937
+ num_images_per_prompt,
938
+ device,
939
+ self.controlnet.dtype,
940
+ )
941
+
942
+ masked_image = image * (mask_image < 0.5)
943
+
944
+ # 5. Prepare timesteps
945
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
946
+ timesteps = self.scheduler.timesteps
947
+
948
+ # 6. Prepare latent variables
949
+ num_channels_latents = self.vae.config.latent_channels
950
+ latents = self.prepare_latents(
951
+ batch_size * num_images_per_prompt,
952
+ num_channels_latents,
953
+ height,
954
+ width,
955
+ prompt_embeds.dtype,
956
+ device,
957
+ generator,
958
+ latents,
959
+ )
960
+
961
+ mask_image_latents = self.prepare_mask_latents(
962
+ mask_image,
963
+ batch_size * num_images_per_prompt,
964
+ height,
965
+ width,
966
+ prompt_embeds.dtype,
967
+ device,
968
+ do_classifier_free_guidance,
969
+ )
970
+
971
+ masked_image_latents = self.prepare_masked_image_latents(
972
+ masked_image,
973
+ batch_size * num_images_per_prompt,
974
+ height,
975
+ width,
976
+ prompt_embeds.dtype,
977
+ device,
978
+ generator,
979
+ do_classifier_free_guidance,
980
+ )
981
+
982
+ if do_classifier_free_guidance:
983
+ controlnet_conditioning_image = torch.cat([controlnet_conditioning_image] * 2)
984
+
985
+ # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
986
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
987
+
988
+ # 8. Denoising loop
989
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
990
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
991
+ for i, t in enumerate(timesteps):
992
+ # expand the latents if we are doing classifier free guidance
993
+ non_inpainting_latent_model_input = (
994
+ torch.cat([latents] * 2) if do_classifier_free_guidance else latents
995
+ )
996
+
997
+ non_inpainting_latent_model_input = self.scheduler.scale_model_input(
998
+ non_inpainting_latent_model_input, t
999
+ )
1000
+
1001
+ inpainting_latent_model_input = torch.cat(
1002
+ [non_inpainting_latent_model_input, mask_image_latents, masked_image_latents], dim=1
1003
+ )
1004
+
1005
+ down_block_res_samples, mid_block_res_sample = self.controlnet(
1006
+ non_inpainting_latent_model_input,
1007
+ t,
1008
+ encoder_hidden_states=prompt_embeds,
1009
+ controlnet_cond=controlnet_conditioning_image,
1010
+ return_dict=False,
1011
+ )
1012
+
1013
+ down_block_res_samples = [
1014
+ down_block_res_sample * controlnet_conditioning_scale
1015
+ for down_block_res_sample in down_block_res_samples
1016
+ ]
1017
+ mid_block_res_sample *= controlnet_conditioning_scale
1018
+
1019
+ # predict the noise residual
1020
+ noise_pred = self.unet(
1021
+ inpainting_latent_model_input,
1022
+ t,
1023
+ encoder_hidden_states=prompt_embeds,
1024
+ cross_attention_kwargs=cross_attention_kwargs,
1025
+ down_block_additional_residuals=down_block_res_samples,
1026
+ mid_block_additional_residual=mid_block_res_sample,
1027
+ ).sample
1028
+
1029
+ # perform guidance
1030
+ if do_classifier_free_guidance:
1031
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
1032
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
1033
+
1034
+ # compute the previous noisy sample x_t -> x_t-1
1035
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
1036
+
1037
+ # call the callback, if provided
1038
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
1039
+ progress_bar.update()
1040
+ if callback is not None and i % callback_steps == 0:
1041
+ callback(i, t, latents)
1042
+
1043
+ # If we do sequential model offloading, let's offload unet and controlnet
1044
+ # manually for max memory savings
1045
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
1046
+ self.unet.to("cpu")
1047
+ self.controlnet.to("cpu")
1048
+ torch.cuda.empty_cache()
1049
+
1050
+ if output_type == "latent":
1051
+ image = latents
1052
+ has_nsfw_concept = None
1053
+ elif output_type == "pil":
1054
+ # 8. Post-processing
1055
+ image = self.decode_latents(latents)
1056
+
1057
+ # 9. Run safety checker
1058
+ image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
1059
+
1060
+ # 10. Convert to PIL
1061
+ image = self.numpy_to_pil(image)
1062
+ else:
1063
+ # 8. Post-processing
1064
+ image = self.decode_latents(latents)
1065
+
1066
+ # 9. Run safety checker
1067
+ image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
1068
+
1069
+ # Offload last model to CPU
1070
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
1071
+ self.final_offload_hook.offload()
1072
+
1073
+ if not return_dict:
1074
+ return (image, has_nsfw_concept)
1075
+
1076
+ return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)