okaris commited on
Commit
1d4b9ba
β€’
1 Parent(s): 6a9c64d
Files changed (8) hide show
  1. .gitmodules +3 -0
  2. README.md +11 -1
  3. demo.py +37 -0
  4. diffusers +1 -0
  5. omni_zero.py +159 -0
  6. pipeline.py +1848 -0
  7. requirements.txt +12 -0
  8. utils.py +99 -0
.gitmodules ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ [submodule "diffusers"]
2
+ path = diffusers
3
+ url = https://github.com/huggingface/diffusers.git
README.md CHANGED
@@ -1,10 +1,20 @@
1
  # Omni-Zero: A diffusion pipeline for zero-shot stylized portrait creation.
2
  ## Try our free demo on [StyleOf](https://styleof.com/s/remix-yourself)
 
3
  ### Single Identity and Style
4
  ![Omni-Zero](https://github.com/okaris/omni-zero/assets/1448702/2c51fb77-a810-4c0a-9555-791a294455ca)
5
- ### Multiple Identities and Styles
 
6
  ![Frame 7-3](https://github.com/okaris/omni-zero/assets/1448702/c5c20961-83bc-47f7-86ed-5948d5590f07)
7
 
 
 
 
 
 
 
 
 
8
  ### Credits
9
  - Special thanks to [fal.ai](https://fal.ai) for providing compute for the research and hosting
10
  - This project wouldn't be possible without the great work of the [InstantX Team](https://github.com/InstantID)
 
1
  # Omni-Zero: A diffusion pipeline for zero-shot stylized portrait creation.
2
  ## Try our free demo on [StyleOf](https://styleof.com/s/remix-yourself)
3
+
4
  ### Single Identity and Style
5
  ![Omni-Zero](https://github.com/okaris/omni-zero/assets/1448702/2c51fb77-a810-4c0a-9555-791a294455ca)
6
+
7
+ ### Multiple Identities and Styles (WIP)
8
  ![Frame 7-3](https://github.com/okaris/omni-zero/assets/1448702/c5c20961-83bc-47f7-86ed-5948d5590f07)
9
 
10
+ ### How to run
11
+ ```
12
+ git clone --recursive https://github.com/okaris/omni-zero.git
13
+ cd omni-zero
14
+ pip install -r requirements.txt
15
+ python demo.py
16
+ ```
17
+
18
  ### Credits
19
  - Special thanks to [fal.ai](https://fal.ai) for providing compute for the research and hosting
20
  - This project wouldn't be possible without the great work of the [InstantX Team](https://github.com/InstantID)
demo.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from omni_zero import OmniZeroSingle
2
+
3
+ def main():
4
+
5
+ omni_zero = OmniZeroSingle(
6
+ base_model="frankjoshua/albedobaseXL_v13",
7
+ )
8
+
9
+ base_image="https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f"
10
+ composition_image="https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f"
11
+ style_image="https://github.com/okaris/omni-zero/assets/1448702/64dc150b-f683-41b1-be23-b6a52c771584"
12
+ identity_image="https://github.com/okaris/omni-zero/assets/1448702/ba193a3a-f90e-4461-848a-560454531c58"
13
+
14
+ images = omni_zero.generate(
15
+ seed=42,
16
+ prompt="A person",
17
+ negative_prompt="blurry, out of focus",
18
+ guidance_scale=3.0,
19
+ number_of_images=1,
20
+ number_of_steps=10,
21
+ base_image=base_image,
22
+ base_image_strength=0.15,
23
+ composition_image=composition_image,
24
+ composition_image_strength=1.0,
25
+ style_image=style_image,
26
+ style_image_strength=1.0,
27
+ identity_image=identity_image,
28
+ identity_image_strength=1.0,
29
+ depth_image=None,
30
+ depth_image_strength=0.5,
31
+ )
32
+
33
+ for i, image in enumerate(images):
34
+ image.save(f"oz_output_{i}.jpg")
35
+
36
+ if __name__ == "__main__":
37
+ main()
diffusers ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit 0d7c4790235ac00b4524b492bc2a680dcc5cf6b0
omni_zero.py ADDED
@@ -0,0 +1,159 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
3
+
4
+ import sys
5
+ sys.path.insert(0, './diffusers/src')
6
+
7
+ import torch
8
+ import torch.nn as nn
9
+
10
+ from huggingface_hub import snapshot_download
11
+ from diffusers import DPMSolverMultistepScheduler
12
+ from diffusers.models import ControlNetModel
13
+
14
+ from transformers import CLIPVisionModelWithProjection
15
+
16
+ from pipeline import OmniZeroPipeline
17
+ from insightface.app import FaceAnalysis
18
+ from controlnet_aux import ZoeDetector
19
+ from utils import draw_kps, load_and_resize_image, align_images
20
+
21
+ import cv2
22
+ import numpy as np
23
+
24
+ class OmniZeroSingle():
25
+ def __init__(self,
26
+ base_model="stabilityai/stable-diffusion-xl-base-1.0",
27
+ ):
28
+ snapshot_download("okaris/antelopev2", local_dir="./models/antelopev2")
29
+ self.face_analysis = FaceAnalysis(name='antelopev2', root='./', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
30
+ self.face_analysis.prepare(ctx_id=0, det_size=(640, 640))
31
+
32
+ dtype = torch.float16
33
+
34
+ ip_adapter_plus_image_encoder = CLIPVisionModelWithProjection.from_pretrained(
35
+ "h94/IP-Adapter",
36
+ subfolder="models/image_encoder",
37
+ torch_dtype=dtype,
38
+ ).to("cuda")
39
+
40
+ zoedepthnet_path = "okaris/zoe-depth-controlnet-xl"
41
+ zoedepthnet = ControlNetModel.from_pretrained(zoedepthnet_path,torch_dtype=dtype).to("cuda")
42
+
43
+ identitiynet_path = "okaris/face-controlnet-xl"
44
+ identitynet = ControlNetModel.from_pretrained(identitiynet_path, torch_dtype=dtype).to("cuda")
45
+
46
+ self.zoe_depth_detector = ZoeDetector.from_pretrained("lllyasviel/Annotators").to("cuda")
47
+
48
+ self.pipeline = OmniZeroPipeline.from_pretrained(
49
+ base_model,
50
+ controlnet=[identitynet, zoedepthnet],
51
+ torch_dtype=dtype,
52
+ image_encoder=ip_adapter_plus_image_encoder,
53
+ ).to("cuda")
54
+
55
+ config = self.pipeline.scheduler.config
56
+ config["timestep_spacing"] = "trailing"
57
+ self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++", final_sigmas_type="zero")
58
+
59
+ self.pipeline.load_ip_adapter(["okaris/ip-adapter-instantid", "h94/IP-Adapter", "h94/IP-Adapter"], subfolder=[None, "sdxl_models", "sdxl_models"], weight_name=["ip-adapter-instantid.bin", "ip-adapter-plus_sdxl_vit-h.safetensors", "ip-adapter-plus_sdxl_vit-h.safetensors"])
60
+ def get_largest_face_embedding_and_kps(self, image, target_image=None):
61
+ face_info = self.face_analysis.get(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR))
62
+ if len(face_info) == 0:
63
+ return None, None
64
+ largest_face = sorted(face_info, key=lambda x: x['bbox'][2] * x['bbox'][3], reverse=True)[0]
65
+ face_embedding = torch.tensor(largest_face['embedding']).to("cuda")
66
+ if target_image is None:
67
+ target_image = image
68
+ zeros = np.zeros((target_image.size[1], target_image.size[0], 3), dtype=np.uint8)
69
+ face_kps_image = draw_kps(zeros, largest_face['kps'])
70
+ return face_embedding, face_kps_image
71
+
72
+ def generate(self,
73
+ seed=42,
74
+ prompt="A person",
75
+ negative_prompt="blurry, out of focus",
76
+ guidance_scale=3.0,
77
+ number_of_images=1,
78
+ number_of_steps=10,
79
+ base_image=None,
80
+ base_image_strength=0.15,
81
+ composition_image=None,
82
+ composition_image_strength=1.0,
83
+ style_image=None,
84
+ style_image_strength=1.0,
85
+ identity_image=None,
86
+ identity_image_strength=1.0,
87
+ depth_image=None,
88
+ depth_image_strength=0.5,
89
+ ):
90
+ resolution = 1024
91
+
92
+ if base_image is not None:
93
+ base_image = load_and_resize_image(base_image, resolution, resolution)
94
+ else:
95
+ if composition_image is not None:
96
+ base_image = load_and_resize_image(composition_image, resolution, resolution)
97
+ else:
98
+ raise ValueError("You must provide a base image or a composition image")
99
+
100
+ if depth_image is None:
101
+ depth_image = self.zoe_depth_detector(base_image, detect_resolution=resolution, image_resolution=resolution)
102
+ else:
103
+ depth_image = load_and_resize_image(depth_image, resolution, resolution)
104
+
105
+ base_image, depth_image = align_images(base_image, depth_image)
106
+
107
+ if composition_image is not None:
108
+ composition_image = load_and_resize_image(composition_image, resolution, resolution)
109
+ else:
110
+ composition_image = base_image
111
+
112
+ if style_image is not None:
113
+ style_image = load_and_resize_image(style_image, resolution, resolution)
114
+ else:
115
+ raise ValueError("You must provide a style image")
116
+
117
+ if identity_image is not None:
118
+ identity_image = load_and_resize_image(identity_image, resolution, resolution)
119
+ else:
120
+ raise ValueError("You must provide an identity image")
121
+
122
+ face_embedding_identity_image, target_kps = self.get_largest_face_embedding_and_kps(identity_image, base_image)
123
+ if face_embedding_identity_image is None:
124
+ raise ValueError("No face found in the identity image, the image might be cropped too tightly or the face is too small")
125
+
126
+ face_embedding_base_image, face_kps_base_image = self.get_largest_face_embedding_and_kps(base_image)
127
+ if face_embedding_base_image is not None:
128
+ target_kps = face_kps_base_image
129
+
130
+ self.pipeline.set_ip_adapter_scale([identity_image_strength,
131
+ {
132
+ "down": { "block_2": [0.0, 0.0] },
133
+ "up": { "block_0": [0.0, style_image_strength, 0.0] }
134
+ },
135
+ {
136
+ "down": { "block_2": [0.0, composition_image_strength] },
137
+ "up": { "block_0": [0.0, 0.0, 0.0] }
138
+ }
139
+ ])
140
+
141
+ generator = torch.Generator(device="cpu").manual_seed(seed)
142
+
143
+ images = self.pipeline(
144
+ prompt=prompt,
145
+ negative_prompt=negative_prompt,
146
+ guidance_scale=guidance_scale,
147
+ ip_adapter_image=[face_embedding_identity_image, style_image, composition_image],
148
+ image=base_image,
149
+ control_image=[target_kps, depth_image],
150
+ controlnet_conditioning_scale=[identity_image_strength, depth_image_strength],
151
+ identity_control_indices=[(0,0)],
152
+ num_inference_steps=number_of_steps,
153
+ num_images_per_prompt=number_of_images,
154
+ strength=(1-base_image_strength),
155
+ generator=generator,
156
+ seed=seed,
157
+ ).images
158
+
159
+ return images
pipeline.py ADDED
@@ -0,0 +1,1848 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+
16
+ import inspect
17
+ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
18
+
19
+ import numpy as np
20
+ import PIL.Image
21
+ import torch
22
+ import torch.nn.functional as F
23
+ import torchsde
24
+
25
+ from transformers import (
26
+ CLIPImageProcessor,
27
+ CLIPTextModel,
28
+ CLIPTextModelWithProjection,
29
+ CLIPTokenizer,
30
+ CLIPVisionModelWithProjection,
31
+ )
32
+
33
+ from diffusers.utils.import_utils import is_invisible_watermark_available
34
+
35
+ from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
36
+ from diffusers.loaders import (
37
+ FromSingleFileMixin,
38
+ IPAdapterMixin,
39
+ StableDiffusionXLLoraLoaderMixin,
40
+ TextualInversionLoaderMixin,
41
+ )
42
+ from diffusers.models import AutoencoderKL, ControlNetModel, ImageProjection, UNet2DConditionModel
43
+ from diffusers.models.attention_processor import (
44
+ AttnProcessor2_0,
45
+ LoRAAttnProcessor2_0,
46
+ LoRAXFormersAttnProcessor,
47
+ XFormersAttnProcessor,
48
+ )
49
+ from diffusers.models.lora import adjust_lora_scale_text_encoder
50
+ from diffusers.schedulers import KarrasDiffusionSchedulers
51
+ from diffusers.utils import (
52
+ USE_PEFT_BACKEND,
53
+ deprecate,
54
+ logging,
55
+ replace_example_docstring,
56
+ scale_lora_layers,
57
+ unscale_lora_layers,
58
+ )
59
+ from diffusers.utils.torch_utils import is_compiled_module, randn_tensor
60
+ from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
61
+ from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
62
+
63
+
64
+ if is_invisible_watermark_available():
65
+ from diffusers.pipelines.stable_diffusion_xl.watermark import StableDiffusionXLWatermarker
66
+
67
+ from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
68
+
69
+
70
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
71
+
72
+ EXAMPLE_DOC_STRING = """
73
+ Examples:
74
+ ```py
75
+ >>> # pip install accelerate transformers safetensors diffusers
76
+
77
+ >>> import torch
78
+ >>> import numpy as np
79
+ >>> from PIL import Image
80
+
81
+ >>> from transformers import DPTFeatureExtractor, DPTForDepthEstimation
82
+ >>> from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline, AutoencoderKL
83
+ >>> from diffusers.utils import load_image
84
+
85
+
86
+ >>> depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
87
+ >>> feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")
88
+ >>> controlnet = ControlNetModel.from_pretrained(
89
+ ... "diffusers/controlnet-depth-sdxl-1.0-small",
90
+ ... variant="fp16",
91
+ ... use_safetensors=True,
92
+ ... torch_dtype=torch.float16,
93
+ ... ).to("cuda")
94
+ >>> vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to("cuda")
95
+ >>> pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
96
+ ... "stabilityai/stable-diffusion-xl-base-1.0",
97
+ ... controlnet=controlnet,
98
+ ... vae=vae,
99
+ ... variant="fp16",
100
+ ... use_safetensors=True,
101
+ ... torch_dtype=torch.float16,
102
+ ... ).to("cuda")
103
+ >>> pipe.enable_model_cpu_offload()
104
+
105
+
106
+ >>> def get_depth_map(image):
107
+ ... image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
108
+ ... with torch.no_grad(), torch.autocast("cuda"):
109
+ ... depth_map = depth_estimator(image).predicted_depth
110
+
111
+ ... depth_map = torch.nn.functional.interpolate(
112
+ ... depth_map.unsqueeze(1),
113
+ ... size=(1024, 1024),
114
+ ... mode="bicubic",
115
+ ... align_corners=False,
116
+ ... )
117
+ ... depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
118
+ ... depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
119
+ ... depth_map = (depth_map - depth_min) / (depth_max - depth_min)
120
+ ... image = torch.cat([depth_map] * 3, dim=1)
121
+ ... image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
122
+ ... image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
123
+ ... return image
124
+
125
+
126
+ >>> prompt = "A robot, 4k photo"
127
+ >>> image = load_image(
128
+ ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
129
+ ... "/kandinsky/cat.png"
130
+ ... ).resize((1024, 1024))
131
+ >>> controlnet_conditioning_scale = 0.5 # recommended for good generalization
132
+ >>> depth_image = get_depth_map(image)
133
+
134
+ >>> images = pipe(
135
+ ... prompt,
136
+ ... image=image,
137
+ ... control_image=depth_image,
138
+ ... strength=0.99,
139
+ ... num_inference_steps=50,
140
+ ... controlnet_conditioning_scale=controlnet_conditioning_scale,
141
+ ... ).images
142
+ >>> images[0].save(f"robot_cat.png")
143
+ ```
144
+ """
145
+
146
+
147
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
148
+ def retrieve_latents(
149
+ encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
150
+ ):
151
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
152
+ return encoder_output.latent_dist.sample(generator)
153
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
154
+ return encoder_output.latent_dist.mode()
155
+ elif hasattr(encoder_output, "latents"):
156
+ return encoder_output.latents
157
+ else:
158
+ raise AttributeError("Could not access latents of provided encoder_output")
159
+
160
+ class BatchedBrownianTree:
161
+ """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
162
+
163
+ def __init__(self, x, t0, t1, seed=None, **kwargs):
164
+ self.cpu_tree = True
165
+ if "cpu" in kwargs:
166
+ self.cpu_tree = kwargs.pop("cpu")
167
+ t0, t1, self.sign = self.sort(t0, t1)
168
+ w0 = kwargs.get('w0', torch.zeros_like(x))
169
+ if seed is None:
170
+ seed = torch.randint(0, 2 ** 63 - 1, []).item()
171
+ self.batched = True
172
+ try:
173
+ assert len(seed) == x.shape[0]
174
+ w0 = w0[0]
175
+ except TypeError:
176
+ seed = [seed]
177
+ self.batched = False
178
+ if self.cpu_tree:
179
+ self.trees = [torchsde.BrownianTree(t0.cpu(), w0.cpu(), t1.cpu(), entropy=s, **kwargs) for s in seed]
180
+ else:
181
+ self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
182
+
183
+ @staticmethod
184
+ def sort(a, b):
185
+ return (a, b, 1) if a < b else (b, a, -1)
186
+
187
+ def __call__(self, t0, t1):
188
+ t0, t1, sign = self.sort(t0, t1)
189
+ if self.cpu_tree:
190
+ w = torch.stack([tree(t0.cpu().float(), t1.cpu().float()).to(t0.dtype).to(t0.device) for tree in self.trees]) * (self.sign * sign)
191
+ else:
192
+ w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
193
+
194
+ return w if self.batched else w[0]
195
+
196
+
197
+ class BrownianTreeNoiseSampler:
198
+ """A noise sampler backed by a torchsde.BrownianTree.
199
+
200
+ Args:
201
+ x (Tensor): The tensor whose shape, device and dtype to use to generate
202
+ random samples.
203
+ sigma_min (float): The low end of the valid interval.
204
+ sigma_max (float): The high end of the valid interval.
205
+ seed (int or List[int]): The random seed. If a list of seeds is
206
+ supplied instead of a single integer, then the noise sampler will
207
+ use one BrownianTree per batch item, each with its own seed.
208
+ transform (callable): A function that maps sigma to the sampler's
209
+ internal timestep.
210
+ """
211
+
212
+ def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
213
+ self.transform = transform
214
+ t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
215
+ self.tree = BatchedBrownianTree(x, t0, t1, seed, cpu=cpu)
216
+
217
+ def __call__(self, sigma, sigma_next):
218
+ t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
219
+ return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
220
+
221
+
222
+ class OmniZeroPipeline(
223
+ DiffusionPipeline,
224
+ StableDiffusionMixin,
225
+ TextualInversionLoaderMixin,
226
+ StableDiffusionXLLoraLoaderMixin,
227
+ FromSingleFileMixin,
228
+ IPAdapterMixin,
229
+ ):
230
+ r"""
231
+ Pipeline for image-to-image generation using Stable Diffusion XL with ControlNet guidance.
232
+
233
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
234
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
235
+
236
+ The pipeline also inherits the following loading methods:
237
+ - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
238
+ - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights
239
+ - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights
240
+ - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
241
+
242
+ Args:
243
+ vae ([`AutoencoderKL`]):
244
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
245
+ text_encoder ([`CLIPTextModel`]):
246
+ Frozen text-encoder. Stable Diffusion uses the text portion of
247
+ [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
248
+ the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
249
+ text_encoder_2 ([` CLIPTextModelWithProjection`]):
250
+ Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of
251
+ [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
252
+ specifically the
253
+ [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
254
+ variant.
255
+ tokenizer (`CLIPTokenizer`):
256
+ Tokenizer of class
257
+ [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
258
+ tokenizer_2 (`CLIPTokenizer`):
259
+ Second Tokenizer of class
260
+ [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
261
+ unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
262
+ controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
263
+ Provides additional conditioning to the unet during the denoising process. If you set multiple ControlNets
264
+ as a list, the outputs from each ControlNet are added together to create one combined additional
265
+ conditioning.
266
+ scheduler ([`SchedulerMixin`]):
267
+ A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
268
+ [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
269
+ requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`):
270
+ Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the
271
+ config of `stabilityai/stable-diffusion-xl-refiner-1-0`.
272
+ force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`):
273
+ Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of
274
+ `stabilityai/stable-diffusion-xl-base-1-0`.
275
+ add_watermarker (`bool`, *optional*):
276
+ Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to
277
+ watermark output images. If not defined, it will default to True if the package is installed, otherwise no
278
+ watermarker will be used.
279
+ feature_extractor ([`~transformers.CLIPImageProcessor`]):
280
+ A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
281
+ """
282
+
283
+ model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae"
284
+ _optional_components = [
285
+ "tokenizer",
286
+ "tokenizer_2",
287
+ "text_encoder",
288
+ "text_encoder_2",
289
+ "feature_extractor",
290
+ "image_encoder",
291
+ ]
292
+ _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
293
+
294
+ def __init__(
295
+ self,
296
+ vae: AutoencoderKL,
297
+ text_encoder: CLIPTextModel,
298
+ text_encoder_2: CLIPTextModelWithProjection,
299
+ tokenizer: CLIPTokenizer,
300
+ tokenizer_2: CLIPTokenizer,
301
+ unet: UNet2DConditionModel,
302
+ controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
303
+ scheduler: KarrasDiffusionSchedulers,
304
+ requires_aesthetics_score: bool = False,
305
+ force_zeros_for_empty_prompt: bool = True,
306
+ add_watermarker: Optional[bool] = None,
307
+ feature_extractor: CLIPImageProcessor = None,
308
+ image_encoder: CLIPVisionModelWithProjection = None,
309
+ ):
310
+ super().__init__()
311
+
312
+ if isinstance(controlnet, (list, tuple)):
313
+ controlnet = MultiControlNetModel(controlnet)
314
+
315
+ self.register_modules(
316
+ vae=vae,
317
+ text_encoder=text_encoder,
318
+ text_encoder_2=text_encoder_2,
319
+ tokenizer=tokenizer,
320
+ tokenizer_2=tokenizer_2,
321
+ unet=unet,
322
+ controlnet=controlnet,
323
+ scheduler=scheduler,
324
+ feature_extractor=feature_extractor,
325
+ image_encoder=image_encoder,
326
+ )
327
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
328
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True)
329
+ self.control_image_processor = VaeImageProcessor(
330
+ vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False
331
+ )
332
+ add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available()
333
+
334
+ if add_watermarker:
335
+ self.watermark = StableDiffusionXLWatermarker()
336
+ else:
337
+ self.watermark = None
338
+
339
+ self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
340
+ self.register_to_config(requires_aesthetics_score=requires_aesthetics_score)
341
+
342
+ self.ays_noise_sigmas = {"SD1": [14.6146412293, 6.4745760956, 3.8636745985, 2.6946151520, 1.8841921177, 1.3943805092, 0.9642583904, 0.6523686016, 0.3977456272, 0.1515232662, 0.0291671582],
343
+ "SDXL":[14.6146412293, 6.3184485287, 3.7681790315, 2.1811480769, 1.3405244945, 0.8620721141, 0.5550693289, 0.3798540708, 0.2332364134, 0.1114188177, 0.0291671582],
344
+ "SVD": [700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.002]}
345
+
346
+ @staticmethod
347
+ def _loglinear_interp(t_steps, num_steps):
348
+ xs = np.linspace(0, 1, len(t_steps))
349
+ ys = np.log(t_steps[::-1])
350
+
351
+ new_xs = np.linspace(0, 1, num_steps)
352
+ new_ys = np.interp(new_xs, xs, ys)
353
+
354
+ return np.exp(new_ys)[::-1].copy()
355
+
356
+ # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt
357
+ def encode_prompt(
358
+ self,
359
+ prompt: str,
360
+ prompt_2: Optional[str] = None,
361
+ device: Optional[torch.device] = None,
362
+ num_images_per_prompt: int = 1,
363
+ do_classifier_free_guidance: bool = True,
364
+ negative_prompt: Optional[str] = None,
365
+ negative_prompt_2: Optional[str] = None,
366
+ prompt_embeds: Optional[torch.FloatTensor] = None,
367
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
368
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
369
+ negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
370
+ lora_scale: Optional[float] = None,
371
+ clip_skip: Optional[int] = None,
372
+ ):
373
+ r"""
374
+ Encodes the prompt into text encoder hidden states.
375
+
376
+ Args:
377
+ prompt (`str` or `List[str]`, *optional*):
378
+ prompt to be encoded
379
+ prompt_2 (`str` or `List[str]`, *optional*):
380
+ The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
381
+ used in both text-encoders
382
+ device: (`torch.device`):
383
+ torch device
384
+ num_images_per_prompt (`int`):
385
+ number of images that should be generated per prompt
386
+ do_classifier_free_guidance (`bool`):
387
+ whether to use classifier free guidance or not
388
+ negative_prompt (`str` or `List[str]`, *optional*):
389
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
390
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
391
+ less than `1`).
392
+ negative_prompt_2 (`str` or `List[str]`, *optional*):
393
+ The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
394
+ `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
395
+ prompt_embeds (`torch.FloatTensor`, *optional*):
396
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
397
+ provided, text embeddings will be generated from `prompt` input argument.
398
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
399
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
400
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
401
+ argument.
402
+ pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
403
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
404
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
405
+ negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
406
+ Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
407
+ weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
408
+ input argument.
409
+ lora_scale (`float`, *optional*):
410
+ A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
411
+ clip_skip (`int`, *optional*):
412
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
413
+ the output of the pre-final layer will be used for computing the prompt embeddings.
414
+ """
415
+ device = device or self._execution_device
416
+
417
+ # set lora scale so that monkey patched LoRA
418
+ # function of text encoder can correctly access it
419
+ if lora_scale is not None and isinstance(self, StableDiffusionXLLoraLoaderMixin):
420
+ self._lora_scale = lora_scale
421
+
422
+ # dynamically adjust the LoRA scale
423
+ if self.text_encoder is not None:
424
+ if not USE_PEFT_BACKEND:
425
+ adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
426
+ else:
427
+ scale_lora_layers(self.text_encoder, lora_scale)
428
+
429
+ if self.text_encoder_2 is not None:
430
+ if not USE_PEFT_BACKEND:
431
+ adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale)
432
+ else:
433
+ scale_lora_layers(self.text_encoder_2, lora_scale)
434
+
435
+ prompt = [prompt] if isinstance(prompt, str) else prompt
436
+
437
+ if prompt is not None:
438
+ batch_size = len(prompt)
439
+ else:
440
+ batch_size = prompt_embeds.shape[0]
441
+
442
+ # Define tokenizers and text encoders
443
+ tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
444
+ text_encoders = (
445
+ [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
446
+ )
447
+
448
+ if prompt_embeds is None:
449
+ prompt_2 = prompt_2 or prompt
450
+ prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
451
+
452
+ # textual inversion: process multi-vector tokens if necessary
453
+ prompt_embeds_list = []
454
+ prompts = [prompt, prompt_2]
455
+ for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
456
+ if isinstance(self, TextualInversionLoaderMixin):
457
+ prompt = self.maybe_convert_prompt(prompt, tokenizer)
458
+
459
+ text_inputs = tokenizer(
460
+ prompt,
461
+ padding="max_length",
462
+ max_length=tokenizer.model_max_length,
463
+ truncation=True,
464
+ return_tensors="pt",
465
+ )
466
+
467
+ text_input_ids = text_inputs.input_ids
468
+ untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
469
+
470
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
471
+ text_input_ids, untruncated_ids
472
+ ):
473
+ removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
474
+ logger.warning(
475
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
476
+ f" {tokenizer.model_max_length} tokens: {removed_text}"
477
+ )
478
+
479
+ prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
480
+
481
+ # We are only ALWAYS interested in the pooled output of the final text encoder
482
+ pooled_prompt_embeds = prompt_embeds[0]
483
+ if clip_skip is None:
484
+ prompt_embeds = prompt_embeds.hidden_states[-2]
485
+ else:
486
+ # "2" because SDXL always indexes from the penultimate layer.
487
+ prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
488
+
489
+ prompt_embeds_list.append(prompt_embeds)
490
+
491
+ prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
492
+
493
+ # get unconditional embeddings for classifier free guidance
494
+ zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
495
+ if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
496
+ negative_prompt_embeds = torch.zeros_like(prompt_embeds)
497
+ negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
498
+ elif do_classifier_free_guidance and negative_prompt_embeds is None:
499
+ negative_prompt = negative_prompt or ""
500
+ negative_prompt_2 = negative_prompt_2 or negative_prompt
501
+
502
+ # normalize str to list
503
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
504
+ negative_prompt_2 = (
505
+ batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
506
+ )
507
+
508
+ uncond_tokens: List[str]
509
+ if prompt is not None and type(prompt) is not type(negative_prompt):
510
+ raise TypeError(
511
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
512
+ f" {type(prompt)}."
513
+ )
514
+ elif batch_size != len(negative_prompt):
515
+ raise ValueError(
516
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
517
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
518
+ " the batch size of `prompt`."
519
+ )
520
+ else:
521
+ uncond_tokens = [negative_prompt, negative_prompt_2]
522
+
523
+ negative_prompt_embeds_list = []
524
+ for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
525
+ if isinstance(self, TextualInversionLoaderMixin):
526
+ negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
527
+
528
+ max_length = prompt_embeds.shape[1]
529
+ uncond_input = tokenizer(
530
+ negative_prompt,
531
+ padding="max_length",
532
+ max_length=max_length,
533
+ truncation=True,
534
+ return_tensors="pt",
535
+ )
536
+
537
+ negative_prompt_embeds = text_encoder(
538
+ uncond_input.input_ids.to(device),
539
+ output_hidden_states=True,
540
+ )
541
+ # We are only ALWAYS interested in the pooled output of the final text encoder
542
+ negative_pooled_prompt_embeds = negative_prompt_embeds[0]
543
+ negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
544
+
545
+ negative_prompt_embeds_list.append(negative_prompt_embeds)
546
+
547
+ negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
548
+
549
+ if self.text_encoder_2 is not None:
550
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
551
+ else:
552
+ prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device)
553
+
554
+ bs_embed, seq_len, _ = prompt_embeds.shape
555
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
556
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
557
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
558
+
559
+ if do_classifier_free_guidance:
560
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
561
+ seq_len = negative_prompt_embeds.shape[1]
562
+
563
+ if self.text_encoder_2 is not None:
564
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
565
+ else:
566
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device)
567
+
568
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
569
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
570
+
571
+ pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
572
+ bs_embed * num_images_per_prompt, -1
573
+ )
574
+ if do_classifier_free_guidance:
575
+ negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
576
+ bs_embed * num_images_per_prompt, -1
577
+ )
578
+
579
+ if self.text_encoder is not None:
580
+ if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
581
+ # Retrieve the original scale by scaling back the LoRA layers
582
+ unscale_lora_layers(self.text_encoder, lora_scale)
583
+
584
+ if self.text_encoder_2 is not None:
585
+ if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
586
+ # Retrieve the original scale by scaling back the LoRA layers
587
+ unscale_lora_layers(self.text_encoder_2, lora_scale)
588
+
589
+ return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
590
+
591
+ # Copied from ..stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
592
+ def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None, unconditional_noising_factor=1.0):
593
+ dtype = next(self.image_encoder.parameters()).dtype
594
+
595
+ needs_encoding = not isinstance(image, torch.Tensor)
596
+ if needs_encoding:
597
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
598
+
599
+ image = image.to(device=device, dtype=dtype)
600
+
601
+ avg_image = torch.mean(image, dim=0, keepdim=True).to(dtype=torch.float32)
602
+ seed = int(torch.sum(avg_image).item()) % 1000000007
603
+ torch.manual_seed(seed)
604
+ additional_noise_for_uncond = torch.rand_like(image) * unconditional_noising_factor
605
+
606
+ if output_hidden_states:
607
+ if needs_encoding:
608
+ image_encoded = self.image_encoder(image, output_hidden_states=True)
609
+ image_enc_hidden_states = image_encoded.hidden_states[-2]
610
+ else:
611
+ image_enc_hidden_states = image.unsqueeze(0).unsqueeze(0)
612
+ image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
613
+
614
+ if needs_encoding:
615
+ uncond_image_encoded = self.image_encoder(additional_noise_for_uncond, output_hidden_states=True)
616
+ uncond_image_enc_hidden_states = uncond_image_encoded.hidden_states[-2]
617
+ else:
618
+ uncond_image_enc_hidden_states = additional_noise_for_uncond.unsqueeze(0).unsqueeze(0)
619
+ uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
620
+ num_images_per_prompt, dim=0
621
+ )
622
+ return image_enc_hidden_states, uncond_image_enc_hidden_states
623
+ else:
624
+ if needs_encoding:
625
+ image_encoded = self.image_encoder(image)
626
+ image_embeds = image_encoded.image_embeds
627
+ else:
628
+ image_embeds = image.unsqueeze(0).unsqueeze(0)
629
+ if needs_encoding:
630
+ uncond_image_encoded = self.image_encoder(additional_noise_for_uncond)
631
+ uncond_image_embeds = uncond_image_encoded.image_embeds
632
+ else:
633
+ uncond_image_embeds = additional_noise_for_uncond.unsqueeze(0).unsqueeze(0)
634
+
635
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
636
+ uncond_image_embeds = uncond_image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
637
+
638
+ return image_embeds, uncond_image_embeds
639
+
640
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
641
+ def prepare_ip_adapter_image_embeds(
642
+ self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
643
+ ):
644
+ if ip_adapter_image_embeds is None:
645
+ if not isinstance(ip_adapter_image, list):
646
+ ip_adapter_image = [ip_adapter_image]
647
+
648
+ if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
649
+ raise ValueError(
650
+ f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
651
+ )
652
+
653
+ image_embeds = []
654
+ for single_ip_adapter_image, image_proj_layer in zip(
655
+ ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
656
+ ):
657
+ output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
658
+ single_image_embeds, single_negative_image_embeds = self.encode_image(
659
+ single_ip_adapter_image, device, 1, output_hidden_state
660
+ )
661
+ single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0)
662
+ single_negative_image_embeds = torch.stack(
663
+ [single_negative_image_embeds] * num_images_per_prompt, dim=0
664
+ )
665
+
666
+ if do_classifier_free_guidance:
667
+ single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
668
+ single_image_embeds = single_image_embeds.to(device)
669
+
670
+ image_embeds.append(single_image_embeds)
671
+ else:
672
+ repeat_dims = [1]
673
+ image_embeds = []
674
+ for single_image_embeds in ip_adapter_image_embeds:
675
+ if do_classifier_free_guidance:
676
+ single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
677
+ single_image_embeds = single_image_embeds.repeat(
678
+ num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
679
+ )
680
+ single_negative_image_embeds = single_negative_image_embeds.repeat(
681
+ num_images_per_prompt, *(repeat_dims * len(single_negative_image_embeds.shape[1:]))
682
+ )
683
+ single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
684
+ else:
685
+ single_image_embeds = single_image_embeds.repeat(
686
+ num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
687
+ )
688
+ image_embeds.append(single_image_embeds)
689
+
690
+ return image_embeds
691
+
692
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
693
+ def prepare_extra_step_kwargs(self, generator, eta):
694
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
695
+ # eta (Ξ·) is only used with the DDIMScheduler, it will be ignored for other schedulers.
696
+ # eta corresponds to Ξ· in DDIM paper: https://arxiv.org/abs/2010.02502
697
+ # and should be between [0, 1]
698
+
699
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
700
+ extra_step_kwargs = {}
701
+ if accepts_eta:
702
+ extra_step_kwargs["eta"] = eta
703
+
704
+ # check if the scheduler accepts generator
705
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
706
+ if accepts_generator:
707
+ extra_step_kwargs["generator"] = generator
708
+ return extra_step_kwargs
709
+
710
+ def check_inputs(
711
+ self,
712
+ prompt,
713
+ prompt_2,
714
+ image,
715
+ strength,
716
+ num_inference_steps,
717
+ callback_steps,
718
+ negative_prompt=None,
719
+ negative_prompt_2=None,
720
+ prompt_embeds=None,
721
+ negative_prompt_embeds=None,
722
+ pooled_prompt_embeds=None,
723
+ negative_pooled_prompt_embeds=None,
724
+ ip_adapter_image=None,
725
+ ip_adapter_image_embeds=None,
726
+ controlnet_conditioning_scale=1.0,
727
+ control_guidance_start=0.0,
728
+ control_guidance_end=1.0,
729
+ callback_on_step_end_tensor_inputs=None,
730
+ ):
731
+ if strength < 0 or strength > 1:
732
+ raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
733
+ if num_inference_steps is None:
734
+ raise ValueError("`num_inference_steps` cannot be None.")
735
+ elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0:
736
+ raise ValueError(
737
+ f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type"
738
+ f" {type(num_inference_steps)}."
739
+ )
740
+
741
+ if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
742
+ raise ValueError(
743
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
744
+ f" {type(callback_steps)}."
745
+ )
746
+
747
+ if callback_on_step_end_tensor_inputs is not None and not all(
748
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
749
+ ):
750
+ raise ValueError(
751
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
752
+ )
753
+
754
+ if prompt is not None and prompt_embeds is not None:
755
+ raise ValueError(
756
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
757
+ " only forward one of the two."
758
+ )
759
+ elif prompt_2 is not None and prompt_embeds is not None:
760
+ raise ValueError(
761
+ f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
762
+ " only forward one of the two."
763
+ )
764
+ elif prompt is None and prompt_embeds is None:
765
+ raise ValueError(
766
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
767
+ )
768
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
769
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
770
+ elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
771
+ raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
772
+
773
+ if negative_prompt is not None and negative_prompt_embeds is not None:
774
+ raise ValueError(
775
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
776
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
777
+ )
778
+ elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
779
+ raise ValueError(
780
+ f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
781
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
782
+ )
783
+
784
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
785
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
786
+ raise ValueError(
787
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
788
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
789
+ f" {negative_prompt_embeds.shape}."
790
+ )
791
+
792
+ if prompt_embeds is not None and pooled_prompt_embeds is None:
793
+ raise ValueError(
794
+ "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
795
+ )
796
+
797
+ if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
798
+ raise ValueError(
799
+ "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
800
+ )
801
+
802
+ # `prompt` needs more sophisticated handling when there are multiple
803
+ # conditionings.
804
+ if isinstance(self.controlnet, MultiControlNetModel):
805
+ if isinstance(prompt, list):
806
+ logger.warning(
807
+ f"You have {len(self.controlnet.nets)} ControlNets and you have passed {len(prompt)}"
808
+ " prompts. The conditionings will be fixed across the prompts."
809
+ )
810
+
811
+ # Check `image`
812
+ is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
813
+ self.controlnet, torch._dynamo.eval_frame.OptimizedModule
814
+ )
815
+ if (
816
+ isinstance(self.controlnet, ControlNetModel)
817
+ or is_compiled
818
+ and isinstance(self.controlnet._orig_mod, ControlNetModel)
819
+ ):
820
+ self.check_image(image, prompt, prompt_embeds)
821
+ elif (
822
+ isinstance(self.controlnet, MultiControlNetModel)
823
+ or is_compiled
824
+ and isinstance(self.controlnet._orig_mod, MultiControlNetModel)
825
+ ):
826
+ if not isinstance(image, list):
827
+ raise TypeError("For multiple controlnets: `image` must be type `list`")
828
+
829
+ # When `image` is a nested list:
830
+ # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]])
831
+ elif any(isinstance(i, list) for i in image):
832
+ raise ValueError("A single batch of multiple conditionings are supported at the moment.")
833
+ elif len(image) != len(self.controlnet.nets):
834
+ raise ValueError(
835
+ f"For multiple controlnets: `image` must have the same length as the number of controlnets, but got {len(image)} images and {len(self.controlnet.nets)} ControlNets."
836
+ )
837
+
838
+ for image_ in image:
839
+ self.check_image(image_, prompt, prompt_embeds)
840
+ else:
841
+ assert False
842
+
843
+ # Check `controlnet_conditioning_scale`
844
+ if (
845
+ isinstance(self.controlnet, ControlNetModel)
846
+ or is_compiled
847
+ and isinstance(self.controlnet._orig_mod, ControlNetModel)
848
+ ):
849
+ if not isinstance(controlnet_conditioning_scale, float):
850
+ raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")
851
+ elif (
852
+ isinstance(self.controlnet, MultiControlNetModel)
853
+ or is_compiled
854
+ and isinstance(self.controlnet._orig_mod, MultiControlNetModel)
855
+ ):
856
+ if isinstance(controlnet_conditioning_scale, list):
857
+ if any(isinstance(i, list) for i in controlnet_conditioning_scale):
858
+ raise ValueError("A single batch of multiple conditionings are supported at the moment.")
859
+ elif isinstance(controlnet_conditioning_scale, list) and len(controlnet_conditioning_scale) != len(
860
+ self.controlnet.nets
861
+ ):
862
+ raise ValueError(
863
+ "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"
864
+ " the same length as the number of controlnets"
865
+ )
866
+ else:
867
+ assert False
868
+
869
+ if not isinstance(control_guidance_start, (tuple, list)):
870
+ control_guidance_start = [control_guidance_start]
871
+
872
+ if not isinstance(control_guidance_end, (tuple, list)):
873
+ control_guidance_end = [control_guidance_end]
874
+
875
+ if len(control_guidance_start) != len(control_guidance_end):
876
+ raise ValueError(
877
+ f"`control_guidance_start` has {len(control_guidance_start)} elements, but `control_guidance_end` has {len(control_guidance_end)} elements. Make sure to provide the same number of elements to each list."
878
+ )
879
+
880
+ if isinstance(self.controlnet, MultiControlNetModel):
881
+ if len(control_guidance_start) != len(self.controlnet.nets):
882
+ raise ValueError(
883
+ f"`control_guidance_start`: {control_guidance_start} has {len(control_guidance_start)} elements but there are {len(self.controlnet.nets)} controlnets available. Make sure to provide {len(self.controlnet.nets)}."
884
+ )
885
+
886
+ for start, end in zip(control_guidance_start, control_guidance_end):
887
+ if start >= end:
888
+ raise ValueError(
889
+ f"control guidance start: {start} cannot be larger or equal to control guidance end: {end}."
890
+ )
891
+ if start < 0.0:
892
+ raise ValueError(f"control guidance start: {start} can't be smaller than 0.")
893
+ if end > 1.0:
894
+ raise ValueError(f"control guidance end: {end} can't be larger than 1.0.")
895
+
896
+ if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
897
+ raise ValueError(
898
+ "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
899
+ )
900
+
901
+ if ip_adapter_image_embeds is not None:
902
+ if not isinstance(ip_adapter_image_embeds, list):
903
+ raise ValueError(
904
+ f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
905
+ )
906
+ elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
907
+ raise ValueError(
908
+ f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
909
+ )
910
+
911
+ # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.check_image
912
+ def check_image(self, image, prompt, prompt_embeds):
913
+ image_is_pil = isinstance(image, PIL.Image.Image)
914
+ image_is_tensor = isinstance(image, torch.Tensor)
915
+ image_is_np = isinstance(image, np.ndarray)
916
+ image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image)
917
+ image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor)
918
+ image_is_np_list = isinstance(image, list) and isinstance(image[0], np.ndarray)
919
+
920
+ if (
921
+ not image_is_pil
922
+ and not image_is_tensor
923
+ and not image_is_np
924
+ and not image_is_pil_list
925
+ and not image_is_tensor_list
926
+ and not image_is_np_list
927
+ ):
928
+ raise TypeError(
929
+ f"image must be passed and be one of PIL image, numpy array, torch tensor, list of PIL images, list of numpy arrays or list of torch tensors, but is {type(image)}"
930
+ )
931
+
932
+ if image_is_pil:
933
+ image_batch_size = 1
934
+ else:
935
+ image_batch_size = len(image)
936
+
937
+ if prompt is not None and isinstance(prompt, str):
938
+ prompt_batch_size = 1
939
+ elif prompt is not None and isinstance(prompt, list):
940
+ prompt_batch_size = len(prompt)
941
+ elif prompt_embeds is not None:
942
+ prompt_batch_size = prompt_embeds.shape[0]
943
+
944
+ if image_batch_size != 1 and image_batch_size != prompt_batch_size:
945
+ raise ValueError(
946
+ f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}"
947
+ )
948
+
949
+ # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image
950
+ def prepare_control_image(
951
+ self,
952
+ image,
953
+ width,
954
+ height,
955
+ batch_size,
956
+ num_images_per_prompt,
957
+ device,
958
+ dtype,
959
+ do_classifier_free_guidance=False,
960
+ guess_mode=False,
961
+ ):
962
+ image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32)
963
+ image_batch_size = image.shape[0]
964
+
965
+ if image_batch_size == 1:
966
+ repeat_by = batch_size
967
+ else:
968
+ # image batch size is the same as prompt batch size
969
+ repeat_by = num_images_per_prompt
970
+
971
+ image = image.repeat_interleave(repeat_by, dim=0)
972
+
973
+ image = image.to(device=device, dtype=dtype)
974
+
975
+ if do_classifier_free_guidance and not guess_mode:
976
+ image = torch.cat([image] * 2)
977
+
978
+ return image
979
+
980
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps
981
+ def get_timesteps(self, num_inference_steps, strength, device):
982
+ # get the original timestep using init_timestep
983
+ init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
984
+
985
+ t_start = max(num_inference_steps - init_timestep, 0)
986
+ timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
987
+ if hasattr(self.scheduler, "set_begin_index"):
988
+ self.scheduler.set_begin_index(t_start * self.scheduler.order)
989
+
990
+ return timesteps, num_inference_steps - t_start
991
+
992
+ # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.prepare_latents
993
+ def prepare_latents(
994
+ self, image, timestep, batch_size, num_channels_latents, height, width, dtype, device, generator=None, add_noise=True, seed=None
995
+ ):
996
+
997
+ if image is None:
998
+ shape = (
999
+ batch_size,
1000
+ num_channels_latents,
1001
+ int(height) // self.vae_scale_factor,
1002
+ int(width) // self.vae_scale_factor,
1003
+ )
1004
+ init_latents = torch.zeros(shape, device=device, dtype=dtype)
1005
+ else:
1006
+ if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
1007
+ raise ValueError(
1008
+ f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
1009
+ )
1010
+
1011
+ latents_mean = latents_std = None
1012
+ if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None:
1013
+ latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1)
1014
+ if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None:
1015
+ latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1)
1016
+
1017
+ # Offload text encoder if `enable_model_cpu_offload` was enabled
1018
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
1019
+ self.text_encoder_2.to("cpu")
1020
+ torch.cuda.empty_cache()
1021
+
1022
+ image = image.to(device=device, dtype=dtype)
1023
+
1024
+ if image.shape[1] == 4:
1025
+ init_latents = image
1026
+
1027
+ else:
1028
+ # make sure the VAE is in float32 mode, as it overflows in float16
1029
+ if self.vae.config.force_upcast:
1030
+ image = image.float()
1031
+ self.vae.to(dtype=torch.float32)
1032
+
1033
+ if isinstance(generator, list) and len(generator) != batch_size:
1034
+ raise ValueError(
1035
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
1036
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
1037
+ )
1038
+
1039
+ elif isinstance(generator, list):
1040
+ init_latents = [
1041
+ retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
1042
+ for i in range(batch_size)
1043
+ ]
1044
+ init_latents = torch.cat(init_latents, dim=0)
1045
+ else:
1046
+ init_latents = retrieve_latents(self.vae.encode(image), generator=generator)
1047
+
1048
+ if self.vae.config.force_upcast:
1049
+ self.vae.to(dtype)
1050
+
1051
+ init_latents = init_latents.to(dtype)
1052
+ if latents_mean is not None and latents_std is not None:
1053
+ latents_mean = latents_mean.to(device=self.device, dtype=dtype)
1054
+ latents_std = latents_std.to(device=self.device, dtype=dtype)
1055
+ init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std
1056
+ else:
1057
+ init_latents = self.vae.config.scaling_factor * init_latents
1058
+
1059
+ if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
1060
+ # expand init_latents for batch_size
1061
+ additional_image_per_prompt = batch_size // init_latents.shape[0]
1062
+ init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
1063
+ elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
1064
+ raise ValueError(
1065
+ f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
1066
+ )
1067
+ else:
1068
+ init_latents = torch.cat([init_latents], dim=0)
1069
+
1070
+ if add_noise:
1071
+ if seed is not None:
1072
+ generator = torch.manual_seed(seed)
1073
+ noise = torch.randn(torch.Size(init_latents.shape), dtype=torch.float32, layout=torch.strided, generator=generator, device="cpu").to(device)
1074
+ init_latents = self.scheduler.add_noise(init_latents.to(device), noise, timestep)
1075
+ return init_latents.to(device, dtype=dtype)
1076
+
1077
+ latents = init_latents
1078
+
1079
+ return latents
1080
+
1081
+ # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids
1082
+ def _get_add_time_ids(
1083
+ self,
1084
+ original_size,
1085
+ crops_coords_top_left,
1086
+ target_size,
1087
+ aesthetic_score,
1088
+ negative_aesthetic_score,
1089
+ negative_original_size,
1090
+ negative_crops_coords_top_left,
1091
+ negative_target_size,
1092
+ dtype,
1093
+ text_encoder_projection_dim=None,
1094
+ ):
1095
+ if self.config.requires_aesthetics_score:
1096
+ add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,))
1097
+ add_neg_time_ids = list(
1098
+ negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,)
1099
+ )
1100
+ else:
1101
+ add_time_ids = list(original_size + crops_coords_top_left + target_size)
1102
+ add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size)
1103
+
1104
+ passed_add_embed_dim = (
1105
+ self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
1106
+ )
1107
+ expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
1108
+
1109
+ if (
1110
+ expected_add_embed_dim > passed_add_embed_dim
1111
+ and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_embed_dim
1112
+ ):
1113
+ raise ValueError(
1114
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model."
1115
+ )
1116
+ elif (
1117
+ expected_add_embed_dim < passed_add_embed_dim
1118
+ and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_embed_dim
1119
+ ):
1120
+ raise ValueError(
1121
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model."
1122
+ )
1123
+ elif expected_add_embed_dim != passed_add_embed_dim:
1124
+ raise ValueError(
1125
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
1126
+ )
1127
+
1128
+ add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
1129
+ add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype)
1130
+
1131
+ return add_time_ids, add_neg_time_ids
1132
+
1133
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
1134
+ def upcast_vae(self):
1135
+ dtype = self.vae.dtype
1136
+ self.vae.to(dtype=torch.float32)
1137
+ use_torch_2_0_or_xformers = isinstance(
1138
+ self.vae.decoder.mid_block.attentions[0].processor,
1139
+ (
1140
+ AttnProcessor2_0,
1141
+ XFormersAttnProcessor,
1142
+ LoRAXFormersAttnProcessor,
1143
+ LoRAAttnProcessor2_0,
1144
+ ),
1145
+ )
1146
+ # if xformers or torch_2_0 is used attention block does not need
1147
+ # to be in float32 which can save lots of memory
1148
+ if use_torch_2_0_or_xformers:
1149
+ self.vae.post_quant_conv.to(dtype)
1150
+ self.vae.decoder.conv_in.to(dtype)
1151
+ self.vae.decoder.mid_block.to(dtype)
1152
+
1153
+ @property
1154
+ def guidance_scale(self):
1155
+ return self._guidance_scale
1156
+
1157
+ @property
1158
+ def clip_skip(self):
1159
+ return self._clip_skip
1160
+
1161
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
1162
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
1163
+ # corresponds to doing no classifier free guidance.
1164
+ @property
1165
+ def do_classifier_free_guidance(self):
1166
+ return self._guidance_scale > 1
1167
+
1168
+ @property
1169
+ def cross_attention_kwargs(self):
1170
+ return self._cross_attention_kwargs
1171
+
1172
+ @property
1173
+ def num_timesteps(self):
1174
+ return self._num_timesteps
1175
+
1176
+ @torch.no_grad()
1177
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
1178
+ def __call__(
1179
+ self,
1180
+ prompt: Union[str, List[str]] = None,
1181
+ prompt_2: Optional[Union[str, List[str]]] = None,
1182
+ image: PipelineImageInput = None,
1183
+ control_image: PipelineImageInput = None,
1184
+ control_mask = None,
1185
+ identity_control_indices = None,
1186
+ height: Optional[int] = None,
1187
+ width: Optional[int] = None,
1188
+ strength: float = 0.8,
1189
+ num_inference_steps: int = 50,
1190
+ timesteps: Optional[List[int]] = None,
1191
+ sigmas: Optional[List[float]] = None,
1192
+ guidance_scale: float = 5.0,
1193
+ negative_prompt: Optional[Union[str, List[str]]] = None,
1194
+ negative_prompt_2: Optional[Union[str, List[str]]] = None,
1195
+ num_images_per_prompt: Optional[int] = 1,
1196
+ eta: float = 0.0,
1197
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
1198
+ seed: Optional[int] = None,
1199
+ latents: Optional[torch.FloatTensor] = None,
1200
+ prompt_embeds: Optional[torch.FloatTensor] = None,
1201
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
1202
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
1203
+ negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
1204
+ ip_adapter_image: Optional[PipelineImageInput] = None,
1205
+ ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
1206
+ output_type: Optional[str] = "pil",
1207
+ return_dict: bool = True,
1208
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
1209
+ controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
1210
+ guess_mode: bool = False,
1211
+ control_guidance_start: Union[float, List[float]] = 0.0,
1212
+ control_guidance_end: Union[float, List[float]] = 1.0,
1213
+ original_size: Tuple[int, int] = None,
1214
+ crops_coords_top_left: Tuple[int, int] = (0, 0),
1215
+ target_size: Tuple[int, int] = None,
1216
+ negative_original_size: Optional[Tuple[int, int]] = None,
1217
+ negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
1218
+ negative_target_size: Optional[Tuple[int, int]] = None,
1219
+ aesthetic_score: float = 6.0,
1220
+ negative_aesthetic_score: float = 2.5,
1221
+ clip_skip: Optional[int] = None,
1222
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
1223
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
1224
+ **kwargs,
1225
+ ):
1226
+ r"""
1227
+ Function invoked when calling the pipeline for generation.
1228
+
1229
+ Args:
1230
+ prompt (`str` or `List[str]`, *optional*):
1231
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
1232
+ instead.
1233
+ prompt_2 (`str` or `List[str]`, *optional*):
1234
+ The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
1235
+ used in both text-encoders
1236
+ image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
1237
+ `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
1238
+ The initial image will be used as the starting point for the image generation process. Can also accept
1239
+ image latents as `image`, if passing latents directly, it will not be encoded again.
1240
+ control_image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
1241
+ `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
1242
+ The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If
1243
+ the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can
1244
+ also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If
1245
+ height and/or width are passed, `image` is resized according to them. If multiple ControlNets are
1246
+ specified in init, images must be passed as a list such that each element of the list can be correctly
1247
+ batched for input to a single controlnet.
1248
+ height (`int`, *optional*, defaults to the size of control_image):
1249
+ The height in pixels of the generated image. Anything below 512 pixels won't work well for
1250
+ [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
1251
+ and checkpoints that are not specifically fine-tuned on low resolutions.
1252
+ width (`int`, *optional*, defaults to the size of control_image):
1253
+ The width in pixels of the generated image. Anything below 512 pixels won't work well for
1254
+ [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
1255
+ and checkpoints that are not specifically fine-tuned on low resolutions.
1256
+ strength (`float`, *optional*, defaults to 0.8):
1257
+ Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
1258
+ starting point and more noise is added the higher the `strength`. The number of denoising steps depends
1259
+ on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
1260
+ process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
1261
+ essentially ignores `image`.
1262
+ num_inference_steps (`int`, *optional*, defaults to 50):
1263
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
1264
+ expense of slower inference.
1265
+ guidance_scale (`float`, *optional*, defaults to 7.5):
1266
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
1267
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
1268
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1269
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
1270
+ usually at the expense of lower image quality.
1271
+ negative_prompt (`str` or `List[str]`, *optional*):
1272
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
1273
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
1274
+ less than `1`).
1275
+ negative_prompt_2 (`str` or `List[str]`, *optional*):
1276
+ The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
1277
+ `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
1278
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
1279
+ The number of images to generate per prompt.
1280
+ eta (`float`, *optional*, defaults to 0.0):
1281
+ Corresponds to parameter eta (Ξ·) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
1282
+ [`schedulers.DDIMScheduler`], will be ignored for others.
1283
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
1284
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
1285
+ to make generation deterministic.
1286
+ latents (`torch.FloatTensor`, *optional*):
1287
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
1288
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
1289
+ tensor will ge generated by sampling using the supplied random `generator`.
1290
+ prompt_embeds (`torch.FloatTensor`, *optional*):
1291
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
1292
+ provided, text embeddings will be generated from `prompt` input argument.
1293
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
1294
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
1295
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
1296
+ argument.
1297
+ pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
1298
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
1299
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
1300
+ negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
1301
+ Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
1302
+ weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
1303
+ input argument.
1304
+ ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
1305
+ ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
1306
+ Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
1307
+ IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
1308
+ contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
1309
+ provided, embeddings are computed from the `ip_adapter_image` input argument.
1310
+ output_type (`str`, *optional*, defaults to `"pil"`):
1311
+ The output format of the generate image. Choose between
1312
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
1313
+ return_dict (`bool`, *optional*, defaults to `True`):
1314
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
1315
+ plain tuple.
1316
+ cross_attention_kwargs (`dict`, *optional*):
1317
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
1318
+ `self.processor` in
1319
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
1320
+ controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
1321
+ The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
1322
+ to the residual in the original unet. If multiple ControlNets are specified in init, you can set the
1323
+ corresponding scale as a list.
1324
+ guess_mode (`bool`, *optional*, defaults to `False`):
1325
+ In this mode, the ControlNet encoder will try best to recognize the content of the input image even if
1326
+ you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.
1327
+ control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
1328
+ The percentage of total steps at which the controlnet starts applying.
1329
+ control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
1330
+ The percentage of total steps at which the controlnet stops applying.
1331
+ original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
1332
+ If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
1333
+ `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
1334
+ explained in section 2.2 of
1335
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
1336
+ crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
1337
+ `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
1338
+ `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
1339
+ `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
1340
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
1341
+ target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
1342
+ For most cases, `target_size` should be set to the desired height and width of the generated image. If
1343
+ not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
1344
+ section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
1345
+ negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
1346
+ To negatively condition the generation process based on a specific image resolution. Part of SDXL's
1347
+ micro-conditioning as explained in section 2.2 of
1348
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
1349
+ information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
1350
+ negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
1351
+ To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
1352
+ micro-conditioning as explained in section 2.2 of
1353
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
1354
+ information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
1355
+ negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
1356
+ To negatively condition the generation process based on a target image resolution. It should be as same
1357
+ as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
1358
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
1359
+ information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
1360
+ aesthetic_score (`float`, *optional*, defaults to 6.0):
1361
+ Used to simulate an aesthetic score of the generated image by influencing the positive text condition.
1362
+ Part of SDXL's micro-conditioning as explained in section 2.2 of
1363
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
1364
+ negative_aesthetic_score (`float`, *optional*, defaults to 2.5):
1365
+ Part of SDXL's micro-conditioning as explained in section 2.2 of
1366
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). Can be used to
1367
+ simulate an aesthetic score of the generated image by influencing the negative text condition.
1368
+ clip_skip (`int`, *optional*):
1369
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
1370
+ the output of the pre-final layer will be used for computing the prompt embeddings.
1371
+ callback_on_step_end (`Callable`, *optional*):
1372
+ A function that calls at the end of each denoising steps during the inference. The function is called
1373
+ with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
1374
+ callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
1375
+ `callback_on_step_end_tensor_inputs`.
1376
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
1377
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
1378
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
1379
+ `._callback_tensor_inputs` attribute of your pipeline class.
1380
+
1381
+ Examples:
1382
+
1383
+ Returns:
1384
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
1385
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple`
1386
+ containing the output images.
1387
+ """
1388
+
1389
+ callback = kwargs.pop("callback", None)
1390
+ callback_steps = kwargs.pop("callback_steps", None)
1391
+
1392
+ if callback is not None:
1393
+ deprecate(
1394
+ "callback",
1395
+ "1.0.0",
1396
+ "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
1397
+ )
1398
+ if callback_steps is not None:
1399
+ deprecate(
1400
+ "callback_steps",
1401
+ "1.0.0",
1402
+ "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
1403
+ )
1404
+
1405
+ controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet
1406
+
1407
+ # align format for control guidance
1408
+ if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
1409
+ control_guidance_start = len(control_guidance_end) * [control_guidance_start]
1410
+ elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):
1411
+ control_guidance_end = len(control_guidance_start) * [control_guidance_end]
1412
+ elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):
1413
+ mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1
1414
+ control_guidance_start, control_guidance_end = (
1415
+ mult * [control_guidance_start],
1416
+ mult * [control_guidance_end],
1417
+ )
1418
+
1419
+ # 1. Check inputs. Raise error if not correct
1420
+ self.check_inputs(
1421
+ prompt,
1422
+ prompt_2,
1423
+ control_image,
1424
+ strength,
1425
+ num_inference_steps,
1426
+ callback_steps,
1427
+ negative_prompt,
1428
+ negative_prompt_2,
1429
+ prompt_embeds,
1430
+ negative_prompt_embeds,
1431
+ pooled_prompt_embeds,
1432
+ negative_pooled_prompt_embeds,
1433
+ ip_adapter_image,
1434
+ ip_adapter_image_embeds,
1435
+ controlnet_conditioning_scale,
1436
+ control_guidance_start,
1437
+ control_guidance_end,
1438
+ callback_on_step_end_tensor_inputs,
1439
+ )
1440
+
1441
+ self._guidance_scale = guidance_scale
1442
+ self._clip_skip = clip_skip
1443
+ self._cross_attention_kwargs = cross_attention_kwargs
1444
+
1445
+ # 2. Define call parameters
1446
+ if prompt is not None and isinstance(prompt, str):
1447
+ batch_size = 1
1448
+ elif prompt is not None and isinstance(prompt, list):
1449
+ batch_size = len(prompt)
1450
+ else:
1451
+ batch_size = prompt_embeds.shape[0]
1452
+
1453
+ device = self._execution_device
1454
+
1455
+ if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):
1456
+ controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)
1457
+
1458
+ global_pool_conditions = (
1459
+ controlnet.config.global_pool_conditions
1460
+ if isinstance(controlnet, ControlNetModel)
1461
+ else controlnet.nets[0].config.global_pool_conditions
1462
+ )
1463
+ guess_mode = guess_mode or global_pool_conditions
1464
+
1465
+ # 3.1. Encode input prompt
1466
+ text_encoder_lora_scale = (
1467
+ self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
1468
+ )
1469
+ (
1470
+ prompt_embeds,
1471
+ negative_prompt_embeds,
1472
+ pooled_prompt_embeds,
1473
+ negative_pooled_prompt_embeds,
1474
+ ) = self.encode_prompt(
1475
+ prompt,
1476
+ prompt_2,
1477
+ device,
1478
+ num_images_per_prompt,
1479
+ self.do_classifier_free_guidance,
1480
+ negative_prompt,
1481
+ negative_prompt_2,
1482
+ prompt_embeds=prompt_embeds,
1483
+ negative_prompt_embeds=negative_prompt_embeds,
1484
+ pooled_prompt_embeds=pooled_prompt_embeds,
1485
+ negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
1486
+ lora_scale=text_encoder_lora_scale,
1487
+ clip_skip=self.clip_skip,
1488
+ )
1489
+
1490
+ # 3.2 Encode ip_adapter_image
1491
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
1492
+ image_embeds = self.prepare_ip_adapter_image_embeds(
1493
+ ip_adapter_image,
1494
+ ip_adapter_image_embeds,
1495
+ device,
1496
+ batch_size * num_images_per_prompt,
1497
+ self.do_classifier_free_guidance,
1498
+ )
1499
+
1500
+ # 4. Prepare image and controlnet_conditioning_image
1501
+ if image is not None:
1502
+ image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32)
1503
+ else:
1504
+ strength = 1.0
1505
+
1506
+ if isinstance(controlnet, ControlNetModel):
1507
+ control_image = self.prepare_control_image(
1508
+ image=control_image,
1509
+ width=width,
1510
+ height=height,
1511
+ batch_size=batch_size * num_images_per_prompt,
1512
+ num_images_per_prompt=num_images_per_prompt,
1513
+ device=device,
1514
+ dtype=controlnet.dtype,
1515
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
1516
+ guess_mode=guess_mode,
1517
+ )
1518
+ height, width = control_image.shape[-2:]
1519
+ elif isinstance(controlnet, MultiControlNetModel):
1520
+ control_images = []
1521
+
1522
+ for control_image_ in control_image:
1523
+ control_image_ = self.prepare_control_image(
1524
+ image=control_image_,
1525
+ width=width,
1526
+ height=height,
1527
+ batch_size=batch_size * num_images_per_prompt,
1528
+ num_images_per_prompt=num_images_per_prompt,
1529
+ device=device,
1530
+ dtype=controlnet.dtype,
1531
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
1532
+ guess_mode=guess_mode,
1533
+ )
1534
+
1535
+ control_images.append(control_image_)
1536
+
1537
+ control_image = control_images
1538
+ height, width = control_image[0].shape[-2:]
1539
+ else:
1540
+ assert False
1541
+
1542
+ # 4.1 Region control
1543
+ controlnet_masks = []
1544
+ if control_mask is not None:
1545
+ for mask in control_mask:
1546
+ mask = np.array(mask)
1547
+ mask_tensor = torch.from_numpy(mask).to(device=device, dtype=prompt_embeds.dtype)
1548
+ mask_tensor = mask_tensor[:, :, 0] / 255.
1549
+ mask_tensor = mask_tensor[None, None]
1550
+ h, w = mask_tensor.shape[-2:]
1551
+ control_mask_list = []
1552
+ for scale in [8, 8, 8, 16, 16, 16, 32, 32, 32]:
1553
+ # Python uses IEEE 754 rounding rules, we need to add a small value to round like the unet model
1554
+ w_n = round((w + 0.01) / 8)
1555
+ h_n = round((h + 0.01) / 8)
1556
+ if scale in [16, 32]:
1557
+ w_n = round((w_n + 0.01) / 2)
1558
+ h_n = round((h_n + 0.01) / 2)
1559
+ if scale == 32:
1560
+ w_n = round((w_n + 0.01) / 2)
1561
+ h_n = round((h_n + 0.01) / 2)
1562
+ scale_mask_weight_image_tensor = F.interpolate(
1563
+ mask_tensor,(h_n, w_n), mode='bilinear')
1564
+ control_mask_list.append(scale_mask_weight_image_tensor)
1565
+ controlnet_masks.append(control_mask_list)
1566
+
1567
+ # 5. Prepare timesteps
1568
+ full_num_inference_steps = int(num_inference_steps / strength) if strength > 0 else num_inference_steps
1569
+
1570
+ if timesteps is None:
1571
+ self.scheduler.set_timesteps(full_num_inference_steps + 1, device=device)
1572
+ sigmas = self._loglinear_interp(self.ays_noise_sigmas["SDXL"], full_num_inference_steps + 1)
1573
+ sigmas[-1] = 0
1574
+ log_sigmas = np.log(np.array((1 - self.scheduler.alphas_cumprod) / self.scheduler.alphas_cumprod) ** 0.5)
1575
+ timesteps = np.array([self.scheduler._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
1576
+ timesteps = timesteps[-(num_inference_steps + 1):-1]
1577
+ if hasattr(self.scheduler, "sigmas"):
1578
+ self.scheduler.sigmas = torch.from_numpy(sigmas)[-(num_inference_steps + 1):]
1579
+ self.scheduler.timesteps = torch.from_numpy(timesteps).to(self.device, dtype=torch.int64)
1580
+ self.scheduler.num_inference_steps = len(self.scheduler.timesteps)
1581
+
1582
+ else:
1583
+ if "timesteps" in inspect.signature(self.scheduler.set_timesteps).parameters:
1584
+ self.scheduler.set_timesteps(full_num_inference_steps + 1, timesteps=timesteps, device=device)
1585
+ else:
1586
+ self.scheduler.set_timesteps(full_num_inference_steps + 1, device=device)
1587
+
1588
+ latent_timestep = self.scheduler.timesteps[:1].repeat(batch_size * num_images_per_prompt)
1589
+ self._num_timesteps = len(self.scheduler.timesteps)
1590
+
1591
+ # 6. Prepare latent variables
1592
+ if latents is None:
1593
+ num_channels_latents = self.unet.config.in_channels
1594
+ latents = self.prepare_latents(
1595
+ image,
1596
+ latent_timestep,
1597
+ batch_size * num_images_per_prompt,
1598
+ num_channels_latents,
1599
+ height,
1600
+ width,
1601
+ prompt_embeds.dtype,
1602
+ device,
1603
+ generator,
1604
+ True,
1605
+ seed
1606
+ )
1607
+
1608
+ if hasattr(self.scheduler, "sigmas"):
1609
+ sigmas = self.scheduler.sigmas
1610
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
1611
+ seeds = [seed] * len(latents) if seed is not None else generator.seed()
1612
+ brownian_tree_noise_sampler = BrownianTreeNoiseSampler(latents, sigma_min, sigma_max, seed=seeds, cpu=False)
1613
+ else:
1614
+ brownian_tree_noise_sampler = None
1615
+
1616
+ # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
1617
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
1618
+
1619
+ # 7.1 Create tensor stating which controlnets to keep
1620
+ controlnet_keep = []
1621
+ for i in range(len(timesteps)):
1622
+ keeps = [
1623
+ 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
1624
+ for s, e in zip(control_guidance_start, control_guidance_end)
1625
+ ]
1626
+ controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps)
1627
+
1628
+ # 7.2 Prepare added time ids & embeddings
1629
+ if isinstance(control_image, list):
1630
+ original_size = original_size or control_image[0].shape[-2:]
1631
+ else:
1632
+ original_size = original_size or control_image.shape[-2:]
1633
+ target_size = target_size or (height, width)
1634
+
1635
+ if negative_original_size is None:
1636
+ negative_original_size = original_size
1637
+ if negative_target_size is None:
1638
+ negative_target_size = target_size
1639
+ add_text_embeds = pooled_prompt_embeds
1640
+
1641
+ if self.text_encoder_2 is None:
1642
+ text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
1643
+ else:
1644
+ text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
1645
+
1646
+ add_time_ids, add_neg_time_ids = self._get_add_time_ids(
1647
+ original_size,
1648
+ crops_coords_top_left,
1649
+ target_size,
1650
+ aesthetic_score,
1651
+ negative_aesthetic_score,
1652
+ negative_original_size,
1653
+ negative_crops_coords_top_left,
1654
+ negative_target_size,
1655
+ dtype=prompt_embeds.dtype,
1656
+ text_encoder_projection_dim=text_encoder_projection_dim,
1657
+ )
1658
+ add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1)
1659
+
1660
+ if self.do_classifier_free_guidance:
1661
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
1662
+ add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
1663
+ add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1)
1664
+ add_time_ids = torch.cat([add_neg_time_ids, add_time_ids], dim=0)
1665
+
1666
+ prompt_embeds = prompt_embeds.to(device)
1667
+ add_text_embeds = add_text_embeds.to(device)
1668
+ add_time_ids = add_time_ids.to(device)
1669
+
1670
+ # 8. Denoising loop
1671
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
1672
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
1673
+ for i, t in enumerate(timesteps):
1674
+ # expand the latents if we are doing classifier free guidance
1675
+ latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
1676
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
1677
+
1678
+ added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
1679
+
1680
+ # controlnet(s) inference
1681
+ if guess_mode and self.do_classifier_free_guidance:
1682
+ # Infer ControlNet only for the conditional batch.
1683
+ control_model_input = latents
1684
+ control_model_input = self.scheduler.scale_model_input(control_model_input, t)
1685
+ controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]
1686
+ controlnet_added_cond_kwargs = {
1687
+ "text_embeds": add_text_embeds.chunk(2)[1],
1688
+ "time_ids": add_time_ids.chunk(2)[1],
1689
+ }
1690
+ else:
1691
+ control_model_input = latent_model_input
1692
+ controlnet_prompt_embeds = prompt_embeds
1693
+ controlnet_added_cond_kwargs = added_cond_kwargs
1694
+
1695
+ if isinstance(controlnet_keep[i], list):
1696
+ cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]
1697
+ else:
1698
+ controlnet_cond_scale = controlnet_conditioning_scale
1699
+ if isinstance(controlnet_cond_scale, list):
1700
+ controlnet_cond_scale = controlnet_cond_scale[0]
1701
+ cond_scale = controlnet_cond_scale * controlnet_keep[i]
1702
+
1703
+ if ip_adapter_image_embeds is None and ip_adapter_image is not None:
1704
+ encoder_hidden_states = self.unet.process_encoder_hidden_states(prompt_embeds, {"image_embeds": image_embeds})
1705
+ ip_adapter_image_embeds = encoder_hidden_states[1]
1706
+
1707
+ down_block_res_samples = None
1708
+ mid_block_res_sample = None
1709
+
1710
+ for controlnet_index in range(len(self.controlnet.nets)):
1711
+ ip_adapter_index = next((y for x, y in identity_control_indices if x == controlnet_index), None)
1712
+ if ip_adapter_index is not None:
1713
+ control_prompt_embeds = ip_adapter_image_embeds[ip_adapter_index].squeeze(1)
1714
+ else:
1715
+ control_prompt_embeds = controlnet_prompt_embeds
1716
+ down_samples, mid_sample = self.controlnet.nets[controlnet_index](
1717
+ control_model_input,
1718
+ t,
1719
+ encoder_hidden_states=control_prompt_embeds,
1720
+ controlnet_cond=control_image[controlnet_index],
1721
+ conditioning_scale=cond_scale[controlnet_index],
1722
+ guess_mode=guess_mode,
1723
+ added_cond_kwargs=controlnet_added_cond_kwargs,
1724
+ return_dict=False,
1725
+ )
1726
+
1727
+ if len(controlnet_masks) > controlnet_index and controlnet_masks[controlnet_index] is not None:
1728
+ down_samples = [
1729
+ down_sample * mask_weight
1730
+ for down_sample, mask_weight in zip(down_samples, controlnet_masks[controlnet_index])
1731
+ ]
1732
+ mid_sample *= controlnet_masks[controlnet_index][-1]
1733
+
1734
+ if down_block_res_samples is None and mid_block_res_sample is None:
1735
+ down_block_res_samples = down_samples
1736
+ mid_block_res_sample = mid_sample
1737
+ else:
1738
+ down_block_res_samples = [
1739
+ samples_prev + samples_curr
1740
+ for samples_prev, samples_curr in zip(down_block_res_samples, down_samples)
1741
+ ]
1742
+ mid_block_res_sample += mid_sample
1743
+
1744
+ if guess_mode and self.do_classifier_free_guidance:
1745
+ # Infered ControlNet only for the conditional batch.
1746
+ # To apply the output of ControlNet to both the unconditional and conditional batches,
1747
+ # add 0 to the unconditional batch to keep it unchanged.
1748
+ down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]
1749
+ mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])
1750
+
1751
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
1752
+ added_cond_kwargs["image_embeds"] = image_embeds
1753
+
1754
+ # predict the noise residual
1755
+ noise_pred = self.unet(
1756
+ latent_model_input,
1757
+ t,
1758
+ encoder_hidden_states=prompt_embeds,
1759
+ cross_attention_kwargs=self.cross_attention_kwargs,
1760
+ down_block_additional_residuals=down_block_res_samples,
1761
+ mid_block_additional_residual=mid_block_res_sample,
1762
+ added_cond_kwargs=added_cond_kwargs,
1763
+ return_dict=False,
1764
+ )[0]
1765
+
1766
+ # perform guidance
1767
+ if self.do_classifier_free_guidance:
1768
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
1769
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
1770
+
1771
+ # compute the previous noisy sample x_t -> x_t-1
1772
+ # check if scheduler.step supports variance noise
1773
+ if "variance_noise" in inspect.signature(self.scheduler.step).parameters and brownian_tree_noise_sampler is not None:
1774
+ sigmas = self.scheduler.sigmas
1775
+ noise = brownian_tree_noise_sampler(sigmas[i], sigmas[i + 1]).to(device=device, dtype=latents.dtype)
1776
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False, variance_noise=noise)[0]
1777
+ else:
1778
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
1779
+
1780
+ if callback_on_step_end is not None:
1781
+ callback_kwargs = {}
1782
+ for k in callback_on_step_end_tensor_inputs:
1783
+ callback_kwargs[k] = locals()[k]
1784
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
1785
+
1786
+ latents = callback_outputs.pop("latents", latents)
1787
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
1788
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
1789
+
1790
+ # call the callback, if provided
1791
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
1792
+ progress_bar.update()
1793
+ if callback is not None and i % callback_steps == 0:
1794
+ step_idx = i // getattr(self.scheduler, "order", 1)
1795
+ callback(step_idx, t, latents)
1796
+
1797
+ # If we do sequential model offloading, let's offload unet and controlnet
1798
+ # manually for max memory savings
1799
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
1800
+ self.unet.to("cpu")
1801
+ self.controlnet.to("cpu")
1802
+ torch.cuda.empty_cache()
1803
+
1804
+ if not output_type == "latent":
1805
+ # make sure the VAE is in float32 mode, as it overflows in float16
1806
+ needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
1807
+
1808
+ if needs_upcasting:
1809
+ self.upcast_vae()
1810
+ latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
1811
+
1812
+ # unscale/denormalize the latents
1813
+ # denormalize with the mean and std if available and not None
1814
+ has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None
1815
+ has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None
1816
+ if has_latents_mean and has_latents_std:
1817
+ latents_mean = (
1818
+ torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)
1819
+ )
1820
+ latents_std = (
1821
+ torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)
1822
+ )
1823
+ latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean
1824
+ else:
1825
+ latents = latents / self.vae.config.scaling_factor
1826
+
1827
+ image = self.vae.decode(latents, return_dict=False)[0]
1828
+
1829
+ # cast back to fp16 if needed
1830
+ if needs_upcasting:
1831
+ self.vae.to(dtype=torch.float16)
1832
+ else:
1833
+ image = latents
1834
+ return StableDiffusionXLPipelineOutput(images=image)
1835
+
1836
+ # apply watermark if available
1837
+ if self.watermark is not None:
1838
+ image = self.watermark.apply_watermark(image)
1839
+
1840
+ image = self.image_processor.postprocess(image, output_type=output_type)
1841
+
1842
+ # Offload all models
1843
+ self.maybe_free_model_hooks()
1844
+
1845
+ if not return_dict:
1846
+ return (image,)
1847
+
1848
+ return StableDiffusionXLPipelineOutput(images=image)
requirements.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # diffusers @ git+https://github.com/huggingface/diffusers.git@0d7c4790235ac00b4524b492bc2a680dcc5cf6b0
2
+ controlnet_aux==0.0.8
3
+ huggingface_hub==0.20.2
4
+ insightface==0.7.3
5
+ numpy==1.26.2
6
+ opencv_contrib_python==4.9.0.80
7
+ opencv_python==4.9.0.80
8
+ opencv_python_headless==4.7.0.72
9
+ Pillow==10.1.0
10
+ torch==2.0.0
11
+ torchsde==0.2.6
12
+ transformers==4.37.1
utils.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import PIL
3
+ import cv2
4
+ import numpy as np
5
+
6
+ from diffusers.utils import load_image
7
+
8
+ def draw_kps(image_pil, kps, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]):
9
+ stickwidth = 4
10
+ limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]])
11
+ kps = np.array(kps)
12
+
13
+ # w, h = image_pil.size
14
+ # out_img = np.zeros([h, w, 3])
15
+ if type(image_pil) == PIL.Image.Image:
16
+ out_img = np.array(image_pil)
17
+ else:
18
+ out_img = image_pil
19
+
20
+ for i in range(len(limbSeq)):
21
+ index = limbSeq[i]
22
+ color = color_list[index[0]]
23
+
24
+ x = kps[index][:, 0]
25
+ y = kps[index][:, 1]
26
+ length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5
27
+ angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1]))
28
+ polygon = cv2.ellipse2Poly(
29
+ (int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1
30
+ )
31
+ out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color)
32
+ out_img = (out_img * 0.6).astype(np.uint8)
33
+
34
+ for idx_kp, kp in enumerate(kps):
35
+ color = color_list[idx_kp]
36
+ x, y = kp
37
+ out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1)
38
+
39
+ out_img_pil = PIL.Image.fromarray(out_img.astype(np.uint8))
40
+ return out_img_pil
41
+
42
+
43
+ def load_and_resize_image(image_path, max_width, max_height, maintain_aspect_ratio=True):
44
+ # Open the image
45
+ # image = Image.open(image_path)
46
+ image = load_image(image_path)
47
+
48
+ # Get the current width and height of the image
49
+ current_width, current_height = image.size
50
+
51
+ if maintain_aspect_ratio:
52
+ # Calculate the aspect ratio of the image
53
+ aspect_ratio = current_width / current_height
54
+
55
+ # Calculate the new dimensions based on the max width and height
56
+ if current_width / max_width > current_height / max_height:
57
+ new_width = max_width
58
+ new_height = int(new_width / aspect_ratio)
59
+ else:
60
+ new_height = max_height
61
+ new_width = int(new_height * aspect_ratio)
62
+ else:
63
+ # Use the max width and height as the new dimensions
64
+ new_width = max_width
65
+ new_height = max_height
66
+
67
+ # Ensure the new dimensions are divisible by 8
68
+ new_width = (new_width // 8) * 8
69
+ new_height = (new_height // 8) * 8
70
+
71
+ # Resize the image
72
+ resized_image = image.resize((new_width, new_height))
73
+
74
+ return resized_image
75
+
76
+ from PIL import Image
77
+
78
+ def align_images(image1, image2):
79
+ """
80
+ Resize two images to the same dimensions by cropping the larger image(s) to match the smaller one.
81
+
82
+ Args:
83
+ image1 (PIL.Image): First image to be aligned.
84
+ image2 (PIL.Image): Second image to be aligned.
85
+
86
+ Returns:
87
+ tuple: A tuple containing two images with the same dimensions.
88
+ """
89
+ # Determine the new size by taking the smaller width and height from both images
90
+ new_width = min(image1.size[0], image2.size[0])
91
+ new_height = min(image1.size[1], image2.size[1])
92
+
93
+ # Crop both images if necessary
94
+ if image1.size != (new_width, new_height):
95
+ image1 = image1.crop((0, 0, new_width, new_height))
96
+ if image2.size != (new_width, new_height):
97
+ image2 = image2.crop((0, 0, new_width, new_height))
98
+
99
+ return image1, image2