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