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.
Files changed (51) hide show
  1. optimum/rbln/__init__.py +115 -0
  2. optimum/rbln/__version__.py +1 -0
  3. optimum/rbln/diffusers/__init__.py +64 -0
  4. optimum/rbln/diffusers/models/__init__.py +26 -0
  5. optimum/rbln/diffusers/models/autoencoder_kl.py +313 -0
  6. optimum/rbln/diffusers/models/controlnet.py +180 -0
  7. optimum/rbln/diffusers/models/unet_2d_condition.py +352 -0
  8. optimum/rbln/diffusers/pipelines/__init__.py +30 -0
  9. optimum/rbln/diffusers/pipelines/controlnet/__init__.py +24 -0
  10. optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +266 -0
  11. optimum/rbln/diffusers/pipelines/stable_diffusion/__init__.py +26 -0
  12. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_controlnet_img2img.py +731 -0
  13. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +106 -0
  14. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +116 -0
  15. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/__init__.py +2 -0
  16. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +109 -0
  17. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +111 -0
  18. optimum/rbln/modeling.py +0 -0
  19. optimum/rbln/modeling_alias.py +49 -0
  20. optimum/rbln/modeling_base.py +645 -0
  21. optimum/rbln/modeling_config.py +169 -0
  22. optimum/rbln/modeling_seq2seq.py +469 -0
  23. optimum/rbln/transformers/__init__.py +59 -0
  24. optimum/rbln/transformers/generation/__init__.py +24 -0
  25. optimum/rbln/transformers/generation/streamers.py +122 -0
  26. optimum/rbln/transformers/models/__init__.py +28 -0
  27. optimum/rbln/transformers/models/bart/__init__.py +24 -0
  28. optimum/rbln/transformers/models/bart/bart_architecture.py +377 -0
  29. optimum/rbln/transformers/models/clip/__init__.py +24 -0
  30. optimum/rbln/transformers/models/clip/modeling_clip.py +116 -0
  31. optimum/rbln/transformers/models/gpt2/__init__.py +24 -0
  32. optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +253 -0
  33. optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +700 -0
  34. optimum/rbln/transformers/models/llama/__init__.py +24 -0
  35. optimum/rbln/transformers/models/llama/llama_architecture.py +607 -0
  36. optimum/rbln/transformers/models/llama/modeling_llama.py +409 -0
  37. optimum/rbln/transformers/models/t5/__init__.py +24 -0
  38. optimum/rbln/transformers/models/t5/t5_architecture.py +439 -0
  39. optimum/rbln/transformers/models/wav2vec2/__init__.py +24 -0
  40. optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py +121 -0
  41. optimum/rbln/transformers/models/whisper/__init__.py +24 -0
  42. optimum/rbln/transformers/models/whisper/modeling_whisper.py +374 -0
  43. optimum/rbln/transformers/models/whisper/whisper_architecture.py +406 -0
  44. optimum/rbln/utils/__init__.py +25 -0
  45. optimum/rbln/utils/import_utils.py +28 -0
  46. optimum/rbln/utils/runtime_utils.py +71 -0
  47. optimum/rbln/utils/save_utils.py +92 -0
  48. optimum_rbln-0.1.0.dist-info/METADATA +144 -0
  49. optimum_rbln-0.1.0.dist-info/RECORD +51 -0
  50. optimum_rbln-0.1.0.dist-info/WHEEL +4 -0
  51. 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)