zhaozhilin commited on
Commit
f77be9b
·
1 Parent(s): 9629b0c

add lpw_stable_diffusion_onnx.py

Browse files
Files changed (2) hide show
  1. pipeline.py +0 -0
  2. pipeline_lpw_stable_diffusion_onnx.py +1149 -0
pipeline.py CHANGED
The diff for this file is too large to render. See raw diff
 
pipeline_lpw_stable_diffusion_onnx.py ADDED
@@ -0,0 +1,1149 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # source https://github.com/huggingface/diffusers/blob/main/examples/community/lpw_stable_diffusion_onnx.py
2
+ import inspect
3
+ import re
4
+ from typing import Callable, List, Optional, Union
5
+
6
+ import numpy as np
7
+ import PIL.Image
8
+ import torch
9
+ from packaging import version
10
+ from transformers import CLIPImageProcessor, CLIPTokenizer
11
+
12
+ import diffusers
13
+ from diffusers import OnnxRuntimeModel, OnnxStableDiffusionPipeline, SchedulerMixin
14
+ from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
15
+ from diffusers.utils import logging
16
+
17
+
18
+ try:
19
+ from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE
20
+ except ImportError:
21
+ ORT_TO_NP_TYPE = {
22
+ "tensor(bool)": np.bool_,
23
+ "tensor(int8)": np.int8,
24
+ "tensor(uint8)": np.uint8,
25
+ "tensor(int16)": np.int16,
26
+ "tensor(uint16)": np.uint16,
27
+ "tensor(int32)": np.int32,
28
+ "tensor(uint32)": np.uint32,
29
+ "tensor(int64)": np.int64,
30
+ "tensor(uint64)": np.uint64,
31
+ "tensor(float16)": np.float16,
32
+ "tensor(float)": np.float32,
33
+ "tensor(double)": np.float64,
34
+ }
35
+
36
+ try:
37
+ from diffusers.utils import PIL_INTERPOLATION
38
+ except ImportError:
39
+ if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"):
40
+ PIL_INTERPOLATION = {
41
+ "linear": PIL.Image.Resampling.BILINEAR,
42
+ "bilinear": PIL.Image.Resampling.BILINEAR,
43
+ "bicubic": PIL.Image.Resampling.BICUBIC,
44
+ "lanczos": PIL.Image.Resampling.LANCZOS,
45
+ "nearest": PIL.Image.Resampling.NEAREST,
46
+ }
47
+ else:
48
+ PIL_INTERPOLATION = {
49
+ "linear": PIL.Image.LINEAR,
50
+ "bilinear": PIL.Image.BILINEAR,
51
+ "bicubic": PIL.Image.BICUBIC,
52
+ "lanczos": PIL.Image.LANCZOS,
53
+ "nearest": PIL.Image.NEAREST,
54
+ }
55
+ # ------------------------------------------------------------------------------
56
+
57
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
58
+
59
+ re_attention = re.compile(
60
+ r"""
61
+ \\\(|
62
+ \\\)|
63
+ \\\[|
64
+ \\]|
65
+ \\\\|
66
+ \\|
67
+ \(|
68
+ \[|
69
+ :([+-]?[.\d]+)\)|
70
+ \)|
71
+ ]|
72
+ [^\\()\[\]:]+|
73
+ :
74
+ """,
75
+ re.X,
76
+ )
77
+
78
+
79
+ def parse_prompt_attention(text):
80
+ """
81
+ Parses a string with attention tokens and returns a list of pairs: text and its associated weight.
82
+ Accepted tokens are:
83
+ (abc) - increases attention to abc by a multiplier of 1.1
84
+ (abc:3.12) - increases attention to abc by a multiplier of 3.12
85
+ [abc] - decreases attention to abc by a multiplier of 1.1
86
+ \\( - literal character '('
87
+ \\[ - literal character '['
88
+ \\) - literal character ')'
89
+ \\] - literal character ']'
90
+ \\ - literal character '\'
91
+ anything else - just text
92
+ >>> parse_prompt_attention('normal text')
93
+ [['normal text', 1.0]]
94
+ >>> parse_prompt_attention('an (important) word')
95
+ [['an ', 1.0], ['important', 1.1], [' word', 1.0]]
96
+ >>> parse_prompt_attention('(unbalanced')
97
+ [['unbalanced', 1.1]]
98
+ >>> parse_prompt_attention('\\(literal\\]')
99
+ [['(literal]', 1.0]]
100
+ >>> parse_prompt_attention('(unnecessary)(parens)')
101
+ [['unnecessaryparens', 1.1]]
102
+ >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')
103
+ [['a ', 1.0],
104
+ ['house', 1.5730000000000004],
105
+ [' ', 1.1],
106
+ ['on', 1.0],
107
+ [' a ', 1.1],
108
+ ['hill', 0.55],
109
+ [', sun, ', 1.1],
110
+ ['sky', 1.4641000000000006],
111
+ ['.', 1.1]]
112
+ """
113
+
114
+ res = []
115
+ round_brackets = []
116
+ square_brackets = []
117
+
118
+ round_bracket_multiplier = 1.1
119
+ square_bracket_multiplier = 1 / 1.1
120
+
121
+ def multiply_range(start_position, multiplier):
122
+ for p in range(start_position, len(res)):
123
+ res[p][1] *= multiplier
124
+
125
+ for m in re_attention.finditer(text):
126
+ text = m.group(0)
127
+ weight = m.group(1)
128
+
129
+ if text.startswith("\\"):
130
+ res.append([text[1:], 1.0])
131
+ elif text == "(":
132
+ round_brackets.append(len(res))
133
+ elif text == "[":
134
+ square_brackets.append(len(res))
135
+ elif weight is not None and len(round_brackets) > 0:
136
+ multiply_range(round_brackets.pop(), float(weight))
137
+ elif text == ")" and len(round_brackets) > 0:
138
+ multiply_range(round_brackets.pop(), round_bracket_multiplier)
139
+ elif text == "]" and len(square_brackets) > 0:
140
+ multiply_range(square_brackets.pop(), square_bracket_multiplier)
141
+ else:
142
+ res.append([text, 1.0])
143
+
144
+ for pos in round_brackets:
145
+ multiply_range(pos, round_bracket_multiplier)
146
+
147
+ for pos in square_brackets:
148
+ multiply_range(pos, square_bracket_multiplier)
149
+
150
+ if len(res) == 0:
151
+ res = [["", 1.0]]
152
+
153
+ # merge runs of identical weights
154
+ i = 0
155
+ while i + 1 < len(res):
156
+ if res[i][1] == res[i + 1][1]:
157
+ res[i][0] += res[i + 1][0]
158
+ res.pop(i + 1)
159
+ else:
160
+ i += 1
161
+
162
+ return res
163
+
164
+
165
+ def get_prompts_with_weights(pipe, prompt: List[str], max_length: int):
166
+ r"""
167
+ Tokenize a list of prompts and return its tokens with weights of each token.
168
+
169
+ No padding, starting or ending token is included.
170
+ """
171
+ tokens = []
172
+ weights = []
173
+ truncated = False
174
+ for text in prompt:
175
+ texts_and_weights = parse_prompt_attention(text)
176
+ text_token = []
177
+ text_weight = []
178
+ for word, weight in texts_and_weights:
179
+ # tokenize and discard the starting and the ending token
180
+ token = pipe.tokenizer(word, return_tensors="np").input_ids[0, 1:-1]
181
+ text_token += list(token)
182
+ # copy the weight by length of token
183
+ text_weight += [weight] * len(token)
184
+ # stop if the text is too long (longer than truncation limit)
185
+ if len(text_token) > max_length:
186
+ truncated = True
187
+ break
188
+ # truncate
189
+ if len(text_token) > max_length:
190
+ truncated = True
191
+ text_token = text_token[:max_length]
192
+ text_weight = text_weight[:max_length]
193
+ tokens.append(text_token)
194
+ weights.append(text_weight)
195
+ if truncated:
196
+ logger.warning("Prompt was truncated. Try to shorten the prompt or increase max_embeddings_multiples")
197
+ return tokens, weights
198
+
199
+
200
+ def pad_tokens_and_weights(tokens, weights, max_length, bos, eos, pad, no_boseos_middle=True, chunk_length=77):
201
+ r"""
202
+ Pad the tokens (with starting and ending tokens) and weights (with 1.0) to max_length.
203
+ """
204
+ max_embeddings_multiples = (max_length - 2) // (chunk_length - 2)
205
+ weights_length = max_length if no_boseos_middle else max_embeddings_multiples * chunk_length
206
+ for i in range(len(tokens)):
207
+ tokens[i] = [bos] + tokens[i] + [pad] * (max_length - 1 - len(tokens[i]) - 1) + [eos]
208
+ if no_boseos_middle:
209
+ weights[i] = [1.0] + weights[i] + [1.0] * (max_length - 1 - len(weights[i]))
210
+ else:
211
+ w = []
212
+ if len(weights[i]) == 0:
213
+ w = [1.0] * weights_length
214
+ else:
215
+ for j in range(max_embeddings_multiples):
216
+ w.append(1.0) # weight for starting token in this chunk
217
+ w += weights[i][j * (chunk_length - 2) : min(len(weights[i]), (j + 1) * (chunk_length - 2))]
218
+ w.append(1.0) # weight for ending token in this chunk
219
+ w += [1.0] * (weights_length - len(w))
220
+ weights[i] = w[:]
221
+
222
+ return tokens, weights
223
+
224
+
225
+ def get_unweighted_text_embeddings(
226
+ pipe,
227
+ text_input: np.array,
228
+ chunk_length: int,
229
+ no_boseos_middle: Optional[bool] = True,
230
+ ):
231
+ """
232
+ When the length of tokens is a multiple of the capacity of the text encoder,
233
+ it should be split into chunks and sent to the text encoder individually.
234
+ """
235
+ max_embeddings_multiples = (text_input.shape[1] - 2) // (chunk_length - 2)
236
+ if max_embeddings_multiples > 1:
237
+ text_embeddings = []
238
+ for i in range(max_embeddings_multiples):
239
+ # extract the i-th chunk
240
+ text_input_chunk = text_input[:, i * (chunk_length - 2) : (i + 1) * (chunk_length - 2) + 2].copy()
241
+
242
+ # cover the head and the tail by the starting and the ending tokens
243
+ text_input_chunk[:, 0] = text_input[0, 0]
244
+ text_input_chunk[:, -1] = text_input[0, -1]
245
+
246
+ text_embedding = pipe.text_encoder(input_ids=text_input_chunk)[0]
247
+
248
+ if no_boseos_middle:
249
+ if i == 0:
250
+ # discard the ending token
251
+ text_embedding = text_embedding[:, :-1]
252
+ elif i == max_embeddings_multiples - 1:
253
+ # discard the starting token
254
+ text_embedding = text_embedding[:, 1:]
255
+ else:
256
+ # discard both starting and ending tokens
257
+ text_embedding = text_embedding[:, 1:-1]
258
+
259
+ text_embeddings.append(text_embedding)
260
+ text_embeddings = np.concatenate(text_embeddings, axis=1)
261
+ else:
262
+ text_embeddings = pipe.text_encoder(input_ids=text_input)[0]
263
+ return text_embeddings
264
+
265
+
266
+ def get_weighted_text_embeddings(
267
+ pipe,
268
+ prompt: Union[str, List[str]],
269
+ uncond_prompt: Optional[Union[str, List[str]]] = None,
270
+ max_embeddings_multiples: Optional[int] = 4,
271
+ no_boseos_middle: Optional[bool] = False,
272
+ skip_parsing: Optional[bool] = False,
273
+ skip_weighting: Optional[bool] = False,
274
+ **kwargs,
275
+ ):
276
+ r"""
277
+ Prompts can be assigned with local weights using brackets. For example,
278
+ prompt 'A (very beautiful) masterpiece' highlights the words 'very beautiful',
279
+ and the embedding tokens corresponding to the words get multiplied by a constant, 1.1.
280
+
281
+ Also, to regularize of the embedding, the weighted embedding would be scaled to preserve the original mean.
282
+
283
+ Args:
284
+ pipe (`OnnxStableDiffusionPipeline`):
285
+ Pipe to provide access to the tokenizer and the text encoder.
286
+ prompt (`str` or `List[str]`):
287
+ The prompt or prompts to guide the image generation.
288
+ uncond_prompt (`str` or `List[str]`):
289
+ The unconditional prompt or prompts for guide the image generation. If unconditional prompt
290
+ is provided, the embeddings of prompt and uncond_prompt are concatenated.
291
+ max_embeddings_multiples (`int`, *optional*, defaults to `1`):
292
+ The max multiple length of prompt embeddings compared to the max output length of text encoder.
293
+ no_boseos_middle (`bool`, *optional*, defaults to `False`):
294
+ If the length of text token is multiples of the capacity of text encoder, whether reserve the starting and
295
+ ending token in each of the chunk in the middle.
296
+ skip_parsing (`bool`, *optional*, defaults to `False`):
297
+ Skip the parsing of brackets.
298
+ skip_weighting (`bool`, *optional*, defaults to `False`):
299
+ Skip the weighting. When the parsing is skipped, it is forced True.
300
+ """
301
+ max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2
302
+ if isinstance(prompt, str):
303
+ prompt = [prompt]
304
+
305
+ if not skip_parsing:
306
+ prompt_tokens, prompt_weights = get_prompts_with_weights(pipe, prompt, max_length - 2)
307
+ if uncond_prompt is not None:
308
+ if isinstance(uncond_prompt, str):
309
+ uncond_prompt = [uncond_prompt]
310
+ uncond_tokens, uncond_weights = get_prompts_with_weights(pipe, uncond_prompt, max_length - 2)
311
+ else:
312
+ prompt_tokens = [
313
+ token[1:-1]
314
+ for token in pipe.tokenizer(prompt, max_length=max_length, truncation=True, return_tensors="np").input_ids
315
+ ]
316
+ prompt_weights = [[1.0] * len(token) for token in prompt_tokens]
317
+ if uncond_prompt is not None:
318
+ if isinstance(uncond_prompt, str):
319
+ uncond_prompt = [uncond_prompt]
320
+ uncond_tokens = [
321
+ token[1:-1]
322
+ for token in pipe.tokenizer(
323
+ uncond_prompt,
324
+ max_length=max_length,
325
+ truncation=True,
326
+ return_tensors="np",
327
+ ).input_ids
328
+ ]
329
+ uncond_weights = [[1.0] * len(token) for token in uncond_tokens]
330
+
331
+ # round up the longest length of tokens to a multiple of (model_max_length - 2)
332
+ max_length = max([len(token) for token in prompt_tokens])
333
+ if uncond_prompt is not None:
334
+ max_length = max(max_length, max([len(token) for token in uncond_tokens]))
335
+
336
+ max_embeddings_multiples = min(
337
+ max_embeddings_multiples,
338
+ (max_length - 1) // (pipe.tokenizer.model_max_length - 2) + 1,
339
+ )
340
+ max_embeddings_multiples = max(1, max_embeddings_multiples)
341
+ max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2
342
+
343
+ # pad the length of tokens and weights
344
+ bos = pipe.tokenizer.bos_token_id
345
+ eos = pipe.tokenizer.eos_token_id
346
+ pad = getattr(pipe.tokenizer, "pad_token_id", eos)
347
+ prompt_tokens, prompt_weights = pad_tokens_and_weights(
348
+ prompt_tokens,
349
+ prompt_weights,
350
+ max_length,
351
+ bos,
352
+ eos,
353
+ pad,
354
+ no_boseos_middle=no_boseos_middle,
355
+ chunk_length=pipe.tokenizer.model_max_length,
356
+ )
357
+ prompt_tokens = np.array(prompt_tokens, dtype=np.int32)
358
+ if uncond_prompt is not None:
359
+ uncond_tokens, uncond_weights = pad_tokens_and_weights(
360
+ uncond_tokens,
361
+ uncond_weights,
362
+ max_length,
363
+ bos,
364
+ eos,
365
+ pad,
366
+ no_boseos_middle=no_boseos_middle,
367
+ chunk_length=pipe.tokenizer.model_max_length,
368
+ )
369
+ uncond_tokens = np.array(uncond_tokens, dtype=np.int32)
370
+
371
+ # get the embeddings
372
+ text_embeddings = get_unweighted_text_embeddings(
373
+ pipe,
374
+ prompt_tokens,
375
+ pipe.tokenizer.model_max_length,
376
+ no_boseos_middle=no_boseos_middle,
377
+ )
378
+ prompt_weights = np.array(prompt_weights, dtype=text_embeddings.dtype)
379
+ if uncond_prompt is not None:
380
+ uncond_embeddings = get_unweighted_text_embeddings(
381
+ pipe,
382
+ uncond_tokens,
383
+ pipe.tokenizer.model_max_length,
384
+ no_boseos_middle=no_boseos_middle,
385
+ )
386
+ uncond_weights = np.array(uncond_weights, dtype=uncond_embeddings.dtype)
387
+
388
+ # assign weights to the prompts and normalize in the sense of mean
389
+ # TODO: should we normalize by chunk or in a whole (current implementation)?
390
+ if (not skip_parsing) and (not skip_weighting):
391
+ previous_mean = text_embeddings.mean(axis=(-2, -1))
392
+ text_embeddings *= prompt_weights[:, :, None]
393
+ text_embeddings *= (previous_mean / text_embeddings.mean(axis=(-2, -1)))[:, None, None]
394
+ if uncond_prompt is not None:
395
+ previous_mean = uncond_embeddings.mean(axis=(-2, -1))
396
+ uncond_embeddings *= uncond_weights[:, :, None]
397
+ uncond_embeddings *= (previous_mean / uncond_embeddings.mean(axis=(-2, -1)))[:, None, None]
398
+
399
+ # For classifier free guidance, we need to do two forward passes.
400
+ # Here we concatenate the unconditional and text embeddings into a single batch
401
+ # to avoid doing two forward passes
402
+ if uncond_prompt is not None:
403
+ return text_embeddings, uncond_embeddings
404
+
405
+ return text_embeddings
406
+
407
+
408
+ def preprocess_image(image):
409
+ w, h = image.size
410
+ w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32
411
+ image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"])
412
+ image = np.array(image).astype(np.float32) / 255.0
413
+ image = image[None].transpose(0, 3, 1, 2)
414
+ return 2.0 * image - 1.0
415
+
416
+
417
+ def preprocess_mask(mask, scale_factor=8):
418
+ mask = mask.convert("L")
419
+ w, h = mask.size
420
+ w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32
421
+ mask = mask.resize((w // scale_factor, h // scale_factor), resample=PIL_INTERPOLATION["nearest"])
422
+ mask = np.array(mask).astype(np.float32) / 255.0
423
+ mask = np.tile(mask, (4, 1, 1))
424
+ mask = mask[None].transpose(0, 1, 2, 3) # what does this step do?
425
+ mask = 1 - mask # repaint white, keep black
426
+ return mask
427
+
428
+
429
+ class OnnxStableDiffusionLongPromptWeightingPipeline(OnnxStableDiffusionPipeline):
430
+ r"""
431
+ Pipeline for text-to-image generation using Stable Diffusion without tokens length limit, and support parsing
432
+ weighting in prompt.
433
+
434
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
435
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
436
+ """
437
+
438
+ if version.parse(version.parse(diffusers.__version__).base_version) >= version.parse("0.9.0"):
439
+
440
+ def __init__(
441
+ self,
442
+ vae_encoder: OnnxRuntimeModel,
443
+ vae_decoder: OnnxRuntimeModel,
444
+ text_encoder: OnnxRuntimeModel,
445
+ tokenizer: CLIPTokenizer,
446
+ unet: OnnxRuntimeModel,
447
+ scheduler: SchedulerMixin,
448
+ safety_checker: OnnxRuntimeModel,
449
+ feature_extractor: CLIPImageProcessor,
450
+ requires_safety_checker: bool = True,
451
+ ):
452
+ super().__init__(
453
+ vae_encoder=vae_encoder,
454
+ vae_decoder=vae_decoder,
455
+ text_encoder=text_encoder,
456
+ tokenizer=tokenizer,
457
+ unet=unet,
458
+ scheduler=scheduler,
459
+ safety_checker=safety_checker,
460
+ feature_extractor=feature_extractor,
461
+ requires_safety_checker=requires_safety_checker,
462
+ )
463
+ self.__init__additional__()
464
+
465
+ else:
466
+
467
+ def __init__(
468
+ self,
469
+ vae_encoder: OnnxRuntimeModel,
470
+ vae_decoder: OnnxRuntimeModel,
471
+ text_encoder: OnnxRuntimeModel,
472
+ tokenizer: CLIPTokenizer,
473
+ unet: OnnxRuntimeModel,
474
+ scheduler: SchedulerMixin,
475
+ safety_checker: OnnxRuntimeModel,
476
+ feature_extractor: CLIPImageProcessor,
477
+ ):
478
+ super().__init__(
479
+ vae_encoder=vae_encoder,
480
+ vae_decoder=vae_decoder,
481
+ text_encoder=text_encoder,
482
+ tokenizer=tokenizer,
483
+ unet=unet,
484
+ scheduler=scheduler,
485
+ safety_checker=safety_checker,
486
+ feature_extractor=feature_extractor,
487
+ )
488
+ self.__init__additional__()
489
+
490
+ def __init__additional__(self):
491
+ self.unet.config.in_channels = 4
492
+ self.vae_scale_factor = 8
493
+
494
+ def encode_prompt(
495
+ self,
496
+ prompt,
497
+ num_images_per_prompt,
498
+ do_classifier_free_guidance,
499
+ negative_prompt,
500
+ max_embeddings_multiples,
501
+ ):
502
+ r"""
503
+ Encodes the prompt into text encoder hidden states.
504
+
505
+ Args:
506
+ prompt (`str` or `list(int)`):
507
+ prompt to be encoded
508
+ num_images_per_prompt (`int`):
509
+ number of images that should be generated per prompt
510
+ do_classifier_free_guidance (`bool`):
511
+ whether to use classifier free guidance or not
512
+ negative_prompt (`str` or `List[str]`):
513
+ The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
514
+ if `guidance_scale` is less than `1`).
515
+ max_embeddings_multiples (`int`, *optional*, defaults to `3`):
516
+ The max multiple length of prompt embeddings compared to the max output length of text encoder.
517
+ """
518
+ batch_size = len(prompt) if isinstance(prompt, list) else 1
519
+
520
+ if negative_prompt is None:
521
+ negative_prompt = [""] * batch_size
522
+ elif isinstance(negative_prompt, str):
523
+ negative_prompt = [negative_prompt] * batch_size
524
+ if batch_size != len(negative_prompt):
525
+ raise ValueError(
526
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
527
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
528
+ " the batch size of `prompt`."
529
+ )
530
+
531
+ text_embeddings, uncond_embeddings = get_weighted_text_embeddings(
532
+ pipe=self,
533
+ prompt=prompt,
534
+ uncond_prompt=negative_prompt if do_classifier_free_guidance else None,
535
+ max_embeddings_multiples=max_embeddings_multiples,
536
+ )
537
+
538
+ text_embeddings = text_embeddings.repeat(num_images_per_prompt, 0)
539
+ if do_classifier_free_guidance:
540
+ uncond_embeddings = uncond_embeddings.repeat(num_images_per_prompt, 0)
541
+ text_embeddings = np.concatenate([uncond_embeddings, text_embeddings])
542
+
543
+ return text_embeddings
544
+
545
+ def check_inputs(self, prompt, height, width, strength, callback_steps):
546
+ if not isinstance(prompt, str) and not isinstance(prompt, list):
547
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
548
+
549
+ if strength < 0 or strength > 1:
550
+ raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
551
+
552
+ if height % 8 != 0 or width % 8 != 0:
553
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
554
+
555
+ if (callback_steps is None) or (
556
+ callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
557
+ ):
558
+ raise ValueError(
559
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
560
+ f" {type(callback_steps)}."
561
+ )
562
+
563
+ def get_timesteps(self, num_inference_steps, strength, is_text2img):
564
+ if is_text2img:
565
+ return self.scheduler.timesteps, num_inference_steps
566
+ else:
567
+ # get the original timestep using init_timestep
568
+ offset = self.scheduler.config.get("steps_offset", 0)
569
+ init_timestep = int(num_inference_steps * strength) + offset
570
+ init_timestep = min(init_timestep, num_inference_steps)
571
+
572
+ t_start = max(num_inference_steps - init_timestep + offset, 0)
573
+ timesteps = self.scheduler.timesteps[t_start:]
574
+ return timesteps, num_inference_steps - t_start
575
+
576
+ def run_safety_checker(self, image):
577
+ if self.safety_checker is not None:
578
+ safety_checker_input = self.feature_extractor(
579
+ self.numpy_to_pil(image), return_tensors="np"
580
+ ).pixel_values.astype(image.dtype)
581
+ # There will throw an error if use safety_checker directly and batchsize>1
582
+ images, has_nsfw_concept = [], []
583
+ for i in range(image.shape[0]):
584
+ image_i, has_nsfw_concept_i = self.safety_checker(
585
+ clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1]
586
+ )
587
+ images.append(image_i)
588
+ has_nsfw_concept.append(has_nsfw_concept_i[0])
589
+ image = np.concatenate(images)
590
+ else:
591
+ has_nsfw_concept = None
592
+ return image, has_nsfw_concept
593
+
594
+ def decode_latents(self, latents):
595
+ latents = 1 / 0.18215 * latents
596
+ # image = self.vae_decoder(latent_sample=latents)[0]
597
+ # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
598
+ image = np.concatenate(
599
+ [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]
600
+ )
601
+ image = np.clip(image / 2 + 0.5, 0, 1)
602
+ image = image.transpose((0, 2, 3, 1))
603
+ return image
604
+
605
+ def prepare_extra_step_kwargs(self, generator, eta):
606
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
607
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
608
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
609
+ # and should be between [0, 1]
610
+
611
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
612
+ extra_step_kwargs = {}
613
+ if accepts_eta:
614
+ extra_step_kwargs["eta"] = eta
615
+
616
+ # check if the scheduler accepts generator
617
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
618
+ if accepts_generator:
619
+ extra_step_kwargs["generator"] = generator
620
+ return extra_step_kwargs
621
+
622
+ def prepare_latents(self, image, timestep, batch_size, height, width, dtype, generator, latents=None):
623
+ if image is None:
624
+ shape = (
625
+ batch_size,
626
+ self.unet.config.in_channels,
627
+ height // self.vae_scale_factor,
628
+ width // self.vae_scale_factor,
629
+ )
630
+
631
+ if latents is None:
632
+ latents = torch.randn(shape, generator=generator, device="cpu").numpy().astype(dtype)
633
+ else:
634
+ if latents.shape != shape:
635
+ raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
636
+
637
+ # scale the initial noise by the standard deviation required by the scheduler
638
+ latents = (torch.from_numpy(latents) * self.scheduler.init_noise_sigma).numpy()
639
+ return latents, None, None
640
+ else:
641
+ init_latents = self.vae_encoder(sample=image)[0]
642
+ init_latents = 0.18215 * init_latents
643
+ init_latents = np.concatenate([init_latents] * batch_size, axis=0)
644
+ init_latents_orig = init_latents
645
+ shape = init_latents.shape
646
+
647
+ # add noise to latents using the timesteps
648
+ noise = torch.randn(shape, generator=generator, device="cpu").numpy().astype(dtype)
649
+ latents = self.scheduler.add_noise(
650
+ torch.from_numpy(init_latents), torch.from_numpy(noise), timestep
651
+ ).numpy()
652
+ return latents, init_latents_orig, noise
653
+
654
+ @torch.no_grad()
655
+ def __call__(
656
+ self,
657
+ prompt: Union[str, List[str]],
658
+ negative_prompt: Optional[Union[str, List[str]]] = None,
659
+ image: Union[np.ndarray, PIL.Image.Image] = None,
660
+ mask_image: Union[np.ndarray, PIL.Image.Image] = None,
661
+ height: int = 512,
662
+ width: int = 512,
663
+ num_inference_steps: int = 50,
664
+ guidance_scale: float = 7.5,
665
+ strength: float = 0.8,
666
+ num_images_per_prompt: Optional[int] = 1,
667
+ eta: float = 0.0,
668
+ generator: Optional[torch.Generator] = None,
669
+ latents: Optional[np.ndarray] = None,
670
+ max_embeddings_multiples: Optional[int] = 3,
671
+ output_type: Optional[str] = "pil",
672
+ return_dict: bool = True,
673
+ callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
674
+ is_cancelled_callback: Optional[Callable[[], bool]] = None,
675
+ callback_steps: int = 1,
676
+ **kwargs,
677
+ ):
678
+ r"""
679
+ Function invoked when calling the pipeline for generation.
680
+
681
+ Args:
682
+ prompt (`str` or `List[str]`):
683
+ The prompt or prompts to guide the image generation.
684
+ negative_prompt (`str` or `List[str]`, *optional*):
685
+ The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
686
+ if `guidance_scale` is less than `1`).
687
+ image (`np.ndarray` or `PIL.Image.Image`):
688
+ `Image`, or tensor representing an image batch, that will be used as the starting point for the
689
+ process.
690
+ mask_image (`np.ndarray` or `PIL.Image.Image`):
691
+ `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be
692
+ replaced by noise and therefore repainted, while black pixels will be preserved. If `mask_image` is a
693
+ PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should
694
+ contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`.
695
+ height (`int`, *optional*, defaults to 512):
696
+ The height in pixels of the generated image.
697
+ width (`int`, *optional*, defaults to 512):
698
+ The width in pixels of the generated image.
699
+ num_inference_steps (`int`, *optional*, defaults to 50):
700
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
701
+ expense of slower inference.
702
+ guidance_scale (`float`, *optional*, defaults to 7.5):
703
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
704
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
705
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
706
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
707
+ usually at the expense of lower image quality.
708
+ strength (`float`, *optional*, defaults to 0.8):
709
+ Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1.
710
+ `image` will be used as a starting point, adding more noise to it the larger the `strength`. The
711
+ number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added
712
+ noise will be maximum and the denoising process will run for the full number of iterations specified in
713
+ `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.
714
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
715
+ The number of images to generate per prompt.
716
+ eta (`float`, *optional*, defaults to 0.0):
717
+ Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
718
+ [`schedulers.DDIMScheduler`], will be ignored for others.
719
+ generator (`torch.Generator`, *optional*):
720
+ A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
721
+ deterministic.
722
+ latents (`np.ndarray`, *optional*):
723
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
724
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
725
+ tensor will ge generated by sampling using the supplied random `generator`.
726
+ max_embeddings_multiples (`int`, *optional*, defaults to `3`):
727
+ The max multiple length of prompt embeddings compared to the max output length of text encoder.
728
+ output_type (`str`, *optional*, defaults to `"pil"`):
729
+ The output format of the generate image. Choose between
730
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
731
+ return_dict (`bool`, *optional*, defaults to `True`):
732
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
733
+ plain tuple.
734
+ callback (`Callable`, *optional*):
735
+ A function that will be called every `callback_steps` steps during inference. The function will be
736
+ called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.
737
+ is_cancelled_callback (`Callable`, *optional*):
738
+ A function that will be called every `callback_steps` steps during inference. If the function returns
739
+ `True`, the inference will be cancelled.
740
+ callback_steps (`int`, *optional*, defaults to 1):
741
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
742
+ called at every step.
743
+
744
+ Returns:
745
+ `None` if cancelled by `is_cancelled_callback`,
746
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
747
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
748
+ When returning a tuple, the first element is a list with the generated images, and the second element is a
749
+ list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
750
+ (nsfw) content, according to the `safety_checker`.
751
+ """
752
+ # 0. Default height and width to unet
753
+ height = height or self.unet.config.sample_size * self.vae_scale_factor
754
+ width = width or self.unet.config.sample_size * self.vae_scale_factor
755
+
756
+ # 1. Check inputs. Raise error if not correct
757
+ self.check_inputs(prompt, height, width, strength, callback_steps)
758
+
759
+ # 2. Define call parameters
760
+ batch_size = 1 if isinstance(prompt, str) else len(prompt)
761
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
762
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
763
+ # corresponds to doing no classifier free guidance.
764
+ do_classifier_free_guidance = guidance_scale > 1.0
765
+
766
+ # 3. Encode input prompt
767
+ text_embeddings = self.encode_prompt(
768
+ prompt,
769
+ num_images_per_prompt,
770
+ do_classifier_free_guidance,
771
+ negative_prompt,
772
+ max_embeddings_multiples,
773
+ )
774
+ dtype = text_embeddings.dtype
775
+
776
+ # 4. Preprocess image and mask
777
+ if isinstance(image, PIL.Image.Image):
778
+ image = preprocess_image(image)
779
+ if image is not None:
780
+ image = image.astype(dtype)
781
+ if isinstance(mask_image, PIL.Image.Image):
782
+ mask_image = preprocess_mask(mask_image, self.vae_scale_factor)
783
+ if mask_image is not None:
784
+ mask = mask_image.astype(dtype)
785
+ mask = np.concatenate([mask] * batch_size * num_images_per_prompt)
786
+ else:
787
+ mask = None
788
+
789
+ # 5. set timesteps
790
+ self.scheduler.set_timesteps(num_inference_steps)
791
+ timestep_dtype = next(
792
+ (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)"
793
+ )
794
+ timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype]
795
+ timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, image is None)
796
+ latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
797
+
798
+ # 6. Prepare latent variables
799
+ latents, init_latents_orig, noise = self.prepare_latents(
800
+ image,
801
+ latent_timestep,
802
+ batch_size * num_images_per_prompt,
803
+ height,
804
+ width,
805
+ dtype,
806
+ generator,
807
+ latents,
808
+ )
809
+
810
+ # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
811
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
812
+
813
+ # 8. Denoising loop
814
+ for i, t in enumerate(self.progress_bar(timesteps)):
815
+ # expand the latents if we are doing classifier free guidance
816
+ latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents
817
+ latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t)
818
+ latent_model_input = latent_model_input.numpy()
819
+
820
+ # predict the noise residual
821
+ noise_pred = self.unet(
822
+ sample=latent_model_input,
823
+ timestep=np.array([t], dtype=timestep_dtype),
824
+ encoder_hidden_states=text_embeddings,
825
+ )
826
+ noise_pred = noise_pred[0]
827
+
828
+ # perform guidance
829
+ if do_classifier_free_guidance:
830
+ noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
831
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
832
+
833
+ # compute the previous noisy sample x_t -> x_t-1
834
+ scheduler_output = self.scheduler.step(
835
+ torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs
836
+ )
837
+ latents = scheduler_output.prev_sample.numpy()
838
+
839
+ if mask is not None:
840
+ # masking
841
+ init_latents_proper = self.scheduler.add_noise(
842
+ torch.from_numpy(init_latents_orig),
843
+ torch.from_numpy(noise),
844
+ t,
845
+ ).numpy()
846
+ latents = (init_latents_proper * mask) + (latents * (1 - mask))
847
+
848
+ # call the callback, if provided
849
+ if i % callback_steps == 0:
850
+ if callback is not None:
851
+ step_idx = i // getattr(self.scheduler, "order", 1)
852
+ callback(step_idx, t, latents)
853
+ if is_cancelled_callback is not None and is_cancelled_callback():
854
+ return None
855
+
856
+ # 9. Post-processing
857
+ image = self.decode_latents(latents)
858
+
859
+ # 10. Run safety checker
860
+ image, has_nsfw_concept = self.run_safety_checker(image)
861
+
862
+ # 11. Convert to PIL
863
+ if output_type == "pil":
864
+ image = self.numpy_to_pil(image)
865
+
866
+ if not return_dict:
867
+ return image, has_nsfw_concept
868
+
869
+ return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
870
+
871
+ def text2img(
872
+ self,
873
+ prompt: Union[str, List[str]],
874
+ negative_prompt: Optional[Union[str, List[str]]] = None,
875
+ height: int = 512,
876
+ width: int = 512,
877
+ num_inference_steps: int = 50,
878
+ guidance_scale: float = 7.5,
879
+ num_images_per_prompt: Optional[int] = 1,
880
+ eta: float = 0.0,
881
+ generator: Optional[torch.Generator] = None,
882
+ latents: Optional[np.ndarray] = None,
883
+ max_embeddings_multiples: Optional[int] = 3,
884
+ output_type: Optional[str] = "pil",
885
+ return_dict: bool = True,
886
+ callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
887
+ callback_steps: int = 1,
888
+ **kwargs,
889
+ ):
890
+ r"""
891
+ Function for text-to-image generation.
892
+ Args:
893
+ prompt (`str` or `List[str]`):
894
+ The prompt or prompts to guide the image generation.
895
+ negative_prompt (`str` or `List[str]`, *optional*):
896
+ The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
897
+ if `guidance_scale` is less than `1`).
898
+ height (`int`, *optional*, defaults to 512):
899
+ The height in pixels of the generated image.
900
+ width (`int`, *optional*, defaults to 512):
901
+ The width in pixels of the generated image.
902
+ num_inference_steps (`int`, *optional*, defaults to 50):
903
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
904
+ expense of slower inference.
905
+ guidance_scale (`float`, *optional*, defaults to 7.5):
906
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
907
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
908
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
909
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
910
+ usually at the expense of lower image quality.
911
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
912
+ The number of images to generate per prompt.
913
+ eta (`float`, *optional*, defaults to 0.0):
914
+ Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
915
+ [`schedulers.DDIMScheduler`], will be ignored for others.
916
+ generator (`torch.Generator`, *optional*):
917
+ A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
918
+ deterministic.
919
+ latents (`np.ndarray`, *optional*):
920
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
921
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
922
+ tensor will ge generated by sampling using the supplied random `generator`.
923
+ max_embeddings_multiples (`int`, *optional*, defaults to `3`):
924
+ The max multiple length of prompt embeddings compared to the max output length of text encoder.
925
+ output_type (`str`, *optional*, defaults to `"pil"`):
926
+ The output format of the generate image. Choose between
927
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
928
+ return_dict (`bool`, *optional*, defaults to `True`):
929
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
930
+ plain tuple.
931
+ callback (`Callable`, *optional*):
932
+ A function that will be called every `callback_steps` steps during inference. The function will be
933
+ called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.
934
+ callback_steps (`int`, *optional*, defaults to 1):
935
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
936
+ called at every step.
937
+ Returns:
938
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
939
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
940
+ When returning a tuple, the first element is a list with the generated images, and the second element is a
941
+ list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
942
+ (nsfw) content, according to the `safety_checker`.
943
+ """
944
+ return self.__call__(
945
+ prompt=prompt,
946
+ negative_prompt=negative_prompt,
947
+ height=height,
948
+ width=width,
949
+ num_inference_steps=num_inference_steps,
950
+ guidance_scale=guidance_scale,
951
+ num_images_per_prompt=num_images_per_prompt,
952
+ eta=eta,
953
+ generator=generator,
954
+ latents=latents,
955
+ max_embeddings_multiples=max_embeddings_multiples,
956
+ output_type=output_type,
957
+ return_dict=return_dict,
958
+ callback=callback,
959
+ callback_steps=callback_steps,
960
+ **kwargs,
961
+ )
962
+
963
+ def img2img(
964
+ self,
965
+ image: Union[np.ndarray, PIL.Image.Image],
966
+ prompt: Union[str, List[str]],
967
+ negative_prompt: Optional[Union[str, List[str]]] = None,
968
+ strength: float = 0.8,
969
+ num_inference_steps: Optional[int] = 50,
970
+ guidance_scale: Optional[float] = 7.5,
971
+ num_images_per_prompt: Optional[int] = 1,
972
+ eta: Optional[float] = 0.0,
973
+ generator: Optional[torch.Generator] = None,
974
+ max_embeddings_multiples: Optional[int] = 3,
975
+ output_type: Optional[str] = "pil",
976
+ return_dict: bool = True,
977
+ callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
978
+ callback_steps: int = 1,
979
+ **kwargs,
980
+ ):
981
+ r"""
982
+ Function for image-to-image generation.
983
+ Args:
984
+ image (`np.ndarray` or `PIL.Image.Image`):
985
+ `Image`, or ndarray representing an image batch, that will be used as the starting point for the
986
+ process.
987
+ prompt (`str` or `List[str]`):
988
+ The prompt or prompts to guide the image generation.
989
+ negative_prompt (`str` or `List[str]`, *optional*):
990
+ The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
991
+ if `guidance_scale` is less than `1`).
992
+ strength (`float`, *optional*, defaults to 0.8):
993
+ Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1.
994
+ `image` will be used as a starting point, adding more noise to it the larger the `strength`. The
995
+ number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added
996
+ noise will be maximum and the denoising process will run for the full number of iterations specified in
997
+ `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.
998
+ num_inference_steps (`int`, *optional*, defaults to 50):
999
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
1000
+ expense of slower inference. This parameter will be modulated by `strength`.
1001
+ guidance_scale (`float`, *optional*, defaults to 7.5):
1002
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
1003
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
1004
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1005
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
1006
+ usually at the expense of lower image quality.
1007
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
1008
+ The number of images to generate per prompt.
1009
+ eta (`float`, *optional*, defaults to 0.0):
1010
+ Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
1011
+ [`schedulers.DDIMScheduler`], will be ignored for others.
1012
+ generator (`torch.Generator`, *optional*):
1013
+ A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
1014
+ deterministic.
1015
+ max_embeddings_multiples (`int`, *optional*, defaults to `3`):
1016
+ The max multiple length of prompt embeddings compared to the max output length of text encoder.
1017
+ output_type (`str`, *optional*, defaults to `"pil"`):
1018
+ The output format of the generate image. Choose between
1019
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
1020
+ return_dict (`bool`, *optional*, defaults to `True`):
1021
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
1022
+ plain tuple.
1023
+ callback (`Callable`, *optional*):
1024
+ A function that will be called every `callback_steps` steps during inference. The function will be
1025
+ called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.
1026
+ callback_steps (`int`, *optional*, defaults to 1):
1027
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
1028
+ called at every step.
1029
+ Returns:
1030
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
1031
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
1032
+ When returning a tuple, the first element is a list with the generated images, and the second element is a
1033
+ list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
1034
+ (nsfw) content, according to the `safety_checker`.
1035
+ """
1036
+ return self.__call__(
1037
+ prompt=prompt,
1038
+ negative_prompt=negative_prompt,
1039
+ image=image,
1040
+ num_inference_steps=num_inference_steps,
1041
+ guidance_scale=guidance_scale,
1042
+ strength=strength,
1043
+ num_images_per_prompt=num_images_per_prompt,
1044
+ eta=eta,
1045
+ generator=generator,
1046
+ max_embeddings_multiples=max_embeddings_multiples,
1047
+ output_type=output_type,
1048
+ return_dict=return_dict,
1049
+ callback=callback,
1050
+ callback_steps=callback_steps,
1051
+ **kwargs,
1052
+ )
1053
+
1054
+ def inpaint(
1055
+ self,
1056
+ image: Union[np.ndarray, PIL.Image.Image],
1057
+ mask_image: Union[np.ndarray, PIL.Image.Image],
1058
+ prompt: Union[str, List[str]],
1059
+ negative_prompt: Optional[Union[str, List[str]]] = None,
1060
+ strength: float = 0.8,
1061
+ num_inference_steps: Optional[int] = 50,
1062
+ guidance_scale: Optional[float] = 7.5,
1063
+ num_images_per_prompt: Optional[int] = 1,
1064
+ eta: Optional[float] = 0.0,
1065
+ generator: Optional[torch.Generator] = None,
1066
+ max_embeddings_multiples: Optional[int] = 3,
1067
+ output_type: Optional[str] = "pil",
1068
+ return_dict: bool = True,
1069
+ callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
1070
+ callback_steps: int = 1,
1071
+ **kwargs,
1072
+ ):
1073
+ r"""
1074
+ Function for inpaint.
1075
+ Args:
1076
+ image (`np.ndarray` or `PIL.Image.Image`):
1077
+ `Image`, or tensor representing an image batch, that will be used as the starting point for the
1078
+ process. This is the image whose masked region will be inpainted.
1079
+ mask_image (`np.ndarray` or `PIL.Image.Image`):
1080
+ `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be
1081
+ replaced by noise and therefore repainted, while black pixels will be preserved. If `mask_image` is a
1082
+ PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should
1083
+ contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`.
1084
+ prompt (`str` or `List[str]`):
1085
+ The prompt or prompts to guide the image generation.
1086
+ negative_prompt (`str` or `List[str]`, *optional*):
1087
+ The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
1088
+ if `guidance_scale` is less than `1`).
1089
+ strength (`float`, *optional*, defaults to 0.8):
1090
+ Conceptually, indicates how much to inpaint the masked area. Must be between 0 and 1. When `strength`
1091
+ is 1, the denoising process will be run on the masked area for the full number of iterations specified
1092
+ in `num_inference_steps`. `image` will be used as a reference for the masked area, adding more
1093
+ noise to that region the larger the `strength`. If `strength` is 0, no inpainting will occur.
1094
+ num_inference_steps (`int`, *optional*, defaults to 50):
1095
+ The reference number of denoising steps. More denoising steps usually lead to a higher quality image at
1096
+ the expense of slower inference. This parameter will be modulated by `strength`, as explained above.
1097
+ guidance_scale (`float`, *optional*, defaults to 7.5):
1098
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
1099
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
1100
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1101
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
1102
+ usually at the expense of lower image quality.
1103
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
1104
+ The number of images to generate per prompt.
1105
+ eta (`float`, *optional*, defaults to 0.0):
1106
+ Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
1107
+ [`schedulers.DDIMScheduler`], will be ignored for others.
1108
+ generator (`torch.Generator`, *optional*):
1109
+ A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
1110
+ deterministic.
1111
+ max_embeddings_multiples (`int`, *optional*, defaults to `3`):
1112
+ The max multiple length of prompt embeddings compared to the max output length of text encoder.
1113
+ output_type (`str`, *optional*, defaults to `"pil"`):
1114
+ The output format of the generate image. Choose between
1115
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
1116
+ return_dict (`bool`, *optional*, defaults to `True`):
1117
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
1118
+ plain tuple.
1119
+ callback (`Callable`, *optional*):
1120
+ A function that will be called every `callback_steps` steps during inference. The function will be
1121
+ called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.
1122
+ callback_steps (`int`, *optional*, defaults to 1):
1123
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
1124
+ called at every step.
1125
+ Returns:
1126
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
1127
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
1128
+ When returning a tuple, the first element is a list with the generated images, and the second element is a
1129
+ list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
1130
+ (nsfw) content, according to the `safety_checker`.
1131
+ """
1132
+ return self.__call__(
1133
+ prompt=prompt,
1134
+ negative_prompt=negative_prompt,
1135
+ image=image,
1136
+ mask_image=mask_image,
1137
+ num_inference_steps=num_inference_steps,
1138
+ guidance_scale=guidance_scale,
1139
+ strength=strength,
1140
+ num_images_per_prompt=num_images_per_prompt,
1141
+ eta=eta,
1142
+ generator=generator,
1143
+ max_embeddings_multiples=max_embeddings_multiples,
1144
+ output_type=output_type,
1145
+ return_dict=return_dict,
1146
+ callback=callback,
1147
+ callback_steps=callback_steps,
1148
+ **kwargs,
1149
+ )