optimum-rbln 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- optimum/rbln/__init__.py +115 -0
- optimum/rbln/__version__.py +1 -0
- optimum/rbln/diffusers/__init__.py +64 -0
- optimum/rbln/diffusers/models/__init__.py +26 -0
- optimum/rbln/diffusers/models/autoencoder_kl.py +313 -0
- optimum/rbln/diffusers/models/controlnet.py +180 -0
- optimum/rbln/diffusers/models/unet_2d_condition.py +352 -0
- optimum/rbln/diffusers/pipelines/__init__.py +30 -0
- optimum/rbln/diffusers/pipelines/controlnet/__init__.py +24 -0
- optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +266 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion/__init__.py +26 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_controlnet_img2img.py +731 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +106 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +116 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion_xl/__init__.py +2 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +109 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +111 -0
- optimum/rbln/modeling.py +0 -0
- optimum/rbln/modeling_alias.py +49 -0
- optimum/rbln/modeling_base.py +645 -0
- optimum/rbln/modeling_config.py +169 -0
- optimum/rbln/modeling_seq2seq.py +469 -0
- optimum/rbln/transformers/__init__.py +59 -0
- optimum/rbln/transformers/generation/__init__.py +24 -0
- optimum/rbln/transformers/generation/streamers.py +122 -0
- optimum/rbln/transformers/models/__init__.py +28 -0
- optimum/rbln/transformers/models/bart/__init__.py +24 -0
- optimum/rbln/transformers/models/bart/bart_architecture.py +377 -0
- optimum/rbln/transformers/models/clip/__init__.py +24 -0
- optimum/rbln/transformers/models/clip/modeling_clip.py +116 -0
- optimum/rbln/transformers/models/gpt2/__init__.py +24 -0
- optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +253 -0
- optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +700 -0
- optimum/rbln/transformers/models/llama/__init__.py +24 -0
- optimum/rbln/transformers/models/llama/llama_architecture.py +607 -0
- optimum/rbln/transformers/models/llama/modeling_llama.py +409 -0
- optimum/rbln/transformers/models/t5/__init__.py +24 -0
- optimum/rbln/transformers/models/t5/t5_architecture.py +439 -0
- optimum/rbln/transformers/models/wav2vec2/__init__.py +24 -0
- optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py +121 -0
- optimum/rbln/transformers/models/whisper/__init__.py +24 -0
- optimum/rbln/transformers/models/whisper/modeling_whisper.py +374 -0
- optimum/rbln/transformers/models/whisper/whisper_architecture.py +406 -0
- optimum/rbln/utils/__init__.py +25 -0
- optimum/rbln/utils/import_utils.py +28 -0
- optimum/rbln/utils/runtime_utils.py +71 -0
- optimum/rbln/utils/save_utils.py +92 -0
- optimum_rbln-0.1.0.dist-info/METADATA +144 -0
- optimum_rbln-0.1.0.dist-info/RECORD +51 -0
- optimum_rbln-0.1.0.dist-info/WHEEL +4 -0
- optimum_rbln-0.1.0.dist-info/licenses/LICENSE +201 -0
@@ -0,0 +1,731 @@
|
|
1
|
+
# Copyright 2024 Rebellions Inc.
|
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
|
+
# Portions of this software are licensed under the Apache License,
|
16
|
+
# Version 2.0. See the NOTICE file distributed with this work for
|
17
|
+
# additional information regarding copyright ownership.
|
18
|
+
|
19
|
+
# All other portions of this software, including proprietary code,
|
20
|
+
# are the intellectual property of Rebellions Inc. and may not be
|
21
|
+
# copied, modified, or distributed without prior written permission
|
22
|
+
# from Rebellions Inc.
|
23
|
+
"""RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
|
24
|
+
|
25
|
+
from pathlib import Path
|
26
|
+
from tempfile import TemporaryDirectory
|
27
|
+
from typing import Any, Callable, Dict, List, Optional, Union
|
28
|
+
|
29
|
+
import torch
|
30
|
+
import torch.nn.functional as F
|
31
|
+
from diffusers import StableDiffusionControlNetImg2ImgPipeline
|
32
|
+
from diffusers.image_processor import PipelineImageInput
|
33
|
+
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
|
34
|
+
from diffusers.utils import deprecate, logging
|
35
|
+
from diffusers.utils.torch_utils import is_compiled_module
|
36
|
+
|
37
|
+
from ....modeling_base import RBLNBaseModel
|
38
|
+
from ....transformers import RBLNCLIPTextModel
|
39
|
+
from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
|
40
|
+
from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
|
41
|
+
|
42
|
+
|
43
|
+
logger = logging.get_logger(__name__)
|
44
|
+
|
45
|
+
|
46
|
+
class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2ImgPipeline):
|
47
|
+
@classmethod
|
48
|
+
def from_pretrained(cls, model_id, **kwargs):
|
49
|
+
export = kwargs.pop("export", None)
|
50
|
+
text_encoder = kwargs.pop("text_encoder", None)
|
51
|
+
controlnets = kwargs.pop("controlnet", None)
|
52
|
+
|
53
|
+
kwargs_dict = {
|
54
|
+
"pretrained_model_name_or_path": model_id,
|
55
|
+
"text_encoder": text_encoder,
|
56
|
+
"controlnet": controlnets,
|
57
|
+
**kwargs,
|
58
|
+
}
|
59
|
+
model = super().from_pretrained(**{k: v for k, v in kwargs_dict.items() if v is not None})
|
60
|
+
|
61
|
+
if export is None or export is False:
|
62
|
+
return model
|
63
|
+
|
64
|
+
rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
|
65
|
+
|
66
|
+
rbln_callback = rbln_config_kwargs.pop("rbln_callback", None)
|
67
|
+
img_width = kwargs.pop("img_width", None)
|
68
|
+
img_height = kwargs.pop("img_height", None)
|
69
|
+
|
70
|
+
model = rbln_callback(model)
|
71
|
+
|
72
|
+
save_dir = TemporaryDirectory()
|
73
|
+
save_dir_path = Path(save_dir.name)
|
74
|
+
|
75
|
+
model.save_pretrained(save_directory=save_dir_path, **kwargs)
|
76
|
+
|
77
|
+
# compile model, create runtime
|
78
|
+
vae = RBLNAutoencoderKL.from_pretrained(
|
79
|
+
model_id=save_dir_path / "vae",
|
80
|
+
export=True,
|
81
|
+
rbln_unet_sample_size=model.unet.config.sample_size,
|
82
|
+
rbln_use_encode=True,
|
83
|
+
rbln_img_width=img_width,
|
84
|
+
rbln_img_height=img_height,
|
85
|
+
rbln_vae_scale_factor=model.vae_scale_factor,
|
86
|
+
**rbln_config_kwargs,
|
87
|
+
**rbln_constructor_kwargs,
|
88
|
+
)
|
89
|
+
|
90
|
+
text_encoder = RBLNCLIPTextModel.from_pretrained(
|
91
|
+
model_id=save_dir_path / "text_encoder",
|
92
|
+
export=True,
|
93
|
+
**rbln_config_kwargs,
|
94
|
+
**rbln_constructor_kwargs,
|
95
|
+
)
|
96
|
+
|
97
|
+
batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
|
98
|
+
unet_batch_size = batch_size * 2
|
99
|
+
|
100
|
+
unet = RBLNUNet2DConditionModel.from_pretrained(
|
101
|
+
model_id=save_dir_path / "unet",
|
102
|
+
export=True,
|
103
|
+
rbln_max_seq_len=text_encoder.config.max_position_embeddings,
|
104
|
+
rbln_batch_size=unet_batch_size,
|
105
|
+
rbln_use_encode=True,
|
106
|
+
rbln_img_width=img_width,
|
107
|
+
rbln_img_height=img_height,
|
108
|
+
rbln_vae_scale_factor=model.vae_scale_factor,
|
109
|
+
rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
|
110
|
+
**rbln_config_kwargs,
|
111
|
+
**rbln_constructor_kwargs,
|
112
|
+
)
|
113
|
+
|
114
|
+
if isinstance(controlnets, (list, tuple)):
|
115
|
+
controlnet = RBLNMultiControlNetModel.from_pretrained(
|
116
|
+
model_id=str(save_dir_path / "controlnet"),
|
117
|
+
export=True,
|
118
|
+
rbln_batch_size=unet_batch_size,
|
119
|
+
rbln_img_width=img_width,
|
120
|
+
rbln_img_height=img_height,
|
121
|
+
rbln_vae_scale_factor=model.vae_scale_factor,
|
122
|
+
**rbln_config_kwargs,
|
123
|
+
**rbln_constructor_kwargs,
|
124
|
+
)
|
125
|
+
controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
|
126
|
+
else:
|
127
|
+
controlnet = RBLNControlNetModel.from_pretrained(
|
128
|
+
model_id=save_dir_path / "controlnet",
|
129
|
+
export=True,
|
130
|
+
rbln_batch_size=unet_batch_size,
|
131
|
+
rbln_img_width=img_width,
|
132
|
+
rbln_img_height=img_height,
|
133
|
+
rbln_vae_scale_factor=model.vae_scale_factor,
|
134
|
+
**rbln_config_kwargs,
|
135
|
+
**rbln_constructor_kwargs,
|
136
|
+
)
|
137
|
+
controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
|
138
|
+
|
139
|
+
# replace modules
|
140
|
+
model.vae = vae
|
141
|
+
model.text_encoder = text_encoder
|
142
|
+
model.unet = unet
|
143
|
+
model.controlnet = controlnet
|
144
|
+
|
145
|
+
# update config to be able to load from file.
|
146
|
+
update_dict = {
|
147
|
+
"vae": ("optimum.rbln", "RBLNAutoencoderKL"),
|
148
|
+
"text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
|
149
|
+
"unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
|
150
|
+
"controlnet": controlnet_dict,
|
151
|
+
}
|
152
|
+
model.register_to_config(**update_dict)
|
153
|
+
|
154
|
+
return model
|
155
|
+
|
156
|
+
def check_inputs(
|
157
|
+
self,
|
158
|
+
prompt,
|
159
|
+
image,
|
160
|
+
callback_steps,
|
161
|
+
negative_prompt=None,
|
162
|
+
prompt_embeds=None,
|
163
|
+
negative_prompt_embeds=None,
|
164
|
+
ip_adapter_image=None,
|
165
|
+
ip_adapter_image_embeds=None,
|
166
|
+
controlnet_conditioning_scale=1.0,
|
167
|
+
control_guidance_start=0.0,
|
168
|
+
control_guidance_end=1.0,
|
169
|
+
callback_on_step_end_tensor_inputs=None,
|
170
|
+
):
|
171
|
+
if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
|
172
|
+
raise ValueError(
|
173
|
+
f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
|
174
|
+
f" {type(callback_steps)}."
|
175
|
+
)
|
176
|
+
|
177
|
+
if callback_on_step_end_tensor_inputs is not None and not all(
|
178
|
+
k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
|
179
|
+
):
|
180
|
+
raise ValueError(
|
181
|
+
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]}"
|
182
|
+
)
|
183
|
+
|
184
|
+
if prompt is not None and prompt_embeds is not None:
|
185
|
+
raise ValueError(
|
186
|
+
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
187
|
+
" only forward one of the two."
|
188
|
+
)
|
189
|
+
elif prompt is None and prompt_embeds is None:
|
190
|
+
raise ValueError(
|
191
|
+
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
192
|
+
)
|
193
|
+
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
194
|
+
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
195
|
+
|
196
|
+
if negative_prompt is not None and negative_prompt_embeds is not None:
|
197
|
+
raise ValueError(
|
198
|
+
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
|
199
|
+
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
|
200
|
+
)
|
201
|
+
|
202
|
+
if prompt_embeds is not None and negative_prompt_embeds is not None:
|
203
|
+
if prompt_embeds.shape != negative_prompt_embeds.shape:
|
204
|
+
raise ValueError(
|
205
|
+
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
|
206
|
+
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
|
207
|
+
f" {negative_prompt_embeds.shape}."
|
208
|
+
)
|
209
|
+
|
210
|
+
# `prompt` needs more sophisticated handling when there are multiple
|
211
|
+
# conditionings.
|
212
|
+
if isinstance(self.controlnet, RBLNMultiControlNetModel):
|
213
|
+
if isinstance(prompt, list):
|
214
|
+
logger.warning(
|
215
|
+
f"You have {len(self.controlnet.nets)} ControlNets and you have passed {len(prompt)}"
|
216
|
+
" prompts. The conditionings will be fixed across the prompts."
|
217
|
+
)
|
218
|
+
|
219
|
+
# Check `image`
|
220
|
+
is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
|
221
|
+
self.controlnet, torch._dynamo.eval_frame.OptimizedModule
|
222
|
+
)
|
223
|
+
if (
|
224
|
+
isinstance(self.controlnet, RBLNControlNetModel)
|
225
|
+
or is_compiled
|
226
|
+
and isinstance(self.controlnet._orig_mod, RBLNControlNetModel)
|
227
|
+
):
|
228
|
+
self.check_image(image, prompt, prompt_embeds)
|
229
|
+
elif (
|
230
|
+
isinstance(self.controlnet, RBLNMultiControlNetModel)
|
231
|
+
or is_compiled
|
232
|
+
and isinstance(self.controlnet._orig_mod, RBLNMultiControlNetModel)
|
233
|
+
):
|
234
|
+
if not isinstance(image, list):
|
235
|
+
raise TypeError("For multiple controlnets: `image` must be type `list`")
|
236
|
+
|
237
|
+
# When `image` is a nested list:
|
238
|
+
# (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]])
|
239
|
+
elif any(isinstance(i, list) for i in image):
|
240
|
+
raise ValueError("A single batch of multiple conditionings are supported at the moment.")
|
241
|
+
elif len(image) != len(self.controlnet.nets):
|
242
|
+
raise ValueError(
|
243
|
+
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."
|
244
|
+
)
|
245
|
+
|
246
|
+
for image_ in image:
|
247
|
+
self.check_image(image_, prompt, prompt_embeds)
|
248
|
+
else:
|
249
|
+
assert False
|
250
|
+
|
251
|
+
# Check `controlnet_conditioning_scale`
|
252
|
+
if (
|
253
|
+
isinstance(self.controlnet, RBLNControlNetModel)
|
254
|
+
or is_compiled
|
255
|
+
and isinstance(self.controlnet._orig_mod, RBLNControlNetModel)
|
256
|
+
):
|
257
|
+
if not isinstance(controlnet_conditioning_scale, float):
|
258
|
+
raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")
|
259
|
+
elif (
|
260
|
+
isinstance(self.controlnet, RBLNMultiControlNetModel)
|
261
|
+
or is_compiled
|
262
|
+
and isinstance(self.controlnet._orig_mod, RBLNMultiControlNetModel)
|
263
|
+
):
|
264
|
+
if isinstance(controlnet_conditioning_scale, list):
|
265
|
+
if any(isinstance(i, list) for i in controlnet_conditioning_scale):
|
266
|
+
raise ValueError("A single batch of multiple conditionings are supported at the moment.")
|
267
|
+
elif isinstance(controlnet_conditioning_scale, list) and len(controlnet_conditioning_scale) != len(
|
268
|
+
self.controlnet.nets
|
269
|
+
):
|
270
|
+
raise ValueError(
|
271
|
+
"For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"
|
272
|
+
" the same length as the number of controlnets"
|
273
|
+
)
|
274
|
+
else:
|
275
|
+
assert False
|
276
|
+
|
277
|
+
if len(control_guidance_start) != len(control_guidance_end):
|
278
|
+
raise ValueError(
|
279
|
+
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."
|
280
|
+
)
|
281
|
+
|
282
|
+
if isinstance(self.controlnet, RBLNMultiControlNetModel):
|
283
|
+
if len(control_guidance_start) != len(self.controlnet.nets):
|
284
|
+
raise ValueError(
|
285
|
+
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)}."
|
286
|
+
)
|
287
|
+
|
288
|
+
for start, end in zip(control_guidance_start, control_guidance_end):
|
289
|
+
if start >= end:
|
290
|
+
raise ValueError(
|
291
|
+
f"control guidance start: {start} cannot be larger or equal to control guidance end: {end}."
|
292
|
+
)
|
293
|
+
if start < 0.0:
|
294
|
+
raise ValueError(f"control guidance start: {start} can't be smaller than 0.")
|
295
|
+
if end > 1.0:
|
296
|
+
raise ValueError(f"control guidance end: {end} can't be larger than 1.0.")
|
297
|
+
|
298
|
+
if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
|
299
|
+
raise ValueError(
|
300
|
+
"Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
|
301
|
+
)
|
302
|
+
|
303
|
+
if ip_adapter_image_embeds is not None:
|
304
|
+
if not isinstance(ip_adapter_image_embeds, list):
|
305
|
+
raise ValueError(
|
306
|
+
f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
|
307
|
+
)
|
308
|
+
elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
|
309
|
+
raise ValueError(
|
310
|
+
f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
|
311
|
+
)
|
312
|
+
|
313
|
+
@torch.no_grad()
|
314
|
+
def __call__(
|
315
|
+
self,
|
316
|
+
prompt: Union[str, List[str]] = None,
|
317
|
+
image: PipelineImageInput = None,
|
318
|
+
control_image: PipelineImageInput = None,
|
319
|
+
height: Optional[int] = None,
|
320
|
+
width: Optional[int] = None,
|
321
|
+
strength: float = 0.8,
|
322
|
+
num_inference_steps: int = 50,
|
323
|
+
guidance_scale: float = 7.5,
|
324
|
+
negative_prompt: Optional[Union[str, List[str]]] = None,
|
325
|
+
num_images_per_prompt: Optional[int] = 1,
|
326
|
+
eta: float = 0.0,
|
327
|
+
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
328
|
+
latents: Optional[torch.FloatTensor] = None,
|
329
|
+
prompt_embeds: Optional[torch.FloatTensor] = None,
|
330
|
+
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
331
|
+
ip_adapter_image: Optional[PipelineImageInput] = None,
|
332
|
+
ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
|
333
|
+
output_type: Optional[str] = "pil",
|
334
|
+
return_dict: bool = True,
|
335
|
+
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
336
|
+
controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
|
337
|
+
guess_mode: bool = False,
|
338
|
+
control_guidance_start: Union[float, List[float]] = 0.0,
|
339
|
+
control_guidance_end: Union[float, List[float]] = 1.0,
|
340
|
+
clip_skip: Optional[int] = None,
|
341
|
+
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
|
342
|
+
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
343
|
+
**kwargs,
|
344
|
+
):
|
345
|
+
r"""
|
346
|
+
The call function to the pipeline for generation.
|
347
|
+
|
348
|
+
Args:
|
349
|
+
prompt (`str` or `List[str]`, *optional*):
|
350
|
+
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
|
351
|
+
image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
|
352
|
+
`List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
|
353
|
+
The initial image to be used as the starting point for the image generation process. Can also accept
|
354
|
+
image latents as `image`, and if passing latents directly they are not encoded again.
|
355
|
+
control_image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
|
356
|
+
`List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
|
357
|
+
The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
|
358
|
+
specified as `torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be
|
359
|
+
accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height
|
360
|
+
and/or width are passed, `image` is resized accordingly. If multiple ControlNets are specified in
|
361
|
+
`init`, images must be passed as a list such that each element of the list can be correctly batched for
|
362
|
+
input to a single ControlNet.
|
363
|
+
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
364
|
+
The height in pixels of the generated image.
|
365
|
+
width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
366
|
+
The width in pixels of the generated image.
|
367
|
+
strength (`float`, *optional*, defaults to 0.8):
|
368
|
+
Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
|
369
|
+
starting point and more noise is added the higher the `strength`. The number of denoising steps depends
|
370
|
+
on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
|
371
|
+
process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
|
372
|
+
essentially ignores `image`.
|
373
|
+
num_inference_steps (`int`, *optional*, defaults to 50):
|
374
|
+
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
375
|
+
expense of slower inference.
|
376
|
+
guidance_scale (`float`, *optional*, defaults to 7.5):
|
377
|
+
A higher guidance scale value encourages the model to generate images closely linked to the text
|
378
|
+
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
|
379
|
+
negative_prompt (`str` or `List[str]`, *optional*):
|
380
|
+
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
|
381
|
+
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
|
382
|
+
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
383
|
+
The number of images to generate per prompt.
|
384
|
+
eta (`float`, *optional*, defaults to 0.0):
|
385
|
+
Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
|
386
|
+
to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
|
387
|
+
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
388
|
+
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
|
389
|
+
generation deterministic.
|
390
|
+
latents (`torch.FloatTensor`, *optional*):
|
391
|
+
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
|
392
|
+
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
393
|
+
tensor is generated by sampling using the supplied random `generator`.
|
394
|
+
prompt_embeds (`torch.FloatTensor`, *optional*):
|
395
|
+
Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
|
396
|
+
provided, text embeddings are generated from the `prompt` input argument.
|
397
|
+
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
398
|
+
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
|
399
|
+
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
|
400
|
+
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
|
401
|
+
ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
|
402
|
+
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters.
|
403
|
+
Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding
|
404
|
+
if `do_classifier_free_guidance` is set to `True`.
|
405
|
+
If not provided, embeddings are computed from the `ip_adapter_image` input argument.
|
406
|
+
output_type (`str`, *optional*, defaults to `"pil"`):
|
407
|
+
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
|
408
|
+
return_dict (`bool`, *optional*, defaults to `True`):
|
409
|
+
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
410
|
+
plain tuple.
|
411
|
+
cross_attention_kwargs (`dict`, *optional*):
|
412
|
+
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
|
413
|
+
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
414
|
+
controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
|
415
|
+
The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
|
416
|
+
to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
|
417
|
+
the corresponding scale as a list.
|
418
|
+
guess_mode (`bool`, *optional*, defaults to `False`):
|
419
|
+
The ControlNet encoder tries to recognize the content of the input image even if you remove all
|
420
|
+
prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
|
421
|
+
control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
|
422
|
+
The percentage of total steps at which the ControlNet starts applying.
|
423
|
+
control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
|
424
|
+
The percentage of total steps at which the ControlNet stops applying.
|
425
|
+
clip_skip (`int`, *optional*):
|
426
|
+
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
|
427
|
+
the output of the pre-final layer will be used for computing the prompt embeddings.
|
428
|
+
callback_on_step_end (`Callable`, *optional*):
|
429
|
+
A function that calls at the end of each denoising steps during the inference. The function is called
|
430
|
+
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
|
431
|
+
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
|
432
|
+
`callback_on_step_end_tensor_inputs`.
|
433
|
+
callback_on_step_end_tensor_inputs (`List`, *optional*):
|
434
|
+
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
|
435
|
+
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
|
436
|
+
`._callback_tensor_inputs` attribute of your pipeine class.
|
437
|
+
|
438
|
+
Examples:
|
439
|
+
|
440
|
+
Returns:
|
441
|
+
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
|
442
|
+
If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
|
443
|
+
otherwise a `tuple` is returned where the first element is a list with the generated images and the
|
444
|
+
second element is a list of `bool`s indicating whether the corresponding generated image contains
|
445
|
+
"not-safe-for-work" (nsfw) content.
|
446
|
+
"""
|
447
|
+
|
448
|
+
callback = kwargs.pop("callback", None)
|
449
|
+
callback_steps = kwargs.pop("callback_steps", None)
|
450
|
+
|
451
|
+
if callback is not None:
|
452
|
+
deprecate(
|
453
|
+
"callback",
|
454
|
+
"1.0.0",
|
455
|
+
"Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
|
456
|
+
)
|
457
|
+
if callback_steps is not None:
|
458
|
+
deprecate(
|
459
|
+
"callback_steps",
|
460
|
+
"1.0.0",
|
461
|
+
"Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
|
462
|
+
)
|
463
|
+
|
464
|
+
controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet
|
465
|
+
|
466
|
+
# align format for control guidance
|
467
|
+
if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
|
468
|
+
control_guidance_start = len(control_guidance_end) * [control_guidance_start]
|
469
|
+
elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):
|
470
|
+
control_guidance_end = len(control_guidance_start) * [control_guidance_end]
|
471
|
+
elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):
|
472
|
+
mult = len(controlnet.nets) if isinstance(controlnet, RBLNMultiControlNetModel) else 1
|
473
|
+
control_guidance_start, control_guidance_end = (
|
474
|
+
mult * [control_guidance_start],
|
475
|
+
mult * [control_guidance_end],
|
476
|
+
)
|
477
|
+
|
478
|
+
# 1. Check inputs. Raise error if not correct
|
479
|
+
self.check_inputs(
|
480
|
+
prompt,
|
481
|
+
control_image,
|
482
|
+
callback_steps,
|
483
|
+
negative_prompt,
|
484
|
+
prompt_embeds,
|
485
|
+
negative_prompt_embeds,
|
486
|
+
ip_adapter_image,
|
487
|
+
ip_adapter_image_embeds,
|
488
|
+
controlnet_conditioning_scale,
|
489
|
+
control_guidance_start,
|
490
|
+
control_guidance_end,
|
491
|
+
callback_on_step_end_tensor_inputs,
|
492
|
+
)
|
493
|
+
|
494
|
+
self._guidance_scale = guidance_scale
|
495
|
+
self._clip_skip = clip_skip
|
496
|
+
self._cross_attention_kwargs = cross_attention_kwargs
|
497
|
+
|
498
|
+
# 2. Define call parameters
|
499
|
+
if prompt is not None and isinstance(prompt, str):
|
500
|
+
batch_size = 1
|
501
|
+
elif prompt is not None and isinstance(prompt, list):
|
502
|
+
batch_size = len(prompt)
|
503
|
+
else:
|
504
|
+
batch_size = prompt_embeds.shape[0]
|
505
|
+
|
506
|
+
device = self._execution_device
|
507
|
+
|
508
|
+
if isinstance(controlnet, RBLNMultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):
|
509
|
+
controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)
|
510
|
+
|
511
|
+
global_pool_conditions = (
|
512
|
+
controlnet.config.global_pool_conditions
|
513
|
+
if isinstance(controlnet, RBLNControlNetModel)
|
514
|
+
else controlnet.nets[0].config.global_pool_conditions
|
515
|
+
)
|
516
|
+
guess_mode = guess_mode or global_pool_conditions
|
517
|
+
|
518
|
+
# 3. Encode input prompt
|
519
|
+
text_encoder_lora_scale = (
|
520
|
+
self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
|
521
|
+
)
|
522
|
+
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
523
|
+
prompt,
|
524
|
+
device,
|
525
|
+
num_images_per_prompt,
|
526
|
+
self.do_classifier_free_guidance,
|
527
|
+
negative_prompt,
|
528
|
+
prompt_embeds=prompt_embeds,
|
529
|
+
negative_prompt_embeds=negative_prompt_embeds,
|
530
|
+
lora_scale=text_encoder_lora_scale,
|
531
|
+
clip_skip=self.clip_skip,
|
532
|
+
)
|
533
|
+
# For classifier free guidance, we need to do two forward passes.
|
534
|
+
# Here we concatenate the unconditional and text embeddings into a single batch
|
535
|
+
# to avoid doing two forward passes
|
536
|
+
if self.do_classifier_free_guidance:
|
537
|
+
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
|
538
|
+
|
539
|
+
if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
|
540
|
+
image_embeds = self.prepare_ip_adapter_image_embeds(
|
541
|
+
ip_adapter_image,
|
542
|
+
ip_adapter_image_embeds,
|
543
|
+
device,
|
544
|
+
batch_size * num_images_per_prompt,
|
545
|
+
self.do_classifier_free_guidance,
|
546
|
+
)
|
547
|
+
|
548
|
+
# 4. Prepare image
|
549
|
+
image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32)
|
550
|
+
|
551
|
+
# 5. Prepare controlnet_conditioning_image
|
552
|
+
if isinstance(controlnet, RBLNControlNetModel):
|
553
|
+
control_image = self.prepare_control_image(
|
554
|
+
image=control_image,
|
555
|
+
width=width,
|
556
|
+
height=height,
|
557
|
+
batch_size=batch_size * num_images_per_prompt,
|
558
|
+
num_images_per_prompt=num_images_per_prompt,
|
559
|
+
device=device,
|
560
|
+
dtype=controlnet.dtype,
|
561
|
+
do_classifier_free_guidance=self.do_classifier_free_guidance,
|
562
|
+
guess_mode=guess_mode,
|
563
|
+
)
|
564
|
+
elif isinstance(controlnet, RBLNMultiControlNetModel):
|
565
|
+
control_images = []
|
566
|
+
|
567
|
+
for control_image_ in control_image:
|
568
|
+
control_image_ = self.prepare_control_image(
|
569
|
+
image=control_image_,
|
570
|
+
width=width,
|
571
|
+
height=height,
|
572
|
+
batch_size=batch_size * num_images_per_prompt,
|
573
|
+
num_images_per_prompt=num_images_per_prompt,
|
574
|
+
device=device,
|
575
|
+
dtype=controlnet.dtype,
|
576
|
+
do_classifier_free_guidance=self.do_classifier_free_guidance,
|
577
|
+
guess_mode=guess_mode,
|
578
|
+
)
|
579
|
+
|
580
|
+
control_images.append(control_image_)
|
581
|
+
|
582
|
+
control_image = control_images
|
583
|
+
else:
|
584
|
+
assert False
|
585
|
+
|
586
|
+
# 5. Prepare timesteps
|
587
|
+
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
588
|
+
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
|
589
|
+
latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
|
590
|
+
self._num_timesteps = len(timesteps)
|
591
|
+
|
592
|
+
# 6. Prepare latent variables
|
593
|
+
latents = self.prepare_latents(
|
594
|
+
image,
|
595
|
+
latent_timestep,
|
596
|
+
batch_size,
|
597
|
+
num_images_per_prompt,
|
598
|
+
prompt_embeds.dtype,
|
599
|
+
device,
|
600
|
+
generator,
|
601
|
+
)
|
602
|
+
|
603
|
+
# 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
|
604
|
+
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
605
|
+
|
606
|
+
# 7.1 Add image embeds for IP-Adapter
|
607
|
+
added_cond_kwargs = (
|
608
|
+
{"image_embeds": image_embeds}
|
609
|
+
if ip_adapter_image is not None or ip_adapter_image_embeds is not None
|
610
|
+
else None
|
611
|
+
)
|
612
|
+
|
613
|
+
# 7.2 Create tensor stating which controlnets to keep
|
614
|
+
controlnet_keep = []
|
615
|
+
for i in range(len(timesteps)):
|
616
|
+
keeps = [
|
617
|
+
1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
|
618
|
+
for s, e in zip(control_guidance_start, control_guidance_end)
|
619
|
+
]
|
620
|
+
controlnet_keep.append(keeps[0] if isinstance(controlnet, RBLNControlNetModel) else keeps)
|
621
|
+
|
622
|
+
# 8. Denoising loop
|
623
|
+
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
624
|
+
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
625
|
+
for i, t in enumerate(timesteps):
|
626
|
+
# expand the latents if we are doing classifier free guidance
|
627
|
+
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
|
628
|
+
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
629
|
+
|
630
|
+
# controlnet(s) inference
|
631
|
+
if guess_mode and self.do_classifier_free_guidance:
|
632
|
+
# Infer ControlNet only for the conditional batch.
|
633
|
+
control_model_input = latents
|
634
|
+
control_model_input = self.scheduler.scale_model_input(control_model_input, t)
|
635
|
+
controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]
|
636
|
+
else:
|
637
|
+
control_model_input = latent_model_input
|
638
|
+
controlnet_prompt_embeds = prompt_embeds
|
639
|
+
|
640
|
+
if isinstance(controlnet_keep[i], list):
|
641
|
+
cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]
|
642
|
+
else:
|
643
|
+
controlnet_cond_scale = controlnet_conditioning_scale
|
644
|
+
if isinstance(controlnet_cond_scale, list):
|
645
|
+
controlnet_cond_scale = controlnet_cond_scale[0]
|
646
|
+
cond_scale = controlnet_cond_scale * controlnet_keep[i]
|
647
|
+
|
648
|
+
down_block_res_samples, mid_block_res_sample = self.controlnet(
|
649
|
+
control_model_input,
|
650
|
+
t,
|
651
|
+
encoder_hidden_states=controlnet_prompt_embeds,
|
652
|
+
controlnet_cond=control_image,
|
653
|
+
conditioning_scale=cond_scale,
|
654
|
+
guess_mode=guess_mode,
|
655
|
+
return_dict=False,
|
656
|
+
)
|
657
|
+
|
658
|
+
if guess_mode and self.do_classifier_free_guidance:
|
659
|
+
# Infered ControlNet only for the conditional batch.
|
660
|
+
# To apply the output of ControlNet to both the unconditional and conditional batches,
|
661
|
+
# add 0 to the unconditional batch to keep it unchanged.
|
662
|
+
down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]
|
663
|
+
mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])
|
664
|
+
|
665
|
+
# predict the noise residual
|
666
|
+
noise_pred = self.unet(
|
667
|
+
latent_model_input,
|
668
|
+
t,
|
669
|
+
encoder_hidden_states=prompt_embeds,
|
670
|
+
cross_attention_kwargs=self.cross_attention_kwargs,
|
671
|
+
down_block_additional_residuals=down_block_res_samples,
|
672
|
+
mid_block_additional_residual=mid_block_res_sample,
|
673
|
+
added_cond_kwargs=added_cond_kwargs,
|
674
|
+
return_dict=False,
|
675
|
+
)[0]
|
676
|
+
|
677
|
+
# perform guidance
|
678
|
+
if self.do_classifier_free_guidance:
|
679
|
+
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
680
|
+
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
681
|
+
|
682
|
+
# compute the previous noisy sample x_t -> x_t-1
|
683
|
+
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
|
684
|
+
|
685
|
+
if callback_on_step_end is not None:
|
686
|
+
callback_kwargs = {}
|
687
|
+
for k in callback_on_step_end_tensor_inputs:
|
688
|
+
callback_kwargs[k] = locals()[k]
|
689
|
+
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
690
|
+
|
691
|
+
latents = callback_outputs.pop("latents", latents)
|
692
|
+
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
693
|
+
negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
|
694
|
+
|
695
|
+
# call the callback, if provided
|
696
|
+
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
697
|
+
progress_bar.update()
|
698
|
+
if callback is not None and i % callback_steps == 0:
|
699
|
+
step_idx = i // getattr(self.scheduler, "order", 1)
|
700
|
+
callback(step_idx, t, latents)
|
701
|
+
|
702
|
+
# If we do sequential model offloading, let's offload unet and controlnet
|
703
|
+
# manually for max memory savings
|
704
|
+
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
|
705
|
+
self.unet.to("cpu")
|
706
|
+
self.controlnet.to("cpu")
|
707
|
+
torch.cuda.empty_cache()
|
708
|
+
|
709
|
+
if not output_type == "latent":
|
710
|
+
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
|
711
|
+
0
|
712
|
+
]
|
713
|
+
image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
|
714
|
+
else:
|
715
|
+
image = latents
|
716
|
+
has_nsfw_concept = None
|
717
|
+
|
718
|
+
if has_nsfw_concept is None:
|
719
|
+
do_denormalize = [True] * image.shape[0]
|
720
|
+
else:
|
721
|
+
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
|
722
|
+
|
723
|
+
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
|
724
|
+
|
725
|
+
# Offload all models
|
726
|
+
self.maybe_free_model_hooks()
|
727
|
+
|
728
|
+
if not return_dict:
|
729
|
+
return (image, has_nsfw_concept)
|
730
|
+
|
731
|
+
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
|