hyoungwoncho commited on
Commit
c51dd04
·
verified ·
1 Parent(s): d4ca2e1

Upload pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +1783 -0
pipeline.py ADDED
@@ -0,0 +1,1783 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Implementation of Stable Diffusion Inpaint Pipeline with Perturbed-Attention Guidance
2
+
3
+ import inspect
4
+ from typing import Any, Callable, Dict, List, Optional, Union
5
+
6
+ import numpy as np
7
+ import PIL.Image
8
+ import torch
9
+ import torch.nn.functional as F
10
+ from packaging import version
11
+ from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
12
+
13
+ from diffusers.configuration_utils import FrozenDict
14
+ from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
15
+ from diffusers.loaders import FromSingleFileMixin, IPAdapterMixin, LoraLoaderMixin, TextualInversionLoaderMixin
16
+ from diffusers.models import AsymmetricAutoencoderKL, AutoencoderKL, ImageProjection, UNet2DConditionModel
17
+ from diffusers.models.lora import adjust_lora_scale_text_encoder
18
+ from diffusers.schedulers import KarrasDiffusionSchedulers
19
+ from diffusers.utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
20
+ from diffusers.utils.torch_utils import randn_tensor
21
+ from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
22
+ from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
23
+ from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
24
+ from diffusers.models.attention_processor import Attention, AttnProcessor2_0
25
+
26
+
27
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
28
+
29
+ class PAGIdentitySelfAttnProcessor:
30
+ r"""
31
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
32
+ """
33
+
34
+ def __init__(self):
35
+ if not hasattr(F, "scaled_dot_product_attention"):
36
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
37
+
38
+ def __call__(
39
+ self,
40
+ attn: Attention,
41
+ hidden_states: torch.FloatTensor,
42
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
43
+ attention_mask: Optional[torch.FloatTensor] = None,
44
+ temb: Optional[torch.FloatTensor] = None,
45
+ *args,
46
+ **kwargs,
47
+ ) -> torch.FloatTensor:
48
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
49
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
50
+ deprecate("scale", "1.0.0", deprecation_message)
51
+
52
+ residual = hidden_states
53
+ if attn.spatial_norm is not None:
54
+ hidden_states = attn.spatial_norm(hidden_states, temb)
55
+
56
+ input_ndim = hidden_states.ndim
57
+ if input_ndim == 4:
58
+ batch_size, channel, height, width = hidden_states.shape
59
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
60
+
61
+ # chunk
62
+ hidden_states_org, hidden_states_ptb = hidden_states.chunk(2)
63
+
64
+ # original path
65
+ batch_size, sequence_length, _ = hidden_states_org.shape
66
+
67
+ if attention_mask is not None:
68
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
69
+ # scaled_dot_product_attention expects attention_mask shape to be
70
+ # (batch, heads, source_length, target_length)
71
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
72
+
73
+ if attn.group_norm is not None:
74
+ hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)
75
+
76
+ query = attn.to_q(hidden_states_org)
77
+ key = attn.to_k(hidden_states_org)
78
+ value = attn.to_v(hidden_states_org)
79
+
80
+ inner_dim = key.shape[-1]
81
+ head_dim = inner_dim // attn.heads
82
+
83
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
84
+
85
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
86
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
87
+
88
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
89
+ # TODO: add support for attn.scale when we move to Torch 2.1
90
+ hidden_states_org = F.scaled_dot_product_attention(
91
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
92
+ )
93
+
94
+ hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
95
+ hidden_states_org = hidden_states_org.to(query.dtype)
96
+
97
+ # linear proj
98
+ hidden_states_org = attn.to_out[0](hidden_states_org)
99
+ # dropout
100
+ hidden_states_org = attn.to_out[1](hidden_states_org)
101
+
102
+ if input_ndim == 4:
103
+ hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)
104
+
105
+ # perturbed path (identity attention)
106
+ batch_size, sequence_length, _ = hidden_states_ptb.shape
107
+
108
+ if attention_mask is not None:
109
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
110
+ # scaled_dot_product_attention expects attention_mask shape to be
111
+ # (batch, heads, source_length, target_length)
112
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
113
+
114
+ if attn.group_norm is not None:
115
+ hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)
116
+
117
+ value = attn.to_v(hidden_states_ptb)
118
+
119
+ hidden_states_ptb = torch.zeros(value.shape).to(value.get_device())
120
+ # hidden_states_ptb = value
121
+
122
+ hidden_states_ptb = hidden_states_ptb.to(query.dtype)
123
+
124
+ # linear proj
125
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
126
+ # dropout
127
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
128
+
129
+ if input_ndim == 4:
130
+ hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)
131
+
132
+ # cat
133
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
134
+
135
+ if attn.residual_connection:
136
+ hidden_states = hidden_states + residual
137
+
138
+ hidden_states = hidden_states / attn.rescale_output_factor
139
+
140
+ return hidden_states
141
+
142
+
143
+ class PAGCFGIdentitySelfAttnProcessor:
144
+ r"""
145
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
146
+ """
147
+
148
+ def __init__(self):
149
+ if not hasattr(F, "scaled_dot_product_attention"):
150
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
151
+
152
+ def __call__(
153
+ self,
154
+ attn: Attention,
155
+ hidden_states: torch.FloatTensor,
156
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
157
+ attention_mask: Optional[torch.FloatTensor] = None,
158
+ temb: Optional[torch.FloatTensor] = None,
159
+ *args,
160
+ **kwargs,
161
+ ) -> torch.FloatTensor:
162
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
163
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
164
+ deprecate("scale", "1.0.0", deprecation_message)
165
+
166
+ residual = hidden_states
167
+ if attn.spatial_norm is not None:
168
+ hidden_states = attn.spatial_norm(hidden_states, temb)
169
+
170
+ input_ndim = hidden_states.ndim
171
+ if input_ndim == 4:
172
+ batch_size, channel, height, width = hidden_states.shape
173
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
174
+
175
+ # chunk
176
+ hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3)
177
+ hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org])
178
+
179
+ # original path
180
+ batch_size, sequence_length, _ = hidden_states_org.shape
181
+
182
+ if attention_mask is not None:
183
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
184
+ # scaled_dot_product_attention expects attention_mask shape to be
185
+ # (batch, heads, source_length, target_length)
186
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
187
+
188
+ if attn.group_norm is not None:
189
+ hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)
190
+
191
+ query = attn.to_q(hidden_states_org)
192
+ key = attn.to_k(hidden_states_org)
193
+ value = attn.to_v(hidden_states_org)
194
+
195
+ inner_dim = key.shape[-1]
196
+ head_dim = inner_dim // attn.heads
197
+
198
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
199
+
200
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
201
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
202
+
203
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
204
+ # TODO: add support for attn.scale when we move to Torch 2.1
205
+ hidden_states_org = F.scaled_dot_product_attention(
206
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
207
+ )
208
+
209
+ hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
210
+ hidden_states_org = hidden_states_org.to(query.dtype)
211
+
212
+ # linear proj
213
+ hidden_states_org = attn.to_out[0](hidden_states_org)
214
+ # dropout
215
+ hidden_states_org = attn.to_out[1](hidden_states_org)
216
+
217
+ if input_ndim == 4:
218
+ hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)
219
+
220
+ # perturbed path (identity attention)
221
+ batch_size, sequence_length, _ = hidden_states_ptb.shape
222
+
223
+ if attention_mask is not None:
224
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
225
+ # scaled_dot_product_attention expects attention_mask shape to be
226
+ # (batch, heads, source_length, target_length)
227
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
228
+
229
+ if attn.group_norm is not None:
230
+ hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)
231
+
232
+ value = attn.to_v(hidden_states_ptb)
233
+ hidden_states_ptb = value
234
+ hidden_states_ptb = hidden_states_ptb.to(query.dtype)
235
+
236
+ # linear proj
237
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
238
+ # dropout
239
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
240
+
241
+ if input_ndim == 4:
242
+ hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)
243
+
244
+ # cat
245
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
246
+
247
+ if attn.residual_connection:
248
+ hidden_states = hidden_states + residual
249
+
250
+ hidden_states = hidden_states / attn.rescale_output_factor
251
+
252
+ return hidden_states
253
+
254
+
255
+ def prepare_mask_and_masked_image(image, mask, height, width, return_image: bool = False):
256
+ """
257
+ Prepares a pair (image, mask) to be consumed by the Stable Diffusion pipeline. This means that those inputs will be
258
+ converted to ``torch.Tensor`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for the
259
+ ``image`` and ``1`` for the ``mask``.
260
+
261
+ The ``image`` will be converted to ``torch.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be
262
+ binarized (``mask > 0.5``) and cast to ``torch.float32`` too.
263
+
264
+ Args:
265
+ image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint.
266
+ It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width``
267
+ ``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``.
268
+ mask (_type_): The mask to apply to the image, i.e. regions to inpaint.
269
+ It can be a ``PIL.Image``, or a ``height x width`` ``np.array`` or a ``1 x height x width``
270
+ ``torch.Tensor`` or a ``batch x 1 x height x width`` ``torch.Tensor``.
271
+
272
+
273
+ Raises:
274
+ ValueError: ``torch.Tensor`` images should be in the ``[-1, 1]`` range. ValueError: ``torch.Tensor`` mask
275
+ should be in the ``[0, 1]`` range. ValueError: ``mask`` and ``image`` should have the same spatial dimensions.
276
+ TypeError: ``mask`` is a ``torch.Tensor`` but ``image`` is not
277
+ (ot the other way around).
278
+
279
+ Returns:
280
+ tuple[torch.Tensor]: The pair (mask, masked_image) as ``torch.Tensor`` with 4
281
+ dimensions: ``batch x channels x height x width``.
282
+ """
283
+ deprecation_message = "The prepare_mask_and_masked_image method is deprecated and will be removed in a future version. Please use VaeImageProcessor.preprocess instead"
284
+ deprecate(
285
+ "prepare_mask_and_masked_image",
286
+ "0.30.0",
287
+ deprecation_message,
288
+ )
289
+ if image is None:
290
+ raise ValueError("`image` input cannot be undefined.")
291
+
292
+ if mask is None:
293
+ raise ValueError("`mask_image` input cannot be undefined.")
294
+
295
+ if isinstance(image, torch.Tensor):
296
+ if not isinstance(mask, torch.Tensor):
297
+ raise TypeError(f"`image` is a torch.Tensor but `mask` (type: {type(mask)} is not")
298
+
299
+ # Batch single image
300
+ if image.ndim == 3:
301
+ assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)"
302
+ image = image.unsqueeze(0)
303
+
304
+ # Batch and add channel dim for single mask
305
+ if mask.ndim == 2:
306
+ mask = mask.unsqueeze(0).unsqueeze(0)
307
+
308
+ # Batch single mask or add channel dim
309
+ if mask.ndim == 3:
310
+ # Single batched mask, no channel dim or single mask not batched but channel dim
311
+ if mask.shape[0] == 1:
312
+ mask = mask.unsqueeze(0)
313
+
314
+ # Batched masks no channel dim
315
+ else:
316
+ mask = mask.unsqueeze(1)
317
+
318
+ assert image.ndim == 4 and mask.ndim == 4, "Image and Mask must have 4 dimensions"
319
+ assert image.shape[-2:] == mask.shape[-2:], "Image and Mask must have the same spatial dimensions"
320
+ assert image.shape[0] == mask.shape[0], "Image and Mask must have the same batch size"
321
+
322
+ # Check image is in [-1, 1]
323
+ if image.min() < -1 or image.max() > 1:
324
+ raise ValueError("Image should be in [-1, 1] range")
325
+
326
+ # Check mask is in [0, 1]
327
+ if mask.min() < 0 or mask.max() > 1:
328
+ raise ValueError("Mask should be in [0, 1] range")
329
+
330
+ # Binarize mask
331
+ mask[mask < 0.5] = 0
332
+ mask[mask >= 0.5] = 1
333
+
334
+ # Image as float32
335
+ image = image.to(dtype=torch.float32)
336
+ elif isinstance(mask, torch.Tensor):
337
+ raise TypeError(f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not")
338
+ else:
339
+ # preprocess image
340
+ if isinstance(image, (PIL.Image.Image, np.ndarray)):
341
+ image = [image]
342
+ if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):
343
+ # resize all images w.r.t passed height an width
344
+ image = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in image]
345
+ image = [np.array(i.convert("RGB"))[None, :] for i in image]
346
+ image = np.concatenate(image, axis=0)
347
+ elif isinstance(image, list) and isinstance(image[0], np.ndarray):
348
+ image = np.concatenate([i[None, :] for i in image], axis=0)
349
+
350
+ image = image.transpose(0, 3, 1, 2)
351
+ image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0
352
+
353
+ # preprocess mask
354
+ if isinstance(mask, (PIL.Image.Image, np.ndarray)):
355
+ mask = [mask]
356
+
357
+ if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image):
358
+ mask = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in mask]
359
+ mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0)
360
+ mask = mask.astype(np.float32) / 255.0
361
+ elif isinstance(mask, list) and isinstance(mask[0], np.ndarray):
362
+ mask = np.concatenate([m[None, None, :] for m in mask], axis=0)
363
+
364
+ mask[mask < 0.5] = 0
365
+ mask[mask >= 0.5] = 1
366
+ mask = torch.from_numpy(mask)
367
+
368
+ masked_image = image * (mask < 0.5)
369
+
370
+ # n.b. ensure backwards compatibility as old function does not return image
371
+ if return_image:
372
+ return mask, masked_image, image
373
+
374
+ return mask, masked_image
375
+
376
+
377
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
378
+ def retrieve_latents(
379
+ encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
380
+ ):
381
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
382
+ return encoder_output.latent_dist.sample(generator)
383
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
384
+ return encoder_output.latent_dist.mode()
385
+ elif hasattr(encoder_output, "latents"):
386
+ return encoder_output.latents
387
+ else:
388
+ raise AttributeError("Could not access latents of provided encoder_output")
389
+
390
+
391
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
392
+ def retrieve_timesteps(
393
+ scheduler,
394
+ num_inference_steps: Optional[int] = None,
395
+ device: Optional[Union[str, torch.device]] = None,
396
+ timesteps: Optional[List[int]] = None,
397
+ **kwargs,
398
+ ):
399
+ """
400
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
401
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
402
+
403
+ Args:
404
+ scheduler (`SchedulerMixin`):
405
+ The scheduler to get timesteps from.
406
+ num_inference_steps (`int`):
407
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
408
+ must be `None`.
409
+ device (`str` or `torch.device`, *optional*):
410
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
411
+ timesteps (`List[int]`, *optional*):
412
+ Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
413
+ timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
414
+ must be `None`.
415
+
416
+ Returns:
417
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
418
+ second element is the number of inference steps.
419
+ """
420
+ if timesteps is not None:
421
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
422
+ if not accepts_timesteps:
423
+ raise ValueError(
424
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
425
+ f" timestep schedules. Please check whether you are using the correct scheduler."
426
+ )
427
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
428
+ timesteps = scheduler.timesteps
429
+ num_inference_steps = len(timesteps)
430
+ else:
431
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
432
+ timesteps = scheduler.timesteps
433
+ return timesteps, num_inference_steps
434
+
435
+
436
+ class StableDiffusionInpaintPipeline(
437
+ DiffusionPipeline,
438
+ StableDiffusionMixin,
439
+ TextualInversionLoaderMixin,
440
+ IPAdapterMixin,
441
+ LoraLoaderMixin,
442
+ FromSingleFileMixin,
443
+ ):
444
+ r"""
445
+ Pipeline for text-guided image inpainting using Stable Diffusion.
446
+
447
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
448
+ implemented for all pipelines (downloading, saving, running on a particular device, etc.).
449
+
450
+ The pipeline also inherits the following loading methods:
451
+ - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
452
+ - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
453
+ - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
454
+ - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
455
+ - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
456
+
457
+ Args:
458
+ vae ([`AutoencoderKL`, `AsymmetricAutoencoderKL`]):
459
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
460
+ text_encoder ([`CLIPTextModel`]):
461
+ Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
462
+ tokenizer ([`~transformers.CLIPTokenizer`]):
463
+ A `CLIPTokenizer` to tokenize text.
464
+ unet ([`UNet2DConditionModel`]):
465
+ A `UNet2DConditionModel` to denoise the encoded image latents.
466
+ scheduler ([`SchedulerMixin`]):
467
+ A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
468
+ [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
469
+ safety_checker ([`StableDiffusionSafetyChecker`]):
470
+ Classification module that estimates whether generated images could be considered offensive or harmful.
471
+ Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
472
+ about a model's potential harms.
473
+ feature_extractor ([`~transformers.CLIPImageProcessor`]):
474
+ A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
475
+ """
476
+
477
+ model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
478
+ _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
479
+ _exclude_from_cpu_offload = ["safety_checker"]
480
+ _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "mask", "masked_image_latents"]
481
+
482
+ def __init__(
483
+ self,
484
+ vae: Union[AutoencoderKL, AsymmetricAutoencoderKL],
485
+ text_encoder: CLIPTextModel,
486
+ tokenizer: CLIPTokenizer,
487
+ unet: UNet2DConditionModel,
488
+ scheduler: KarrasDiffusionSchedulers,
489
+ safety_checker: StableDiffusionSafetyChecker,
490
+ feature_extractor: CLIPImageProcessor,
491
+ image_encoder: CLIPVisionModelWithProjection = None,
492
+ requires_safety_checker: bool = True,
493
+ ):
494
+ super().__init__()
495
+
496
+ if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
497
+ deprecation_message = (
498
+ f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
499
+ f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
500
+ "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
501
+ " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
502
+ " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
503
+ " file"
504
+ )
505
+ deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
506
+ new_config = dict(scheduler.config)
507
+ new_config["steps_offset"] = 1
508
+ scheduler._internal_dict = FrozenDict(new_config)
509
+
510
+ if hasattr(scheduler.config, "skip_prk_steps") and scheduler.config.skip_prk_steps is False:
511
+ deprecation_message = (
512
+ f"The configuration file of this scheduler: {scheduler} has not set the configuration"
513
+ " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"
514
+ " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"
515
+ " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"
516
+ " Hub, it would be very nice if you could open a Pull request for the"
517
+ " `scheduler/scheduler_config.json` file"
518
+ )
519
+ deprecate("skip_prk_steps not set", "1.0.0", deprecation_message, standard_warn=False)
520
+ new_config = dict(scheduler.config)
521
+ new_config["skip_prk_steps"] = True
522
+ scheduler._internal_dict = FrozenDict(new_config)
523
+
524
+ if safety_checker is None and requires_safety_checker:
525
+ logger.warning(
526
+ f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
527
+ " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
528
+ " results in services or applications open to the public. Both the diffusers team and Hugging Face"
529
+ " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
530
+ " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
531
+ " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
532
+ )
533
+
534
+ if safety_checker is not None and feature_extractor is None:
535
+ raise ValueError(
536
+ "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
537
+ " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
538
+ )
539
+
540
+ is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
541
+ version.parse(unet.config._diffusers_version).base_version
542
+ ) < version.parse("0.9.0.dev0")
543
+ is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
544
+ if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
545
+ deprecation_message = (
546
+ "The configuration file of the unet has set the default `sample_size` to smaller than"
547
+ " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"
548
+ " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
549
+ " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
550
+ " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
551
+ " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
552
+ " in the config might lead to incorrect results in future versions. If you have downloaded this"
553
+ " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
554
+ " the `unet/config.json` file"
555
+ )
556
+ deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
557
+ new_config = dict(unet.config)
558
+ new_config["sample_size"] = 64
559
+ unet._internal_dict = FrozenDict(new_config)
560
+
561
+ # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4
562
+ if unet.config.in_channels != 9:
563
+ logger.info(f"You have loaded a UNet with {unet.config.in_channels} input channels which.")
564
+
565
+ self.register_modules(
566
+ vae=vae,
567
+ text_encoder=text_encoder,
568
+ tokenizer=tokenizer,
569
+ unet=unet,
570
+ scheduler=scheduler,
571
+ safety_checker=safety_checker,
572
+ feature_extractor=feature_extractor,
573
+ image_encoder=image_encoder,
574
+ )
575
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
576
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
577
+ self.mask_processor = VaeImageProcessor(
578
+ vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True
579
+ )
580
+ self.register_to_config(requires_safety_checker=requires_safety_checker)
581
+
582
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
583
+ def _encode_prompt(
584
+ self,
585
+ prompt,
586
+ device,
587
+ num_images_per_prompt,
588
+ do_classifier_free_guidance,
589
+ negative_prompt=None,
590
+ prompt_embeds: Optional[torch.FloatTensor] = None,
591
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
592
+ lora_scale: Optional[float] = None,
593
+ **kwargs,
594
+ ):
595
+ deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
596
+ deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
597
+
598
+ prompt_embeds_tuple = self.encode_prompt(
599
+ prompt=prompt,
600
+ device=device,
601
+ num_images_per_prompt=num_images_per_prompt,
602
+ do_classifier_free_guidance=do_classifier_free_guidance,
603
+ negative_prompt=negative_prompt,
604
+ prompt_embeds=prompt_embeds,
605
+ negative_prompt_embeds=negative_prompt_embeds,
606
+ lora_scale=lora_scale,
607
+ **kwargs,
608
+ )
609
+
610
+ # concatenate for backwards comp
611
+ prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])
612
+
613
+ return prompt_embeds
614
+
615
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt
616
+ def encode_prompt(
617
+ self,
618
+ prompt,
619
+ device,
620
+ num_images_per_prompt,
621
+ do_classifier_free_guidance,
622
+ negative_prompt=None,
623
+ prompt_embeds: Optional[torch.FloatTensor] = None,
624
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
625
+ lora_scale: Optional[float] = None,
626
+ clip_skip: Optional[int] = None,
627
+ ):
628
+ r"""
629
+ Encodes the prompt into text encoder hidden states.
630
+
631
+ Args:
632
+ prompt (`str` or `List[str]`, *optional*):
633
+ prompt to be encoded
634
+ device: (`torch.device`):
635
+ torch device
636
+ num_images_per_prompt (`int`):
637
+ number of images that should be generated per prompt
638
+ do_classifier_free_guidance (`bool`):
639
+ whether to use classifier free guidance or not
640
+ negative_prompt (`str` or `List[str]`, *optional*):
641
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
642
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
643
+ less than `1`).
644
+ prompt_embeds (`torch.FloatTensor`, *optional*):
645
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
646
+ provided, text embeddings will be generated from `prompt` input argument.
647
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
648
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
649
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
650
+ argument.
651
+ lora_scale (`float`, *optional*):
652
+ A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
653
+ clip_skip (`int`, *optional*):
654
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
655
+ the output of the pre-final layer will be used for computing the prompt embeddings.
656
+ """
657
+ # set lora scale so that monkey patched LoRA
658
+ # function of text encoder can correctly access it
659
+ if lora_scale is not None and isinstance(self, LoraLoaderMixin):
660
+ self._lora_scale = lora_scale
661
+
662
+ # dynamically adjust the LoRA scale
663
+ if not USE_PEFT_BACKEND:
664
+ adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
665
+ else:
666
+ scale_lora_layers(self.text_encoder, lora_scale)
667
+
668
+ if prompt is not None and isinstance(prompt, str):
669
+ batch_size = 1
670
+ elif prompt is not None and isinstance(prompt, list):
671
+ batch_size = len(prompt)
672
+ else:
673
+ batch_size = prompt_embeds.shape[0]
674
+
675
+ if prompt_embeds is None:
676
+ # textual inversion: process multi-vector tokens if necessary
677
+ if isinstance(self, TextualInversionLoaderMixin):
678
+ prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
679
+
680
+ text_inputs = self.tokenizer(
681
+ prompt,
682
+ padding="max_length",
683
+ max_length=self.tokenizer.model_max_length,
684
+ truncation=True,
685
+ return_tensors="pt",
686
+ )
687
+ text_input_ids = text_inputs.input_ids
688
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
689
+
690
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
691
+ text_input_ids, untruncated_ids
692
+ ):
693
+ removed_text = self.tokenizer.batch_decode(
694
+ untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
695
+ )
696
+ logger.warning(
697
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
698
+ f" {self.tokenizer.model_max_length} tokens: {removed_text}"
699
+ )
700
+
701
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
702
+ attention_mask = text_inputs.attention_mask.to(device)
703
+ else:
704
+ attention_mask = None
705
+
706
+ if clip_skip is None:
707
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
708
+ prompt_embeds = prompt_embeds[0]
709
+ else:
710
+ prompt_embeds = self.text_encoder(
711
+ text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
712
+ )
713
+ # Access the `hidden_states` first, that contains a tuple of
714
+ # all the hidden states from the encoder layers. Then index into
715
+ # the tuple to access the hidden states from the desired layer.
716
+ prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]
717
+ # We also need to apply the final LayerNorm here to not mess with the
718
+ # representations. The `last_hidden_states` that we typically use for
719
+ # obtaining the final prompt representations passes through the LayerNorm
720
+ # layer.
721
+ prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)
722
+
723
+ if self.text_encoder is not None:
724
+ prompt_embeds_dtype = self.text_encoder.dtype
725
+ elif self.unet is not None:
726
+ prompt_embeds_dtype = self.unet.dtype
727
+ else:
728
+ prompt_embeds_dtype = prompt_embeds.dtype
729
+
730
+ prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
731
+
732
+ bs_embed, seq_len, _ = prompt_embeds.shape
733
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
734
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
735
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
736
+
737
+ # get unconditional embeddings for classifier free guidance
738
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
739
+ uncond_tokens: List[str]
740
+ if negative_prompt is None:
741
+ uncond_tokens = [""] * batch_size
742
+ elif prompt is not None and type(prompt) is not type(negative_prompt):
743
+ raise TypeError(
744
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
745
+ f" {type(prompt)}."
746
+ )
747
+ elif isinstance(negative_prompt, str):
748
+ uncond_tokens = [negative_prompt]
749
+ elif batch_size != len(negative_prompt):
750
+ raise ValueError(
751
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
752
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
753
+ " the batch size of `prompt`."
754
+ )
755
+ else:
756
+ uncond_tokens = negative_prompt
757
+
758
+ # textual inversion: process multi-vector tokens if necessary
759
+ if isinstance(self, TextualInversionLoaderMixin):
760
+ uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
761
+
762
+ max_length = prompt_embeds.shape[1]
763
+ uncond_input = self.tokenizer(
764
+ uncond_tokens,
765
+ padding="max_length",
766
+ max_length=max_length,
767
+ truncation=True,
768
+ return_tensors="pt",
769
+ )
770
+
771
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
772
+ attention_mask = uncond_input.attention_mask.to(device)
773
+ else:
774
+ attention_mask = None
775
+
776
+ negative_prompt_embeds = self.text_encoder(
777
+ uncond_input.input_ids.to(device),
778
+ attention_mask=attention_mask,
779
+ )
780
+ negative_prompt_embeds = negative_prompt_embeds[0]
781
+
782
+ if do_classifier_free_guidance:
783
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
784
+ seq_len = negative_prompt_embeds.shape[1]
785
+
786
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
787
+
788
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
789
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
790
+
791
+ if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:
792
+ # Retrieve the original scale by scaling back the LoRA layers
793
+ unscale_lora_layers(self.text_encoder, lora_scale)
794
+
795
+ return prompt_embeds, negative_prompt_embeds
796
+
797
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
798
+ def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
799
+ dtype = next(self.image_encoder.parameters()).dtype
800
+
801
+ if not isinstance(image, torch.Tensor):
802
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
803
+
804
+ image = image.to(device=device, dtype=dtype)
805
+ if output_hidden_states:
806
+ image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
807
+ image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
808
+ uncond_image_enc_hidden_states = self.image_encoder(
809
+ torch.zeros_like(image), output_hidden_states=True
810
+ ).hidden_states[-2]
811
+ uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
812
+ num_images_per_prompt, dim=0
813
+ )
814
+ return image_enc_hidden_states, uncond_image_enc_hidden_states
815
+ else:
816
+ image_embeds = self.image_encoder(image).image_embeds
817
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
818
+ uncond_image_embeds = torch.zeros_like(image_embeds)
819
+
820
+ return image_embeds, uncond_image_embeds
821
+
822
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
823
+ def prepare_ip_adapter_image_embeds(
824
+ self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
825
+ ):
826
+ if ip_adapter_image_embeds is None:
827
+ if not isinstance(ip_adapter_image, list):
828
+ ip_adapter_image = [ip_adapter_image]
829
+
830
+ if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
831
+ raise ValueError(
832
+ f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
833
+ )
834
+
835
+ image_embeds = []
836
+ for single_ip_adapter_image, image_proj_layer in zip(
837
+ ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
838
+ ):
839
+ output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
840
+ single_image_embeds, single_negative_image_embeds = self.encode_image(
841
+ single_ip_adapter_image, device, 1, output_hidden_state
842
+ )
843
+ single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0)
844
+ single_negative_image_embeds = torch.stack(
845
+ [single_negative_image_embeds] * num_images_per_prompt, dim=0
846
+ )
847
+
848
+ if do_classifier_free_guidance:
849
+ single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
850
+ single_image_embeds = single_image_embeds.to(device)
851
+
852
+ image_embeds.append(single_image_embeds)
853
+ else:
854
+ repeat_dims = [1]
855
+ image_embeds = []
856
+ for single_image_embeds in ip_adapter_image_embeds:
857
+ if do_classifier_free_guidance:
858
+ single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
859
+ single_image_embeds = single_image_embeds.repeat(
860
+ num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
861
+ )
862
+ single_negative_image_embeds = single_negative_image_embeds.repeat(
863
+ num_images_per_prompt, *(repeat_dims * len(single_negative_image_embeds.shape[1:]))
864
+ )
865
+ single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
866
+ else:
867
+ single_image_embeds = single_image_embeds.repeat(
868
+ num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
869
+ )
870
+ image_embeds.append(single_image_embeds)
871
+
872
+ return image_embeds
873
+
874
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
875
+ def run_safety_checker(self, image, device, dtype):
876
+ if self.safety_checker is None:
877
+ has_nsfw_concept = None
878
+ else:
879
+ if torch.is_tensor(image):
880
+ feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
881
+ else:
882
+ feature_extractor_input = self.image_processor.numpy_to_pil(image)
883
+ safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
884
+ image, has_nsfw_concept = self.safety_checker(
885
+ images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
886
+ )
887
+ return image, has_nsfw_concept
888
+
889
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
890
+ def prepare_extra_step_kwargs(self, generator, eta):
891
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
892
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
893
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
894
+ # and should be between [0, 1]
895
+
896
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
897
+ extra_step_kwargs = {}
898
+ if accepts_eta:
899
+ extra_step_kwargs["eta"] = eta
900
+
901
+ # check if the scheduler accepts generator
902
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
903
+ if accepts_generator:
904
+ extra_step_kwargs["generator"] = generator
905
+ return extra_step_kwargs
906
+
907
+ def check_inputs(
908
+ self,
909
+ prompt,
910
+ image,
911
+ mask_image,
912
+ height,
913
+ width,
914
+ strength,
915
+ callback_steps,
916
+ output_type,
917
+ negative_prompt=None,
918
+ prompt_embeds=None,
919
+ negative_prompt_embeds=None,
920
+ ip_adapter_image=None,
921
+ ip_adapter_image_embeds=None,
922
+ callback_on_step_end_tensor_inputs=None,
923
+ padding_mask_crop=None,
924
+ ):
925
+ if strength < 0 or strength > 1:
926
+ raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
927
+
928
+ if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
929
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
930
+
931
+ if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
932
+ raise ValueError(
933
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
934
+ f" {type(callback_steps)}."
935
+ )
936
+
937
+ if callback_on_step_end_tensor_inputs is not None and not all(
938
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
939
+ ):
940
+ raise ValueError(
941
+ 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]}"
942
+ )
943
+
944
+ if prompt is not None and prompt_embeds is not None:
945
+ raise ValueError(
946
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
947
+ " only forward one of the two."
948
+ )
949
+ elif prompt is None and prompt_embeds is None:
950
+ raise ValueError(
951
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
952
+ )
953
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
954
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
955
+
956
+ if negative_prompt is not None and negative_prompt_embeds is not None:
957
+ raise ValueError(
958
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
959
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
960
+ )
961
+
962
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
963
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
964
+ raise ValueError(
965
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
966
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
967
+ f" {negative_prompt_embeds.shape}."
968
+ )
969
+ if padding_mask_crop is not None:
970
+ if not isinstance(image, PIL.Image.Image):
971
+ raise ValueError(
972
+ f"The image should be a PIL image when inpainting mask crop, but is of type" f" {type(image)}."
973
+ )
974
+ if not isinstance(mask_image, PIL.Image.Image):
975
+ raise ValueError(
976
+ f"The mask image should be a PIL image when inpainting mask crop, but is of type"
977
+ f" {type(mask_image)}."
978
+ )
979
+ if output_type != "pil":
980
+ raise ValueError(f"The output type should be PIL when inpainting mask crop, but is" f" {output_type}.")
981
+
982
+ if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
983
+ raise ValueError(
984
+ "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
985
+ )
986
+
987
+ if ip_adapter_image_embeds is not None:
988
+ if not isinstance(ip_adapter_image_embeds, list):
989
+ raise ValueError(
990
+ f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
991
+ )
992
+ elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
993
+ raise ValueError(
994
+ f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
995
+ )
996
+
997
+ def prepare_latents(
998
+ self,
999
+ batch_size,
1000
+ num_channels_latents,
1001
+ height,
1002
+ width,
1003
+ dtype,
1004
+ device,
1005
+ generator,
1006
+ latents=None,
1007
+ image=None,
1008
+ timestep=None,
1009
+ is_strength_max=True,
1010
+ return_noise=False,
1011
+ return_image_latents=False,
1012
+ ):
1013
+ shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
1014
+ if isinstance(generator, list) and len(generator) != batch_size:
1015
+ raise ValueError(
1016
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
1017
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
1018
+ )
1019
+
1020
+ if (image is None or timestep is None) and not is_strength_max:
1021
+ raise ValueError(
1022
+ "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise."
1023
+ "However, either the image or the noise timestep has not been provided."
1024
+ )
1025
+
1026
+ if return_image_latents or (latents is None and not is_strength_max):
1027
+ image = image.to(device=device, dtype=dtype)
1028
+
1029
+ if image.shape[1] == 4:
1030
+ image_latents = image
1031
+ else:
1032
+ image_latents = self._encode_vae_image(image=image, generator=generator)
1033
+ image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1)
1034
+
1035
+ if latents is None:
1036
+ noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
1037
+ # if strength is 1. then initialise the latents to noise, else initial to image + noise
1038
+ latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep)
1039
+ # if pure noise then scale the initial latents by the Scheduler's init sigma
1040
+ latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents
1041
+ else:
1042
+ noise = latents.to(device)
1043
+ latents = noise * self.scheduler.init_noise_sigma
1044
+
1045
+ outputs = (latents,)
1046
+
1047
+ if return_noise:
1048
+ outputs += (noise,)
1049
+
1050
+ if return_image_latents:
1051
+ outputs += (image_latents,)
1052
+
1053
+ return outputs
1054
+
1055
+ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
1056
+ if isinstance(generator, list):
1057
+ image_latents = [
1058
+ retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
1059
+ for i in range(image.shape[0])
1060
+ ]
1061
+ image_latents = torch.cat(image_latents, dim=0)
1062
+ else:
1063
+ image_latents = retrieve_latents(self.vae.encode(image), generator=generator)
1064
+
1065
+ image_latents = self.vae.config.scaling_factor * image_latents
1066
+
1067
+ return image_latents
1068
+
1069
+ def prepare_mask_latents(
1070
+ self, mask, masked_image, batch_size, height, width, dtype, device, generator
1071
+ ):
1072
+ # resize the mask to latents shape as we concatenate the mask to the latents
1073
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
1074
+ # and half precision
1075
+ mask = torch.nn.functional.interpolate(
1076
+ mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor)
1077
+ )
1078
+ mask = mask.to(device=device, dtype=dtype)
1079
+
1080
+ masked_image = masked_image.to(device=device, dtype=dtype)
1081
+
1082
+ if masked_image.shape[1] == 4:
1083
+ masked_image_latents = masked_image
1084
+ else:
1085
+ masked_image_latents = self._encode_vae_image(masked_image, generator=generator)
1086
+
1087
+ # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
1088
+ if mask.shape[0] < batch_size:
1089
+ if not batch_size % mask.shape[0] == 0:
1090
+ raise ValueError(
1091
+ "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
1092
+ f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"
1093
+ " of masks that you pass is divisible by the total requested batch size."
1094
+ )
1095
+ mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
1096
+ if masked_image_latents.shape[0] < batch_size:
1097
+ if not batch_size % masked_image_latents.shape[0] == 0:
1098
+ raise ValueError(
1099
+ "The passed images and the required batch size don't match. Images are supposed to be duplicated"
1100
+ f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
1101
+ " Make sure the number of images that you pass is divisible by the total requested batch size."
1102
+ )
1103
+ masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)
1104
+
1105
+ # aligning device to prevent device errors when concating it with the latent model input
1106
+ masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
1107
+ return mask, masked_image_latents
1108
+
1109
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps
1110
+ def get_timesteps(self, num_inference_steps, strength, device):
1111
+ # get the original timestep using init_timestep
1112
+ init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
1113
+
1114
+ t_start = max(num_inference_steps - init_timestep, 0)
1115
+ timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
1116
+ if hasattr(self.scheduler, "set_begin_index"):
1117
+ self.scheduler.set_begin_index(t_start * self.scheduler.order)
1118
+
1119
+ return timesteps, num_inference_steps - t_start
1120
+
1121
+ # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
1122
+ def get_guidance_scale_embedding(
1123
+ self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
1124
+ ) -> torch.FloatTensor:
1125
+ """
1126
+ See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
1127
+
1128
+ Args:
1129
+ w (`torch.Tensor`):
1130
+ Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
1131
+ embedding_dim (`int`, *optional*, defaults to 512):
1132
+ Dimension of the embeddings to generate.
1133
+ dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
1134
+ Data type of the generated embeddings.
1135
+
1136
+ Returns:
1137
+ `torch.FloatTensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
1138
+ """
1139
+ assert len(w.shape) == 1
1140
+ w = w * 1000.0
1141
+
1142
+ half_dim = embedding_dim // 2
1143
+ emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
1144
+ emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
1145
+ emb = w.to(dtype)[:, None] * emb[None, :]
1146
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
1147
+ if embedding_dim % 2 == 1: # zero pad
1148
+ emb = torch.nn.functional.pad(emb, (0, 1))
1149
+ assert emb.shape == (w.shape[0], embedding_dim)
1150
+ return emb
1151
+
1152
+ @property
1153
+ def guidance_scale(self):
1154
+ return self._guidance_scale
1155
+
1156
+ @property
1157
+ def clip_skip(self):
1158
+ return self._clip_skip
1159
+
1160
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
1161
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
1162
+ # corresponds to doing no classifier free guidance.
1163
+ @property
1164
+ def do_classifier_free_guidance(self):
1165
+ return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
1166
+
1167
+ @property
1168
+ def cross_attention_kwargs(self):
1169
+ return self._cross_attention_kwargs
1170
+
1171
+ @property
1172
+ def num_timesteps(self):
1173
+ return self._num_timesteps
1174
+
1175
+ @property
1176
+ def interrupt(self):
1177
+ return self._interrupt
1178
+
1179
+ @property
1180
+ def pag_scale(self):
1181
+ return self._pag_scale
1182
+
1183
+ @property
1184
+ def do_perturbed_attention_guidance(self):
1185
+ return self._pag_scale > 0
1186
+
1187
+ @property
1188
+ def pag_adaptive_scaling(self):
1189
+ return self._pag_adaptive_scaling
1190
+
1191
+ @property
1192
+ def do_pag_adaptive_scaling(self):
1193
+ return self._pag_adaptive_scaling > 0
1194
+
1195
+ @property
1196
+ def pag_applied_layers_index(self):
1197
+ return self._pag_applied_layers_index
1198
+
1199
+ @torch.no_grad()
1200
+ def __call__(
1201
+ self,
1202
+ prompt: Union[str, List[str]] = None,
1203
+ image: PipelineImageInput = None,
1204
+ mask_image: PipelineImageInput = None,
1205
+ masked_image_latents: torch.FloatTensor = None,
1206
+ height: Optional[int] = None,
1207
+ width: Optional[int] = None,
1208
+ padding_mask_crop: Optional[int] = None,
1209
+ strength: float = 1.0,
1210
+ num_inference_steps: int = 50,
1211
+ timesteps: List[int] = None,
1212
+ guidance_scale: float = 7.5,
1213
+ pag_scale: float = 0.0,
1214
+ pag_adaptive_scaling: float = 0.0,
1215
+ pag_applied_layers_index: List[str] = ["d4"],
1216
+ negative_prompt: Optional[Union[str, List[str]]] = None,
1217
+ num_images_per_prompt: Optional[int] = 1,
1218
+ eta: float = 0.0,
1219
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
1220
+ latents: Optional[torch.FloatTensor] = None,
1221
+ prompt_embeds: Optional[torch.FloatTensor] = None,
1222
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
1223
+ ip_adapter_image: Optional[PipelineImageInput] = None,
1224
+ ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
1225
+ output_type: Optional[str] = "pil",
1226
+ return_dict: bool = True,
1227
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
1228
+ clip_skip: int = None,
1229
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
1230
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
1231
+ **kwargs,
1232
+ ):
1233
+ r"""
1234
+ The call function to the pipeline for generation.
1235
+
1236
+ Args:
1237
+ prompt (`str` or `List[str]`, *optional*):
1238
+ The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
1239
+ image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
1240
+ `Image`, numpy array or tensor representing an image batch to be inpainted (which parts of the image to
1241
+ be masked out with `mask_image` and repainted according to `prompt`). For both numpy array and pytorch
1242
+ tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the
1243
+ expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the
1244
+ expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but
1245
+ if passing latents directly it is not encoded again.
1246
+ mask_image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
1247
+ `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask
1248
+ are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
1249
+ single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
1250
+ color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
1251
+ H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
1252
+ 1)`, or `(H, W)`.
1253
+ height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
1254
+ The height in pixels of the generated image.
1255
+ width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
1256
+ The width in pixels of the generated image.
1257
+ padding_mask_crop (`int`, *optional*, defaults to `None`):
1258
+ The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to
1259
+ image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region
1260
+ with the same aspect ration of the image and contains all masked area, and then expand that area based
1261
+ on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before
1262
+ resizing to the original image size for inpainting. This is useful when the masked area is small while
1263
+ the image is large and contain information irrelevant for inpainting, such as background.
1264
+ strength (`float`, *optional*, defaults to 1.0):
1265
+ Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
1266
+ starting point and more noise is added the higher the `strength`. The number of denoising steps depends
1267
+ on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
1268
+ process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
1269
+ essentially ignores `image`.
1270
+ num_inference_steps (`int`, *optional*, defaults to 50):
1271
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
1272
+ expense of slower inference. This parameter is modulated by `strength`.
1273
+ timesteps (`List[int]`, *optional*):
1274
+ Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
1275
+ in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
1276
+ passed will be used. Must be in descending order.
1277
+ guidance_scale (`float`, *optional*, defaults to 7.5):
1278
+ A higher guidance scale value encourages the model to generate images closely linked to the text
1279
+ `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
1280
+ negative_prompt (`str` or `List[str]`, *optional*):
1281
+ The prompt or prompts to guide what to not include in image generation. If not defined, you need to
1282
+ pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
1283
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
1284
+ The number of images to generate per prompt.
1285
+ eta (`float`, *optional*, defaults to 0.0):
1286
+ Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
1287
+ to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
1288
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
1289
+ A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
1290
+ generation deterministic.
1291
+ latents (`torch.FloatTensor`, *optional*):
1292
+ Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
1293
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
1294
+ tensor is generated by sampling using the supplied random `generator`.
1295
+ prompt_embeds (`torch.FloatTensor`, *optional*):
1296
+ Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
1297
+ provided, text embeddings are generated from the `prompt` input argument.
1298
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
1299
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
1300
+ not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
1301
+ ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
1302
+ ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
1303
+ Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
1304
+ IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
1305
+ contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
1306
+ provided, embeddings are computed from the `ip_adapter_image` input argument.
1307
+ output_type (`str`, *optional*, defaults to `"pil"`):
1308
+ The output format of the generated image. Choose between `PIL.Image` or `np.array`.
1309
+ return_dict (`bool`, *optional*, defaults to `True`):
1310
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
1311
+ plain tuple.
1312
+ cross_attention_kwargs (`dict`, *optional*):
1313
+ A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
1314
+ [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
1315
+ clip_skip (`int`, *optional*):
1316
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
1317
+ the output of the pre-final layer will be used for computing the prompt embeddings.
1318
+ callback_on_step_end (`Callable`, *optional*):
1319
+ A function that calls at the end of each denoising steps during the inference. The function is called
1320
+ with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
1321
+ callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
1322
+ `callback_on_step_end_tensor_inputs`.
1323
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
1324
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
1325
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
1326
+ `._callback_tensor_inputs` attribute of your pipeline class.
1327
+ Examples:
1328
+
1329
+ ```py
1330
+ >>> import PIL
1331
+ >>> import requests
1332
+ >>> import torch
1333
+ >>> from io import BytesIO
1334
+
1335
+ >>> from diffusers import StableDiffusionInpaintPipeline
1336
+
1337
+
1338
+ >>> def download_image(url):
1339
+ ... response = requests.get(url)
1340
+ ... return PIL.Image.open(BytesIO(response.content)).convert("RGB")
1341
+
1342
+
1343
+ >>> img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
1344
+ >>> mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
1345
+
1346
+ >>> init_image = download_image(img_url).resize((512, 512))
1347
+ >>> mask_image = download_image(mask_url).resize((512, 512))
1348
+
1349
+ >>> pipe = StableDiffusionInpaintPipeline.from_pretrained(
1350
+ ... "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16
1351
+ ... )
1352
+ >>> pipe = pipe.to("cuda")
1353
+
1354
+ >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
1355
+ >>> image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
1356
+ ```
1357
+
1358
+ Returns:
1359
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
1360
+ If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
1361
+ otherwise a `tuple` is returned where the first element is a list with the generated images and the
1362
+ second element is a list of `bool`s indicating whether the corresponding generated image contains
1363
+ "not-safe-for-work" (nsfw) content.
1364
+ """
1365
+
1366
+ callback = kwargs.pop("callback", None)
1367
+ callback_steps = kwargs.pop("callback_steps", None)
1368
+
1369
+ if callback is not None:
1370
+ deprecate(
1371
+ "callback",
1372
+ "1.0.0",
1373
+ "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
1374
+ )
1375
+ if callback_steps is not None:
1376
+ deprecate(
1377
+ "callback_steps",
1378
+ "1.0.0",
1379
+ "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
1380
+ )
1381
+
1382
+ # 0. Default height and width to unet
1383
+ height = height or self.unet.config.sample_size * self.vae_scale_factor
1384
+ width = width or self.unet.config.sample_size * self.vae_scale_factor
1385
+
1386
+ # 1. Check inputs
1387
+ self.check_inputs(
1388
+ prompt,
1389
+ image,
1390
+ mask_image,
1391
+ height,
1392
+ width,
1393
+ strength,
1394
+ callback_steps,
1395
+ output_type,
1396
+ negative_prompt,
1397
+ prompt_embeds,
1398
+ negative_prompt_embeds,
1399
+ ip_adapter_image,
1400
+ ip_adapter_image_embeds,
1401
+ callback_on_step_end_tensor_inputs,
1402
+ padding_mask_crop,
1403
+ )
1404
+
1405
+ self._guidance_scale = guidance_scale
1406
+ self._clip_skip = clip_skip
1407
+ self._cross_attention_kwargs = cross_attention_kwargs
1408
+ self._interrupt = False
1409
+
1410
+ self._pag_scale = pag_scale
1411
+ self._pag_adaptive_scaling = pag_adaptive_scaling
1412
+ self._pag_applied_layers_index = pag_applied_layers_index
1413
+
1414
+ # 2. Define call parameters
1415
+ if prompt is not None and isinstance(prompt, str):
1416
+ batch_size = 1
1417
+ elif prompt is not None and isinstance(prompt, list):
1418
+ batch_size = len(prompt)
1419
+ else:
1420
+ batch_size = prompt_embeds.shape[0]
1421
+
1422
+ device = self._execution_device
1423
+
1424
+ # 3. Encode input prompt
1425
+ text_encoder_lora_scale = (
1426
+ cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
1427
+ )
1428
+ prompt_embeds, negative_prompt_embeds = self.encode_prompt(
1429
+ prompt,
1430
+ device,
1431
+ num_images_per_prompt,
1432
+ self.do_classifier_free_guidance,
1433
+ negative_prompt,
1434
+ prompt_embeds=prompt_embeds,
1435
+ negative_prompt_embeds=negative_prompt_embeds,
1436
+ lora_scale=text_encoder_lora_scale,
1437
+ clip_skip=self.clip_skip,
1438
+ )
1439
+ # For classifier free guidance, we need to do two forward passes.
1440
+ # Here we concatenate the unconditional and text embeddings into a single batch
1441
+ # to avoid doing two forward passes
1442
+
1443
+ # cfg
1444
+ if self.do_classifier_free_guidance and not self.do_perturbed_attention_guidance:
1445
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
1446
+ # pag
1447
+ elif not self.do_classifier_free_guidance and self.do_perturbed_attention_guidance:
1448
+ prompt_embeds = torch.cat([prompt_embeds, prompt_embeds])
1449
+ # both
1450
+ elif self.do_classifier_free_guidance and self.do_perturbed_attention_guidance:
1451
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds, prompt_embeds])
1452
+
1453
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
1454
+ image_embeds = self.prepare_ip_adapter_image_embeds(
1455
+ ip_adapter_image,
1456
+ ip_adapter_image_embeds,
1457
+ device,
1458
+ batch_size * num_images_per_prompt
1459
+ )
1460
+
1461
+ # 4. set timesteps
1462
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
1463
+ timesteps, num_inference_steps = self.get_timesteps(
1464
+ num_inference_steps=num_inference_steps, strength=strength, device=device
1465
+ )
1466
+ # check that number of inference steps is not < 1 - as this doesn't make sense
1467
+ if num_inference_steps < 1:
1468
+ raise ValueError(
1469
+ f"After adjusting the num_inference_steps by strength parameter: {strength}, the number of pipeline"
1470
+ f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline."
1471
+ )
1472
+ # at which timestep to set the initial noise (n.b. 50% if strength is 0.5)
1473
+ latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
1474
+ # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise
1475
+ is_strength_max = strength == 1.0
1476
+
1477
+ # 5. Preprocess mask and image
1478
+
1479
+ if padding_mask_crop is not None:
1480
+ crops_coords = self.mask_processor.get_crop_region(mask_image, width, height, pad=padding_mask_crop)
1481
+ resize_mode = "fill"
1482
+ else:
1483
+ crops_coords = None
1484
+ resize_mode = "default"
1485
+
1486
+ original_image = image
1487
+ init_image = self.image_processor.preprocess(
1488
+ image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode
1489
+ )
1490
+ init_image = init_image.to(dtype=torch.float32)
1491
+
1492
+ # 6. Prepare latent variables
1493
+ num_channels_latents = self.vae.config.latent_channels
1494
+ num_channels_unet = self.unet.config.in_channels
1495
+ return_image_latents = num_channels_unet == 4
1496
+
1497
+ latents_outputs = self.prepare_latents(
1498
+ batch_size * num_images_per_prompt,
1499
+ num_channels_latents,
1500
+ height,
1501
+ width,
1502
+ prompt_embeds.dtype,
1503
+ device,
1504
+ generator,
1505
+ latents,
1506
+ image=init_image,
1507
+ timestep=latent_timestep,
1508
+ is_strength_max=is_strength_max,
1509
+ return_noise=True,
1510
+ return_image_latents=return_image_latents,
1511
+ )
1512
+
1513
+ if return_image_latents:
1514
+ latents, noise, image_latents = latents_outputs
1515
+ else:
1516
+ latents, noise = latents_outputs
1517
+
1518
+ # 7. Prepare mask latent variables
1519
+ mask_condition = self.mask_processor.preprocess(
1520
+ mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords
1521
+ )
1522
+
1523
+ if masked_image_latents is None:
1524
+ masked_image = init_image * (mask_condition < 0.5)
1525
+ else:
1526
+ masked_image = masked_image_latents
1527
+
1528
+ mask, masked_image_latents = self.prepare_mask_latents(
1529
+ mask_condition,
1530
+ masked_image,
1531
+ batch_size * num_images_per_prompt,
1532
+ height,
1533
+ width,
1534
+ prompt_embeds.dtype,
1535
+ device,
1536
+ generator
1537
+ )
1538
+
1539
+ # 8. Check that sizes of mask, masked image and latents match
1540
+ if num_channels_unet == 9:
1541
+ # default case for runwayml/stable-diffusion-inpainting
1542
+ num_channels_mask = mask.shape[1]
1543
+ num_channels_masked_image = masked_image_latents.shape[1]
1544
+ if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels:
1545
+ raise ValueError(
1546
+ f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects"
1547
+ f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +"
1548
+ f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}"
1549
+ f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of"
1550
+ " `pipeline.unet` or your `mask_image` or `image` input."
1551
+ )
1552
+ elif num_channels_unet != 4:
1553
+ raise ValueError(
1554
+ f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}."
1555
+ )
1556
+
1557
+ # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
1558
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
1559
+
1560
+ # 9.1 Add image embeds for IP-Adapter
1561
+ added_cond_kwargs = (
1562
+ {"image_embeds": image_embeds}
1563
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None
1564
+ else None
1565
+ )
1566
+
1567
+ # 9.2 Optionally get Guidance Scale Embedding
1568
+ timestep_cond = None
1569
+ if self.unet.config.time_cond_proj_dim is not None:
1570
+ guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
1571
+ timestep_cond = self.get_guidance_scale_embedding(
1572
+ guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
1573
+ ).to(device=device, dtype=latents.dtype)
1574
+
1575
+ # 10. Denoising loop
1576
+ if self.do_perturbed_attention_guidance:
1577
+ down_layers = []
1578
+ mid_layers = []
1579
+ up_layers = []
1580
+ for name, module in self.unet.named_modules():
1581
+ if "attn1" in name and "to" not in name:
1582
+ layer_type = name.split(".")[0].split("_")[0]
1583
+ if layer_type == "down":
1584
+ down_layers.append(module)
1585
+ elif layer_type == "mid":
1586
+ mid_layers.append(module)
1587
+ elif layer_type == "up":
1588
+ up_layers.append(module)
1589
+ else:
1590
+ raise ValueError(f"Invalid layer type: {layer_type}")
1591
+
1592
+ # change attention layer in UNet if use PAG
1593
+ if self.do_perturbed_attention_guidance:
1594
+ if self.do_classifier_free_guidance:
1595
+ replace_processor = PAGCFGIdentitySelfAttnProcessor()
1596
+ else:
1597
+ replace_processor = PAGIdentitySelfAttnProcessor()
1598
+
1599
+ drop_layers = self.pag_applied_layers_index
1600
+ for drop_layer in drop_layers:
1601
+ try:
1602
+ if drop_layer[0] == "d":
1603
+ down_layers[int(drop_layer[1])].processor = replace_processor
1604
+ elif drop_layer[0] == "m":
1605
+ mid_layers[int(drop_layer[1])].processor = replace_processor
1606
+ elif drop_layer[0] == "u":
1607
+ up_layers[int(drop_layer[1])].processor = replace_processor
1608
+ else:
1609
+ raise ValueError(f"Invalid layer type: {drop_layer[0]}")
1610
+ except IndexError:
1611
+ raise ValueError(
1612
+ f"Invalid layer index: {drop_layer}. Available layers: {len(down_layers)} down layers, {len(mid_layers)} mid layers, {len(up_layers)} up layers."
1613
+ )
1614
+
1615
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
1616
+ self._num_timesteps = len(timesteps)
1617
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
1618
+ for i, t in enumerate(timesteps):
1619
+ if self.interrupt:
1620
+ continue
1621
+
1622
+ # cfg
1623
+ if self.do_classifier_free_guidance and not self.do_perturbed_attention_guidance:
1624
+ latent_model_input = torch.cat([latents] * 2)
1625
+ mask_input = torch.cat([mask] * 2)
1626
+ masked_image_latents_input = torch.cat([masked_image_latents] * 2)
1627
+ # pag
1628
+ elif not self.do_classifier_free_guidance and self.do_perturbed_attention_guidance:
1629
+ latent_model_input = torch.cat([latents] * 2)
1630
+ mask_input = torch.cat([mask] * 2)
1631
+ masked_image_latents_input = torch.cat([masked_image_latents] * 2)
1632
+ # both
1633
+ elif self.do_classifier_free_guidance and self.do_perturbed_attention_guidance:
1634
+ latent_model_input = torch.cat([latents] * 3)
1635
+ mask_input = torch.cat([mask] * 3)
1636
+ masked_image_latents_input = torch.cat([masked_image_latents] * 3)
1637
+ # no
1638
+ else:
1639
+ latent_model_input = latents
1640
+ mask_input = mask
1641
+ masked_image_latents_input = masked_image_latents
1642
+
1643
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
1644
+
1645
+ if num_channels_unet == 9:
1646
+
1647
+ latent_model_input = torch.cat([latent_model_input, mask_input, masked_image_latents_input], dim=1)
1648
+
1649
+ # predict the noise residual
1650
+ noise_pred = self.unet(
1651
+ latent_model_input,
1652
+ t,
1653
+ encoder_hidden_states=prompt_embeds,
1654
+ timestep_cond=timestep_cond,
1655
+ cross_attention_kwargs=self.cross_attention_kwargs,
1656
+ added_cond_kwargs=added_cond_kwargs,
1657
+ return_dict=False,
1658
+ )[0]
1659
+
1660
+ # perform guidance
1661
+
1662
+ # cfg
1663
+ if self.do_classifier_free_guidance and not self.do_perturbed_attention_guidance:
1664
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
1665
+
1666
+ delta = noise_pred_text - noise_pred_uncond
1667
+ noise_pred = noise_pred_uncond + self.guidance_scale * delta
1668
+
1669
+ # pag
1670
+ elif not self.do_classifier_free_guidance and self.do_perturbed_attention_guidance:
1671
+ noise_pred_original, noise_pred_perturb = noise_pred.chunk(2)
1672
+
1673
+ signal_scale = self.pag_scale
1674
+ if self.do_pag_adaptive_scaling:
1675
+ signal_scale = self.pag_scale - self.pag_adaptive_scaling * (1000 - t)
1676
+ if signal_scale < 0:
1677
+ signal_scale = 0
1678
+
1679
+ noise_pred = noise_pred_original + signal_scale * (noise_pred_original - noise_pred_perturb)
1680
+
1681
+ # both
1682
+ elif self.do_classifier_free_guidance and self.do_perturbed_attention_guidance:
1683
+ noise_pred_uncond, noise_pred_text, noise_pred_text_perturb = noise_pred.chunk(3)
1684
+
1685
+ signal_scale = self.pag_scale
1686
+ if self.do_pag_adaptive_scaling:
1687
+ signal_scale = self.pag_scale - self.pag_adaptive_scaling * (1000 - t)
1688
+ if signal_scale < 0:
1689
+ signal_scale = 0
1690
+
1691
+ noise_pred = (
1692
+ noise_pred_text
1693
+ + (self.guidance_scale - 1.0) * (noise_pred_text - noise_pred_uncond)
1694
+ + signal_scale * (noise_pred_text - noise_pred_text_perturb)
1695
+ )
1696
+
1697
+ # compute the previous noisy sample x_t -> x_t-1
1698
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
1699
+ if num_channels_unet == 4:
1700
+ init_latents_proper = image_latents
1701
+ if self.do_classifier_free_guidance:
1702
+ init_mask, _ = mask.chunk(2)
1703
+ else:
1704
+ init_mask = mask
1705
+
1706
+ if i < len(timesteps) - 1:
1707
+ noise_timestep = timesteps[i + 1]
1708
+ init_latents_proper = self.scheduler.add_noise(
1709
+ init_latents_proper, noise, torch.tensor([noise_timestep])
1710
+ )
1711
+
1712
+ latents = (1 - init_mask) * init_latents_proper + init_mask * latents
1713
+
1714
+ if callback_on_step_end is not None:
1715
+ callback_kwargs = {}
1716
+ for k in callback_on_step_end_tensor_inputs:
1717
+ callback_kwargs[k] = locals()[k]
1718
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
1719
+
1720
+ latents = callback_outputs.pop("latents", latents)
1721
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
1722
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
1723
+ mask = callback_outputs.pop("mask", mask)
1724
+ masked_image_latents = callback_outputs.pop("masked_image_latents", masked_image_latents)
1725
+
1726
+ # call the callback, if provided
1727
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
1728
+ progress_bar.update()
1729
+ if callback is not None and i % callback_steps == 0:
1730
+ step_idx = i // getattr(self.scheduler, "order", 1)
1731
+ callback(step_idx, t, latents)
1732
+
1733
+ if not output_type == "latent":
1734
+ condition_kwargs = {}
1735
+ if isinstance(self.vae, AsymmetricAutoencoderKL):
1736
+ init_image = init_image.to(device=device, dtype=masked_image_latents.dtype)
1737
+ init_image_condition = init_image.clone()
1738
+ init_image = self._encode_vae_image(init_image, generator=generator)
1739
+ mask_condition = mask_condition.to(device=device, dtype=masked_image_latents.dtype)
1740
+ condition_kwargs = {"image": init_image_condition, "mask": mask_condition}
1741
+ image = self.vae.decode(
1742
+ latents / self.vae.config.scaling_factor, return_dict=False, generator=generator, **condition_kwargs
1743
+ )[0]
1744
+ image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
1745
+ else:
1746
+ image = latents
1747
+ has_nsfw_concept = None
1748
+
1749
+ if has_nsfw_concept is None:
1750
+ do_denormalize = [True] * image.shape[0]
1751
+ else:
1752
+ do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
1753
+
1754
+ image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
1755
+
1756
+ if padding_mask_crop is not None:
1757
+ image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image]
1758
+
1759
+ # Offload all models
1760
+ self.maybe_free_model_hooks()
1761
+
1762
+ # change attention layer in UNet if use PAG
1763
+ if self.do_perturbed_attention_guidance:
1764
+ drop_layers = self.pag_applied_layers_index
1765
+ for drop_layer in drop_layers:
1766
+ try:
1767
+ if drop_layer[0] == "d":
1768
+ down_layers[int(drop_layer[1])].processor = AttnProcessor2_0()
1769
+ elif drop_layer[0] == "m":
1770
+ mid_layers[int(drop_layer[1])].processor = AttnProcessor2_0()
1771
+ elif drop_layer[0] == "u":
1772
+ up_layers[int(drop_layer[1])].processor = AttnProcessor2_0()
1773
+ else:
1774
+ raise ValueError(f"Invalid layer type: {drop_layer[0]}")
1775
+ except IndexError:
1776
+ raise ValueError(
1777
+ f"Invalid layer index: {drop_layer}. Available layers: {len(down_layers)} down layers, {len(mid_layers)} mid layers, {len(up_layers)} up layers."
1778
+ )
1779
+
1780
+ if not return_dict:
1781
+ return (image, has_nsfw_concept)
1782
+
1783
+ return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)