multimodalart HF staff commited on
Commit
ca51647
1 Parent(s): 8ddce9c

Delete ip_adapter

Browse files
ip_adapter/__init__.py DELETED
@@ -1,9 +0,0 @@
1
- from .ip_adapter import IPAdapter, IPAdapterPlus, IPAdapterPlusXL, IPAdapterXL, IPAdapterFull
2
-
3
- __all__ = [
4
- "IPAdapter",
5
- "IPAdapterPlus",
6
- "IPAdapterPlusXL",
7
- "IPAdapterXL",
8
- "IPAdapterFull",
9
- ]
 
 
 
 
 
 
 
 
 
 
ip_adapter/attention_processor.py DELETED
@@ -1,554 +0,0 @@
1
- # modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
2
- import torch
3
- import torch.nn as nn
4
- import torch.nn.functional as F
5
-
6
-
7
- class AttnProcessor(nn.Module):
8
- r"""
9
- Default processor for performing attention-related computations.
10
- """
11
-
12
- def __init__(
13
- self,
14
- hidden_size=None,
15
- cross_attention_dim=None,
16
- ):
17
- super().__init__()
18
-
19
- def __call__(
20
- self,
21
- attn,
22
- hidden_states,
23
- encoder_hidden_states=None,
24
- attention_mask=None,
25
- temb=None,
26
- ):
27
- residual = hidden_states
28
-
29
- if attn.spatial_norm is not None:
30
- hidden_states = attn.spatial_norm(hidden_states, temb)
31
-
32
- input_ndim = hidden_states.ndim
33
-
34
- if input_ndim == 4:
35
- batch_size, channel, height, width = hidden_states.shape
36
- hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
37
-
38
- batch_size, sequence_length, _ = (
39
- hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
40
- )
41
- attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
42
-
43
- if attn.group_norm is not None:
44
- hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
45
-
46
- query = attn.to_q(hidden_states)
47
-
48
- if encoder_hidden_states is None:
49
- encoder_hidden_states = hidden_states
50
- elif attn.norm_cross:
51
- encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
52
-
53
- key = attn.to_k(encoder_hidden_states)
54
- value = attn.to_v(encoder_hidden_states)
55
-
56
- query = attn.head_to_batch_dim(query)
57
- key = attn.head_to_batch_dim(key)
58
- value = attn.head_to_batch_dim(value)
59
-
60
- attention_probs = attn.get_attention_scores(query, key, attention_mask)
61
- hidden_states = torch.bmm(attention_probs, value)
62
- hidden_states = attn.batch_to_head_dim(hidden_states)
63
-
64
- # linear proj
65
- hidden_states = attn.to_out[0](hidden_states)
66
- # dropout
67
- hidden_states = attn.to_out[1](hidden_states)
68
-
69
- if input_ndim == 4:
70
- hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
71
-
72
- if attn.residual_connection:
73
- hidden_states = hidden_states + residual
74
-
75
- hidden_states = hidden_states / attn.rescale_output_factor
76
-
77
- return hidden_states
78
-
79
-
80
- class IPAttnProcessor(nn.Module):
81
- r"""
82
- Attention processor for IP-Adapater.
83
- Args:
84
- hidden_size (`int`):
85
- The hidden size of the attention layer.
86
- cross_attention_dim (`int`):
87
- The number of channels in the `encoder_hidden_states`.
88
- scale (`float`, defaults to 1.0):
89
- the weight scale of image prompt.
90
- num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
91
- The context length of the image features.
92
- """
93
-
94
- def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4):
95
- super().__init__()
96
-
97
- self.hidden_size = hidden_size
98
- self.cross_attention_dim = cross_attention_dim
99
- self.scale = scale
100
- self.num_tokens = num_tokens
101
-
102
- self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
103
- self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
104
-
105
- def __call__(
106
- self,
107
- attn,
108
- hidden_states,
109
- encoder_hidden_states=None,
110
- attention_mask=None,
111
- temb=None,
112
- ):
113
- residual = hidden_states
114
-
115
- if attn.spatial_norm is not None:
116
- hidden_states = attn.spatial_norm(hidden_states, temb)
117
-
118
- input_ndim = hidden_states.ndim
119
-
120
- if input_ndim == 4:
121
- batch_size, channel, height, width = hidden_states.shape
122
- hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
123
-
124
- batch_size, sequence_length, _ = (
125
- hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
126
- )
127
- attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
128
-
129
- if attn.group_norm is not None:
130
- hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
131
-
132
- query = attn.to_q(hidden_states)
133
-
134
- if encoder_hidden_states is None:
135
- encoder_hidden_states = hidden_states
136
- else:
137
- # get encoder_hidden_states, ip_hidden_states
138
- end_pos = encoder_hidden_states.shape[1] - self.num_tokens
139
- encoder_hidden_states, ip_hidden_states = (
140
- encoder_hidden_states[:, :end_pos, :],
141
- encoder_hidden_states[:, end_pos:, :],
142
- )
143
- if attn.norm_cross:
144
- encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
145
-
146
- key = attn.to_k(encoder_hidden_states)
147
- value = attn.to_v(encoder_hidden_states)
148
-
149
- query = attn.head_to_batch_dim(query)
150
- key = attn.head_to_batch_dim(key)
151
- value = attn.head_to_batch_dim(value)
152
-
153
- attention_probs = attn.get_attention_scores(query, key, attention_mask)
154
- hidden_states = torch.bmm(attention_probs, value)
155
- hidden_states = attn.batch_to_head_dim(hidden_states)
156
-
157
- # for ip-adapter
158
- ip_key = self.to_k_ip(ip_hidden_states)
159
- ip_value = self.to_v_ip(ip_hidden_states)
160
-
161
- ip_key = attn.head_to_batch_dim(ip_key)
162
- ip_value = attn.head_to_batch_dim(ip_value)
163
-
164
- ip_attention_probs = attn.get_attention_scores(query, ip_key, None)
165
- ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
166
- ip_hidden_states = attn.batch_to_head_dim(ip_hidden_states)
167
-
168
- hidden_states = hidden_states + self.scale * ip_hidden_states
169
-
170
- # linear proj
171
- hidden_states = attn.to_out[0](hidden_states)
172
- # dropout
173
- hidden_states = attn.to_out[1](hidden_states)
174
-
175
- if input_ndim == 4:
176
- hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
177
-
178
- if attn.residual_connection:
179
- hidden_states = hidden_states + residual
180
-
181
- hidden_states = hidden_states / attn.rescale_output_factor
182
-
183
- return hidden_states
184
-
185
-
186
- class AttnProcessor2_0(torch.nn.Module):
187
- r"""
188
- Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
189
- """
190
-
191
- def __init__(
192
- self,
193
- hidden_size=None,
194
- cross_attention_dim=None,
195
- ):
196
- super().__init__()
197
- if not hasattr(F, "scaled_dot_product_attention"):
198
- raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
199
-
200
- def __call__(
201
- self,
202
- attn,
203
- hidden_states,
204
- encoder_hidden_states=None,
205
- attention_mask=None,
206
- temb=None,
207
- ):
208
- residual = hidden_states
209
-
210
- if attn.spatial_norm is not None:
211
- hidden_states = attn.spatial_norm(hidden_states, temb)
212
-
213
- input_ndim = hidden_states.ndim
214
-
215
- if input_ndim == 4:
216
- batch_size, channel, height, width = hidden_states.shape
217
- hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
218
-
219
- batch_size, sequence_length, _ = (
220
- hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
221
- )
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 = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
231
-
232
- query = attn.to_q(hidden_states)
233
-
234
- if encoder_hidden_states is None:
235
- encoder_hidden_states = hidden_states
236
- elif attn.norm_cross:
237
- encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
238
-
239
- key = attn.to_k(encoder_hidden_states)
240
- value = attn.to_v(encoder_hidden_states)
241
-
242
- inner_dim = key.shape[-1]
243
- head_dim = inner_dim // attn.heads
244
-
245
- query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
246
-
247
- key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
248
- value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
249
-
250
- # the output of sdp = (batch, num_heads, seq_len, head_dim)
251
- # TODO: add support for attn.scale when we move to Torch 2.1
252
- hidden_states = F.scaled_dot_product_attention(
253
- query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
254
- )
255
-
256
- hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
257
- hidden_states = hidden_states.to(query.dtype)
258
-
259
- # linear proj
260
- hidden_states = attn.to_out[0](hidden_states)
261
- # dropout
262
- hidden_states = attn.to_out[1](hidden_states)
263
-
264
- if input_ndim == 4:
265
- hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
266
-
267
- if attn.residual_connection:
268
- hidden_states = hidden_states + residual
269
-
270
- hidden_states = hidden_states / attn.rescale_output_factor
271
-
272
- return hidden_states
273
-
274
-
275
- class IPAttnProcessor2_0(torch.nn.Module):
276
- r"""
277
- Attention processor for IP-Adapater for PyTorch 2.0.
278
- Args:
279
- hidden_size (`int`):
280
- The hidden size of the attention layer.
281
- cross_attention_dim (`int`):
282
- The number of channels in the `encoder_hidden_states`.
283
- scale (`float`, defaults to 1.0):
284
- the weight scale of image prompt.
285
- num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
286
- The context length of the image features.
287
- """
288
-
289
- def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4):
290
- super().__init__()
291
-
292
- if not hasattr(F, "scaled_dot_product_attention"):
293
- raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
294
-
295
- self.hidden_size = hidden_size
296
- self.cross_attention_dim = cross_attention_dim
297
- self.scale = scale
298
- self.num_tokens = num_tokens
299
-
300
- self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
301
- self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
302
-
303
- def __call__(
304
- self,
305
- attn,
306
- hidden_states,
307
- encoder_hidden_states=None,
308
- attention_mask=None,
309
- temb=None,
310
- ):
311
- residual = hidden_states
312
-
313
- if attn.spatial_norm is not None:
314
- hidden_states = attn.spatial_norm(hidden_states, temb)
315
-
316
- input_ndim = hidden_states.ndim
317
-
318
- if input_ndim == 4:
319
- batch_size, channel, height, width = hidden_states.shape
320
- hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
321
-
322
- batch_size, sequence_length, _ = (
323
- hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
324
- )
325
-
326
- if attention_mask is not None:
327
- attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
328
- # scaled_dot_product_attention expects attention_mask shape to be
329
- # (batch, heads, source_length, target_length)
330
- attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
331
-
332
- if attn.group_norm is not None:
333
- hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
334
-
335
- query = attn.to_q(hidden_states)
336
-
337
- if encoder_hidden_states is None:
338
- encoder_hidden_states = hidden_states
339
- else:
340
- # get encoder_hidden_states, ip_hidden_states
341
- end_pos = encoder_hidden_states.shape[1] - self.num_tokens
342
- encoder_hidden_states, ip_hidden_states = (
343
- encoder_hidden_states[:, :end_pos, :],
344
- encoder_hidden_states[:, end_pos:, :],
345
- )
346
- if attn.norm_cross:
347
- encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
348
-
349
- key = attn.to_k(encoder_hidden_states)
350
- value = attn.to_v(encoder_hidden_states)
351
-
352
- inner_dim = key.shape[-1]
353
- head_dim = inner_dim // attn.heads
354
-
355
- query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
356
-
357
- key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
358
- value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
359
-
360
- # the output of sdp = (batch, num_heads, seq_len, head_dim)
361
- # TODO: add support for attn.scale when we move to Torch 2.1
362
- hidden_states = F.scaled_dot_product_attention(
363
- query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
364
- )
365
-
366
- hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
367
- hidden_states = hidden_states.to(query.dtype)
368
-
369
- # for ip-adapter
370
- ip_key = self.to_k_ip(ip_hidden_states)
371
- ip_value = self.to_v_ip(ip_hidden_states)
372
-
373
- ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
374
- ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
375
-
376
- # the output of sdp = (batch, num_heads, seq_len, head_dim)
377
- # TODO: add support for attn.scale when we move to Torch 2.1
378
- ip_hidden_states = F.scaled_dot_product_attention(
379
- query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
380
- )
381
-
382
- ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
383
- ip_hidden_states = ip_hidden_states.to(query.dtype)
384
-
385
- hidden_states = hidden_states + self.scale * ip_hidden_states
386
-
387
- # linear proj
388
- hidden_states = attn.to_out[0](hidden_states)
389
- # dropout
390
- hidden_states = attn.to_out[1](hidden_states)
391
-
392
- if input_ndim == 4:
393
- hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
394
-
395
- if attn.residual_connection:
396
- hidden_states = hidden_states + residual
397
-
398
- hidden_states = hidden_states / attn.rescale_output_factor
399
-
400
- return hidden_states
401
-
402
-
403
- ## for controlnet
404
- class CNAttnProcessor:
405
- r"""
406
- Default processor for performing attention-related computations.
407
- """
408
-
409
- def __init__(self, num_tokens=4):
410
- self.num_tokens = num_tokens
411
-
412
- def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None):
413
- residual = hidden_states
414
-
415
- if attn.spatial_norm is not None:
416
- hidden_states = attn.spatial_norm(hidden_states, temb)
417
-
418
- input_ndim = hidden_states.ndim
419
-
420
- if input_ndim == 4:
421
- batch_size, channel, height, width = hidden_states.shape
422
- hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
423
-
424
- batch_size, sequence_length, _ = (
425
- hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
426
- )
427
- attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
428
-
429
- if attn.group_norm is not None:
430
- hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
431
-
432
- query = attn.to_q(hidden_states)
433
-
434
- if encoder_hidden_states is None:
435
- encoder_hidden_states = hidden_states
436
- else:
437
- end_pos = encoder_hidden_states.shape[1] - self.num_tokens
438
- encoder_hidden_states = encoder_hidden_states[:, :end_pos] # only use text
439
- if attn.norm_cross:
440
- encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
441
-
442
- key = attn.to_k(encoder_hidden_states)
443
- value = attn.to_v(encoder_hidden_states)
444
-
445
- query = attn.head_to_batch_dim(query)
446
- key = attn.head_to_batch_dim(key)
447
- value = attn.head_to_batch_dim(value)
448
-
449
- attention_probs = attn.get_attention_scores(query, key, attention_mask)
450
- hidden_states = torch.bmm(attention_probs, value)
451
- hidden_states = attn.batch_to_head_dim(hidden_states)
452
-
453
- # linear proj
454
- hidden_states = attn.to_out[0](hidden_states)
455
- # dropout
456
- hidden_states = attn.to_out[1](hidden_states)
457
-
458
- if input_ndim == 4:
459
- hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
460
-
461
- if attn.residual_connection:
462
- hidden_states = hidden_states + residual
463
-
464
- hidden_states = hidden_states / attn.rescale_output_factor
465
-
466
- return hidden_states
467
-
468
-
469
- class CNAttnProcessor2_0:
470
- r"""
471
- Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
472
- """
473
-
474
- def __init__(self, num_tokens=4):
475
- if not hasattr(F, "scaled_dot_product_attention"):
476
- raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
477
- self.num_tokens = num_tokens
478
-
479
- def __call__(
480
- self,
481
- attn,
482
- hidden_states,
483
- encoder_hidden_states=None,
484
- attention_mask=None,
485
- temb=None,
486
- ):
487
- residual = hidden_states
488
-
489
- if attn.spatial_norm is not None:
490
- hidden_states = attn.spatial_norm(hidden_states, temb)
491
-
492
- input_ndim = hidden_states.ndim
493
-
494
- if input_ndim == 4:
495
- batch_size, channel, height, width = hidden_states.shape
496
- hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
497
-
498
- batch_size, sequence_length, _ = (
499
- hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
500
- )
501
-
502
- if attention_mask is not None:
503
- attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
504
- # scaled_dot_product_attention expects attention_mask shape to be
505
- # (batch, heads, source_length, target_length)
506
- attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
507
-
508
- if attn.group_norm is not None:
509
- hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
510
-
511
- query = attn.to_q(hidden_states)
512
-
513
- if encoder_hidden_states is None:
514
- encoder_hidden_states = hidden_states
515
- else:
516
- end_pos = encoder_hidden_states.shape[1] - self.num_tokens
517
- encoder_hidden_states = encoder_hidden_states[:, :end_pos] # only use text
518
- if attn.norm_cross:
519
- encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
520
-
521
- key = attn.to_k(encoder_hidden_states)
522
- value = attn.to_v(encoder_hidden_states)
523
-
524
- inner_dim = key.shape[-1]
525
- head_dim = inner_dim // attn.heads
526
-
527
- query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
528
-
529
- key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
530
- value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
531
-
532
- # the output of sdp = (batch, num_heads, seq_len, head_dim)
533
- # TODO: add support for attn.scale when we move to Torch 2.1
534
- hidden_states = F.scaled_dot_product_attention(
535
- query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
536
- )
537
-
538
- hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
539
- hidden_states = hidden_states.to(query.dtype)
540
-
541
- # linear proj
542
- hidden_states = attn.to_out[0](hidden_states)
543
- # dropout
544
- hidden_states = attn.to_out[1](hidden_states)
545
-
546
- if input_ndim == 4:
547
- hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
548
-
549
- if attn.residual_connection:
550
- hidden_states = hidden_states + residual
551
-
552
- hidden_states = hidden_states / attn.rescale_output_factor
553
-
554
- return hidden_states
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ip_adapter/attention_processor_faceid.py DELETED
@@ -1,204 +0,0 @@
1
- # modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
2
- import torch
3
- import torch.nn as nn
4
- import torch.nn.functional as F
5
-
6
- from diffusers.models.lora import LoRALinearLayer
7
-
8
-
9
- class LoRAAttnProcessor(nn.Module):
10
- r"""
11
- Default processor for performing attention-related computations.
12
- """
13
-
14
- def __init__(
15
- self,
16
- hidden_size=None,
17
- cross_attention_dim=None,
18
- rank=4,
19
- network_alpha=None,
20
- lora_scale=1.0,
21
- ):
22
- super().__init__()
23
-
24
- self.rank = rank
25
- self.lora_scale = lora_scale
26
-
27
- self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
28
- self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
29
- self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
30
- self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
31
-
32
- def __call__(
33
- self,
34
- attn,
35
- hidden_states,
36
- encoder_hidden_states=None,
37
- attention_mask=None,
38
- temb=None,
39
- ):
40
- residual = hidden_states
41
-
42
- if attn.spatial_norm is not None:
43
- hidden_states = attn.spatial_norm(hidden_states, temb)
44
-
45
- input_ndim = hidden_states.ndim
46
-
47
- if input_ndim == 4:
48
- batch_size, channel, height, width = hidden_states.shape
49
- hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
50
-
51
- batch_size, sequence_length, _ = (
52
- hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
53
- )
54
- attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
55
-
56
- if attn.group_norm is not None:
57
- hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
58
-
59
- query = attn.to_q(hidden_states) + self.lora_scale * self.to_q_lora(hidden_states)
60
-
61
- if encoder_hidden_states is None:
62
- encoder_hidden_states = hidden_states
63
- elif attn.norm_cross:
64
- encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
65
-
66
- key = attn.to_k(encoder_hidden_states) + self.lora_scale * self.to_k_lora(encoder_hidden_states)
67
- value = attn.to_v(encoder_hidden_states) + self.lora_scale * self.to_v_lora(encoder_hidden_states)
68
-
69
- query = attn.head_to_batch_dim(query)
70
- key = attn.head_to_batch_dim(key)
71
- value = attn.head_to_batch_dim(value)
72
-
73
- attention_probs = attn.get_attention_scores(query, key, attention_mask)
74
- hidden_states = torch.bmm(attention_probs, value)
75
- hidden_states = attn.batch_to_head_dim(hidden_states)
76
-
77
- # linear proj
78
- hidden_states = attn.to_out[0](hidden_states) + self.lora_scale * self.to_out_lora(hidden_states)
79
- # dropout
80
- hidden_states = attn.to_out[1](hidden_states)
81
-
82
- if input_ndim == 4:
83
- hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
84
-
85
- if attn.residual_connection:
86
- hidden_states = hidden_states + residual
87
-
88
- hidden_states = hidden_states / attn.rescale_output_factor
89
-
90
- return hidden_states
91
-
92
-
93
- class LoRAIPAttnProcessor(nn.Module):
94
- r"""
95
- Attention processor for IP-Adapater.
96
- Args:
97
- hidden_size (`int`):
98
- The hidden size of the attention layer.
99
- cross_attention_dim (`int`):
100
- The number of channels in the `encoder_hidden_states`.
101
- scale (`float`, defaults to 1.0):
102
- the weight scale of image prompt.
103
- num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
104
- The context length of the image features.
105
- """
106
-
107
- def __init__(self, hidden_size, cross_attention_dim=None, rank=4, network_alpha=None, lora_scale=1.0, scale=1.0, num_tokens=4):
108
- super().__init__()
109
-
110
- self.rank = rank
111
- self.lora_scale = lora_scale
112
-
113
- self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
114
- self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
115
- self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
116
- self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
117
-
118
- self.hidden_size = hidden_size
119
- self.cross_attention_dim = cross_attention_dim
120
- self.scale = scale
121
- self.num_tokens = num_tokens
122
-
123
- self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
124
- self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
125
-
126
- def __call__(
127
- self,
128
- attn,
129
- hidden_states,
130
- encoder_hidden_states=None,
131
- attention_mask=None,
132
- temb=None,
133
- ):
134
- residual = hidden_states
135
-
136
- if attn.spatial_norm is not None:
137
- hidden_states = attn.spatial_norm(hidden_states, temb)
138
-
139
- input_ndim = hidden_states.ndim
140
-
141
- if input_ndim == 4:
142
- batch_size, channel, height, width = hidden_states.shape
143
- hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
144
-
145
- batch_size, sequence_length, _ = (
146
- hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
147
- )
148
- attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
149
-
150
- if attn.group_norm is not None:
151
- hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
152
-
153
- query = attn.to_q(hidden_states) + self.lora_scale * self.to_q_lora(hidden_states)
154
-
155
- if encoder_hidden_states is None:
156
- encoder_hidden_states = hidden_states
157
- else:
158
- # get encoder_hidden_states, ip_hidden_states
159
- end_pos = encoder_hidden_states.shape[1] - self.num_tokens
160
- encoder_hidden_states, ip_hidden_states = (
161
- encoder_hidden_states[:, :end_pos, :],
162
- encoder_hidden_states[:, end_pos:, :],
163
- )
164
- if attn.norm_cross:
165
- encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
166
-
167
- key = attn.to_k(encoder_hidden_states) + self.lora_scale * self.to_k_lora(encoder_hidden_states)
168
- value = attn.to_v(encoder_hidden_states) + self.lora_scale * self.to_v_lora(encoder_hidden_states)
169
-
170
- query = attn.head_to_batch_dim(query)
171
- key = attn.head_to_batch_dim(key)
172
- value = attn.head_to_batch_dim(value)
173
-
174
- attention_probs = attn.get_attention_scores(query, key, attention_mask)
175
- hidden_states = torch.bmm(attention_probs, value)
176
- hidden_states = attn.batch_to_head_dim(hidden_states)
177
-
178
- # for ip-adapter
179
- ip_key = self.to_k_ip(ip_hidden_states)
180
- ip_value = self.to_v_ip(ip_hidden_states)
181
-
182
- ip_key = attn.head_to_batch_dim(ip_key)
183
- ip_value = attn.head_to_batch_dim(ip_value)
184
-
185
- ip_attention_probs = attn.get_attention_scores(query, ip_key, None)
186
- ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
187
- ip_hidden_states = attn.batch_to_head_dim(ip_hidden_states)
188
-
189
- hidden_states = hidden_states + self.scale * ip_hidden_states
190
-
191
- # linear proj
192
- hidden_states = attn.to_out[0](hidden_states) + self.lora_scale * self.to_out_lora(hidden_states)
193
- # dropout
194
- hidden_states = attn.to_out[1](hidden_states)
195
-
196
- if input_ndim == 4:
197
- hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
198
-
199
- if attn.residual_connection:
200
- hidden_states = hidden_states + residual
201
-
202
- hidden_states = hidden_states / attn.rescale_output_factor
203
-
204
- return hidden_states
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ip_adapter/custom_pipelines.py DELETED
@@ -1,394 +0,0 @@
1
- from typing import Any, Callable, Dict, List, Optional, Tuple, Union
2
-
3
- import torch
4
- from diffusers import StableDiffusionXLPipeline
5
- from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput
6
- from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import rescale_noise_cfg
7
-
8
- from .utils import is_torch2_available
9
-
10
- if is_torch2_available():
11
- from .attention_processor import IPAttnProcessor2_0 as IPAttnProcessor
12
- else:
13
- from .attention_processor import IPAttnProcessor
14
-
15
-
16
- class StableDiffusionXLCustomPipeline(StableDiffusionXLPipeline):
17
- def set_scale(self, scale):
18
- for attn_processor in self.unet.attn_processors.values():
19
- if isinstance(attn_processor, IPAttnProcessor):
20
- attn_processor.scale = scale
21
-
22
- @torch.no_grad()
23
- def __call__( # noqa: C901
24
- self,
25
- prompt: Optional[Union[str, List[str]]] = None,
26
- prompt_2: Optional[Union[str, List[str]]] = None,
27
- height: Optional[int] = None,
28
- width: Optional[int] = None,
29
- num_inference_steps: int = 50,
30
- denoising_end: Optional[float] = None,
31
- guidance_scale: float = 5.0,
32
- negative_prompt: Optional[Union[str, List[str]]] = None,
33
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
34
- num_images_per_prompt: Optional[int] = 1,
35
- eta: float = 0.0,
36
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
37
- latents: Optional[torch.FloatTensor] = None,
38
- prompt_embeds: Optional[torch.FloatTensor] = None,
39
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
40
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
41
- negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
42
- output_type: Optional[str] = "pil",
43
- return_dict: bool = True,
44
- callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
45
- callback_steps: int = 1,
46
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
47
- guidance_rescale: float = 0.0,
48
- original_size: Optional[Tuple[int, int]] = None,
49
- crops_coords_top_left: Tuple[int, int] = (0, 0),
50
- target_size: Optional[Tuple[int, int]] = None,
51
- negative_original_size: Optional[Tuple[int, int]] = None,
52
- negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
53
- negative_target_size: Optional[Tuple[int, int]] = None,
54
- control_guidance_start: float = 0.0,
55
- control_guidance_end: float = 1.0,
56
- ):
57
- r"""
58
- Function invoked when calling the pipeline for generation.
59
-
60
- Args:
61
- prompt (`str` or `List[str]`, *optional*):
62
- The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
63
- instead.
64
- prompt_2 (`str` or `List[str]`, *optional*):
65
- The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
66
- used in both text-encoders
67
- height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
68
- The height in pixels of the generated image. This is set to 1024 by default for the best results.
69
- Anything below 512 pixels won't work well for
70
- [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
71
- and checkpoints that are not specifically fine-tuned on low resolutions.
72
- width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
73
- The width in pixels of the generated image. This is set to 1024 by default for the best results.
74
- Anything below 512 pixels won't work well for
75
- [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
76
- and checkpoints that are not specifically fine-tuned on low resolutions.
77
- num_inference_steps (`int`, *optional*, defaults to 50):
78
- The number of denoising steps. More denoising steps usually lead to a higher quality image at the
79
- expense of slower inference.
80
- denoising_end (`float`, *optional*):
81
- When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
82
- completed before it is intentionally prematurely terminated. As a result, the returned sample will
83
- still retain a substantial amount of noise as determined by the discrete timesteps selected by the
84
- scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
85
- "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
86
- Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
87
- guidance_scale (`float`, *optional*, defaults to 5.0):
88
- Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
89
- `guidance_scale` is defined as `w` of equation 2. of [Imagen
90
- Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
91
- 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
92
- usually at the expense of lower image quality.
93
- negative_prompt (`str` or `List[str]`, *optional*):
94
- The prompt or prompts not to guide the image generation. If not defined, one has to pass
95
- `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
96
- less than `1`).
97
- negative_prompt_2 (`str` or `List[str]`, *optional*):
98
- The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
99
- `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
100
- num_images_per_prompt (`int`, *optional*, defaults to 1):
101
- The number of images to generate per prompt.
102
- eta (`float`, *optional*, defaults to 0.0):
103
- Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
104
- [`schedulers.DDIMScheduler`], will be ignored for others.
105
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
106
- One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
107
- to make generation deterministic.
108
- latents (`torch.FloatTensor`, *optional*):
109
- Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
110
- generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
111
- tensor will ge generated by sampling using the supplied random `generator`.
112
- prompt_embeds (`torch.FloatTensor`, *optional*):
113
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
114
- provided, text embeddings will be generated from `prompt` input argument.
115
- negative_prompt_embeds (`torch.FloatTensor`, *optional*):
116
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
117
- weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
118
- argument.
119
- pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
120
- Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
121
- If not provided, pooled text embeddings will be generated from `prompt` input argument.
122
- negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
123
- Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
124
- weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
125
- input argument.
126
- output_type (`str`, *optional*, defaults to `"pil"`):
127
- The output format of the generate image. Choose between
128
- [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
129
- return_dict (`bool`, *optional*, defaults to `True`):
130
- Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
131
- of a plain tuple.
132
- callback (`Callable`, *optional*):
133
- A function that will be called every `callback_steps` steps during inference. The function will be
134
- called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
135
- callback_steps (`int`, *optional*, defaults to 1):
136
- The frequency at which the `callback` function will be called. If not specified, the callback will be
137
- called at every step.
138
- cross_attention_kwargs (`dict`, *optional*):
139
- A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
140
- `self.processor` in
141
- [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
142
- guidance_rescale (`float`, *optional*, defaults to 0.7):
143
- Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
144
- Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
145
- [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
146
- Guidance rescale factor should fix overexposure when using zero terminal SNR.
147
- original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
148
- If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
149
- `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
150
- explained in section 2.2 of
151
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
152
- crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
153
- `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
154
- `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
155
- `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
156
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
157
- target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
158
- For most cases, `target_size` should be set to the desired height and width of the generated image. If
159
- not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
160
- section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
161
- negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
162
- To negatively condition the generation process based on a specific image resolution. Part of SDXL's
163
- micro-conditioning as explained in section 2.2 of
164
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
165
- information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
166
- negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
167
- To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
168
- micro-conditioning as explained in section 2.2 of
169
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
170
- information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
171
- negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
172
- To negatively condition the generation process based on a target image resolution. It should be as same
173
- as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
174
- [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
175
- information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
176
- control_guidance_start (`float`, *optional*, defaults to 0.0):
177
- The percentage of total steps at which the ControlNet starts applying.
178
- control_guidance_end (`float`, *optional*, defaults to 1.0):
179
- The percentage of total steps at which the ControlNet stops applying.
180
-
181
- Examples:
182
-
183
- Returns:
184
- [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
185
- [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
186
- `tuple`. When returning a tuple, the first element is a list with the generated images.
187
- """
188
- # 0. Default height and width to unet
189
- height = height or self.default_sample_size * self.vae_scale_factor
190
- width = width or self.default_sample_size * self.vae_scale_factor
191
-
192
- original_size = original_size or (height, width)
193
- target_size = target_size or (height, width)
194
-
195
- # 1. Check inputs. Raise error if not correct
196
- self.check_inputs(
197
- prompt,
198
- prompt_2,
199
- height,
200
- width,
201
- callback_steps,
202
- negative_prompt,
203
- negative_prompt_2,
204
- prompt_embeds,
205
- negative_prompt_embeds,
206
- pooled_prompt_embeds,
207
- negative_pooled_prompt_embeds,
208
- )
209
-
210
- # 2. Define call parameters
211
- if prompt is not None and isinstance(prompt, str):
212
- batch_size = 1
213
- elif prompt is not None and isinstance(prompt, list):
214
- batch_size = len(prompt)
215
- else:
216
- batch_size = prompt_embeds.shape[0]
217
-
218
- device = self._execution_device
219
-
220
- # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
221
- # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
222
- # corresponds to doing no classifier free guidance.
223
- do_classifier_free_guidance = guidance_scale > 1.0
224
-
225
- # 3. Encode input prompt
226
- text_encoder_lora_scale = (
227
- cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
228
- )
229
- (
230
- prompt_embeds,
231
- negative_prompt_embeds,
232
- pooled_prompt_embeds,
233
- negative_pooled_prompt_embeds,
234
- ) = self.encode_prompt(
235
- prompt=prompt,
236
- prompt_2=prompt_2,
237
- device=device,
238
- num_images_per_prompt=num_images_per_prompt,
239
- do_classifier_free_guidance=do_classifier_free_guidance,
240
- negative_prompt=negative_prompt,
241
- negative_prompt_2=negative_prompt_2,
242
- prompt_embeds=prompt_embeds,
243
- negative_prompt_embeds=negative_prompt_embeds,
244
- pooled_prompt_embeds=pooled_prompt_embeds,
245
- negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
246
- lora_scale=text_encoder_lora_scale,
247
- )
248
-
249
- # 4. Prepare timesteps
250
- self.scheduler.set_timesteps(num_inference_steps, device=device)
251
-
252
- timesteps = self.scheduler.timesteps
253
-
254
- # 5. Prepare latent variables
255
- num_channels_latents = self.unet.config.in_channels
256
- latents = self.prepare_latents(
257
- batch_size * num_images_per_prompt,
258
- num_channels_latents,
259
- height,
260
- width,
261
- prompt_embeds.dtype,
262
- device,
263
- generator,
264
- latents,
265
- )
266
-
267
- # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
268
- extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
269
-
270
- # 7. Prepare added time ids & embeddings
271
- add_text_embeds = pooled_prompt_embeds
272
- if self.text_encoder_2 is None:
273
- text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
274
- else:
275
- text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
276
-
277
- add_time_ids = self._get_add_time_ids(
278
- original_size,
279
- crops_coords_top_left,
280
- target_size,
281
- dtype=prompt_embeds.dtype,
282
- text_encoder_projection_dim=text_encoder_projection_dim,
283
- )
284
- if negative_original_size is not None and negative_target_size is not None:
285
- negative_add_time_ids = self._get_add_time_ids(
286
- negative_original_size,
287
- negative_crops_coords_top_left,
288
- negative_target_size,
289
- dtype=prompt_embeds.dtype,
290
- text_encoder_projection_dim=text_encoder_projection_dim,
291
- )
292
- else:
293
- negative_add_time_ids = add_time_ids
294
-
295
- if do_classifier_free_guidance:
296
- prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
297
- add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
298
- add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
299
-
300
- prompt_embeds = prompt_embeds.to(device)
301
- add_text_embeds = add_text_embeds.to(device)
302
- add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
303
-
304
- # 8. Denoising loop
305
- num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
306
-
307
- # 7.1 Apply denoising_end
308
- if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1:
309
- discrete_timestep_cutoff = int(
310
- round(
311
- self.scheduler.config.num_train_timesteps
312
- - (denoising_end * self.scheduler.config.num_train_timesteps)
313
- )
314
- )
315
- num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
316
- timesteps = timesteps[:num_inference_steps]
317
-
318
- # get init conditioning scale
319
- for attn_processor in self.unet.attn_processors.values():
320
- if isinstance(attn_processor, IPAttnProcessor):
321
- conditioning_scale = attn_processor.scale
322
- break
323
-
324
- with self.progress_bar(total=num_inference_steps) as progress_bar:
325
- for i, t in enumerate(timesteps):
326
- if (i / len(timesteps) < control_guidance_start) or ((i + 1) / len(timesteps) > control_guidance_end):
327
- self.set_scale(0.0)
328
- else:
329
- self.set_scale(conditioning_scale)
330
-
331
- # expand the latents if we are doing classifier free guidance
332
- latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
333
-
334
- latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
335
-
336
- # predict the noise residual
337
- added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
338
- noise_pred = self.unet(
339
- latent_model_input,
340
- t,
341
- encoder_hidden_states=prompt_embeds,
342
- cross_attention_kwargs=cross_attention_kwargs,
343
- added_cond_kwargs=added_cond_kwargs,
344
- return_dict=False,
345
- )[0]
346
-
347
- # perform guidance
348
- if do_classifier_free_guidance:
349
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
350
- noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
351
-
352
- if do_classifier_free_guidance and guidance_rescale > 0.0:
353
- # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
354
- noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
355
-
356
- # compute the previous noisy sample x_t -> x_t-1
357
- latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
358
-
359
- # call the callback, if provided
360
- if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
361
- progress_bar.update()
362
- if callback is not None and i % callback_steps == 0:
363
- callback(i, t, latents)
364
-
365
- if not output_type == "latent":
366
- # make sure the VAE is in float32 mode, as it overflows in float16
367
- needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
368
-
369
- if needs_upcasting:
370
- self.upcast_vae()
371
- latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
372
-
373
- image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
374
-
375
- # cast back to fp16 if needed
376
- if needs_upcasting:
377
- self.vae.to(dtype=torch.float16)
378
- else:
379
- image = latents
380
-
381
- if output_type != "latent":
382
- # apply watermark if available
383
- if self.watermark is not None:
384
- image = self.watermark.apply_watermark(image)
385
-
386
- image = self.image_processor.postprocess(image, output_type=output_type)
387
-
388
- # Offload all models
389
- self.maybe_free_model_hooks()
390
-
391
- if not return_dict:
392
- return (image,)
393
-
394
- return StableDiffusionXLPipelineOutput(images=image)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ip_adapter/ip_adapter.py DELETED
@@ -1,413 +0,0 @@
1
- import os
2
- from typing import List
3
-
4
- import torch
5
- from diffusers import StableDiffusionPipeline
6
- from diffusers.pipelines.controlnet import MultiControlNetModel
7
- from PIL import Image
8
- from safetensors import safe_open
9
- from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
10
-
11
- from .utils import is_torch2_available
12
-
13
- if is_torch2_available():
14
- from .attention_processor import (
15
- AttnProcessor2_0 as AttnProcessor,
16
- )
17
- from .attention_processor import (
18
- CNAttnProcessor2_0 as CNAttnProcessor,
19
- )
20
- from .attention_processor import (
21
- IPAttnProcessor2_0 as IPAttnProcessor,
22
- )
23
- else:
24
- from .attention_processor import AttnProcessor, CNAttnProcessor, IPAttnProcessor
25
- from .resampler import Resampler
26
-
27
-
28
- class ImageProjModel(torch.nn.Module):
29
- """Projection Model"""
30
-
31
- def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
32
- super().__init__()
33
-
34
- self.cross_attention_dim = cross_attention_dim
35
- self.clip_extra_context_tokens = clip_extra_context_tokens
36
- self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
37
- self.norm = torch.nn.LayerNorm(cross_attention_dim)
38
-
39
- def forward(self, image_embeds):
40
- embeds = image_embeds
41
- clip_extra_context_tokens = self.proj(embeds).reshape(
42
- -1, self.clip_extra_context_tokens, self.cross_attention_dim
43
- )
44
- clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
45
- return clip_extra_context_tokens
46
-
47
-
48
- class MLPProjModel(torch.nn.Module):
49
- """SD model with image prompt"""
50
- def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024):
51
- super().__init__()
52
-
53
- self.proj = torch.nn.Sequential(
54
- torch.nn.Linear(clip_embeddings_dim, clip_embeddings_dim),
55
- torch.nn.GELU(),
56
- torch.nn.Linear(clip_embeddings_dim, cross_attention_dim),
57
- torch.nn.LayerNorm(cross_attention_dim)
58
- )
59
-
60
- def forward(self, image_embeds):
61
- clip_extra_context_tokens = self.proj(image_embeds)
62
- return clip_extra_context_tokens
63
-
64
-
65
- class IPAdapter:
66
- def __init__(self, sd_pipe, image_encoder_path, ip_ckpt, device, num_tokens=4):
67
- self.device = device
68
- self.image_encoder_path = image_encoder_path
69
- self.ip_ckpt = ip_ckpt
70
- self.num_tokens = num_tokens
71
-
72
- self.pipe = sd_pipe.to(self.device)
73
- self.set_ip_adapter()
74
-
75
- # load image encoder
76
- self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(self.image_encoder_path).to(
77
- self.device, dtype=torch.float16
78
- )
79
- self.clip_image_processor = CLIPImageProcessor()
80
- # image proj model
81
- self.image_proj_model = self.init_proj()
82
-
83
- self.load_ip_adapter()
84
-
85
- def init_proj(self):
86
- image_proj_model = ImageProjModel(
87
- cross_attention_dim=self.pipe.unet.config.cross_attention_dim,
88
- clip_embeddings_dim=self.image_encoder.config.projection_dim,
89
- clip_extra_context_tokens=self.num_tokens,
90
- ).to(self.device, dtype=torch.float16)
91
- return image_proj_model
92
-
93
- def set_ip_adapter(self):
94
- unet = self.pipe.unet
95
- attn_procs = {}
96
- for name in unet.attn_processors.keys():
97
- cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
98
- if name.startswith("mid_block"):
99
- hidden_size = unet.config.block_out_channels[-1]
100
- elif name.startswith("up_blocks"):
101
- block_id = int(name[len("up_blocks.")])
102
- hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
103
- elif name.startswith("down_blocks"):
104
- block_id = int(name[len("down_blocks.")])
105
- hidden_size = unet.config.block_out_channels[block_id]
106
- if cross_attention_dim is None:
107
- attn_procs[name] = AttnProcessor()
108
- else:
109
- attn_procs[name] = IPAttnProcessor(
110
- hidden_size=hidden_size,
111
- cross_attention_dim=cross_attention_dim,
112
- scale=1.0,
113
- num_tokens=self.num_tokens,
114
- ).to(self.device, dtype=torch.float16)
115
- unet.set_attn_processor(attn_procs)
116
- if hasattr(self.pipe, "controlnet"):
117
- if isinstance(self.pipe.controlnet, MultiControlNetModel):
118
- for controlnet in self.pipe.controlnet.nets:
119
- controlnet.set_attn_processor(CNAttnProcessor(num_tokens=self.num_tokens))
120
- else:
121
- self.pipe.controlnet.set_attn_processor(CNAttnProcessor(num_tokens=self.num_tokens))
122
-
123
- def load_ip_adapter(self):
124
- if os.path.splitext(self.ip_ckpt)[-1] == ".safetensors":
125
- state_dict = {"image_proj": {}, "ip_adapter": {}}
126
- with safe_open(self.ip_ckpt, framework="pt", device="cpu") as f:
127
- for key in f.keys():
128
- if key.startswith("image_proj."):
129
- state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key)
130
- elif key.startswith("ip_adapter."):
131
- state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key)
132
- else:
133
- state_dict = torch.load(self.ip_ckpt, map_location="cpu")
134
- self.image_proj_model.load_state_dict(state_dict["image_proj"])
135
- ip_layers = torch.nn.ModuleList(self.pipe.unet.attn_processors.values())
136
- ip_layers.load_state_dict(state_dict["ip_adapter"])
137
-
138
- @torch.inference_mode()
139
- def get_image_embeds(self, pil_image=None, clip_image_embeds=None):
140
- if pil_image is not None:
141
- if isinstance(pil_image, Image.Image):
142
- pil_image = [pil_image]
143
- clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
144
- clip_image_embeds = self.image_encoder(clip_image.to(self.device, dtype=torch.float16)).image_embeds
145
- else:
146
- clip_image_embeds = clip_image_embeds.to(self.device, dtype=torch.float16)
147
- image_prompt_embeds = self.image_proj_model(clip_image_embeds)
148
- uncond_image_prompt_embeds = self.image_proj_model(torch.zeros_like(clip_image_embeds))
149
- return image_prompt_embeds, uncond_image_prompt_embeds
150
-
151
- def set_scale(self, scale):
152
- for attn_processor in self.pipe.unet.attn_processors.values():
153
- if isinstance(attn_processor, IPAttnProcessor):
154
- attn_processor.scale = scale
155
-
156
- def generate(
157
- self,
158
- pil_image=None,
159
- clip_image_embeds=None,
160
- prompt=None,
161
- negative_prompt=None,
162
- scale=1.0,
163
- num_samples=4,
164
- seed=None,
165
- guidance_scale=7.5,
166
- num_inference_steps=30,
167
- **kwargs,
168
- ):
169
- self.set_scale(scale)
170
-
171
- if pil_image is not None:
172
- num_prompts = 1 if isinstance(pil_image, Image.Image) else len(pil_image)
173
- else:
174
- num_prompts = clip_image_embeds.size(0)
175
-
176
- if prompt is None:
177
- prompt = "best quality, high quality"
178
- if negative_prompt is None:
179
- negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
180
-
181
- if not isinstance(prompt, List):
182
- prompt = [prompt] * num_prompts
183
- if not isinstance(negative_prompt, List):
184
- negative_prompt = [negative_prompt] * num_prompts
185
-
186
- image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(
187
- pil_image=pil_image, clip_image_embeds=clip_image_embeds
188
- )
189
- bs_embed, seq_len, _ = image_prompt_embeds.shape
190
- image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
191
- image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
192
- uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
193
- uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
194
-
195
- with torch.inference_mode():
196
- prompt_embeds_, negative_prompt_embeds_ = self.pipe.encode_prompt(
197
- prompt,
198
- device=self.device,
199
- num_images_per_prompt=num_samples,
200
- do_classifier_free_guidance=True,
201
- negative_prompt=negative_prompt,
202
- )
203
- prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1)
204
- negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1)
205
-
206
- generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None
207
- images = self.pipe(
208
- prompt_embeds=prompt_embeds,
209
- negative_prompt_embeds=negative_prompt_embeds,
210
- guidance_scale=guidance_scale,
211
- num_inference_steps=num_inference_steps,
212
- generator=generator,
213
- **kwargs,
214
- ).images
215
-
216
- return images
217
-
218
-
219
- class IPAdapterXL(IPAdapter):
220
- """SDXL"""
221
-
222
- def generate(
223
- self,
224
- pil_image,
225
- prompt=None,
226
- negative_prompt=None,
227
- scale=1.0,
228
- num_samples=4,
229
- seed=None,
230
- num_inference_steps=30,
231
- **kwargs,
232
- ):
233
- self.set_scale(scale)
234
-
235
- num_prompts = 1 if isinstance(pil_image, Image.Image) else len(pil_image)
236
-
237
- if prompt is None:
238
- prompt = "best quality, high quality"
239
- if negative_prompt is None:
240
- negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
241
-
242
- if not isinstance(prompt, List):
243
- prompt = [prompt] * num_prompts
244
- if not isinstance(negative_prompt, List):
245
- negative_prompt = [negative_prompt] * num_prompts
246
-
247
- image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(pil_image)
248
- bs_embed, seq_len, _ = image_prompt_embeds.shape
249
- image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
250
- image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
251
- uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
252
- uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
253
-
254
- with torch.inference_mode():
255
- (
256
- prompt_embeds,
257
- negative_prompt_embeds,
258
- pooled_prompt_embeds,
259
- negative_pooled_prompt_embeds,
260
- ) = self.pipe.encode_prompt(
261
- prompt,
262
- num_images_per_prompt=num_samples,
263
- do_classifier_free_guidance=True,
264
- negative_prompt=negative_prompt,
265
- )
266
- prompt_embeds = torch.cat([prompt_embeds, image_prompt_embeds], dim=1)
267
- negative_prompt_embeds = torch.cat([negative_prompt_embeds, uncond_image_prompt_embeds], dim=1)
268
-
269
- generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None
270
- images = self.pipe(
271
- prompt_embeds=prompt_embeds,
272
- negative_prompt_embeds=negative_prompt_embeds,
273
- pooled_prompt_embeds=pooled_prompt_embeds,
274
- negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
275
- num_inference_steps=num_inference_steps,
276
- generator=generator,
277
- **kwargs,
278
- ).images
279
-
280
- return images
281
-
282
-
283
- class IPAdapterPlus(IPAdapter):
284
- """IP-Adapter with fine-grained features"""
285
-
286
- def init_proj(self):
287
- image_proj_model = Resampler(
288
- dim=self.pipe.unet.config.cross_attention_dim,
289
- depth=4,
290
- dim_head=64,
291
- heads=12,
292
- num_queries=self.num_tokens,
293
- embedding_dim=self.image_encoder.config.hidden_size,
294
- output_dim=self.pipe.unet.config.cross_attention_dim,
295
- ff_mult=4,
296
- ).to(self.device, dtype=torch.float16)
297
- return image_proj_model
298
-
299
- @torch.inference_mode()
300
- def get_image_embeds(self, pil_image=None, clip_image_embeds=None):
301
- if isinstance(pil_image, Image.Image):
302
- pil_image = [pil_image]
303
- clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
304
- clip_image = clip_image.to(self.device, dtype=torch.float16)
305
- clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]
306
- image_prompt_embeds = self.image_proj_model(clip_image_embeds)
307
- uncond_clip_image_embeds = self.image_encoder(
308
- torch.zeros_like(clip_image), output_hidden_states=True
309
- ).hidden_states[-2]
310
- uncond_image_prompt_embeds = self.image_proj_model(uncond_clip_image_embeds)
311
- return image_prompt_embeds, uncond_image_prompt_embeds
312
-
313
-
314
- class IPAdapterFull(IPAdapterPlus):
315
- """IP-Adapter with full features"""
316
-
317
- def init_proj(self):
318
- image_proj_model = MLPProjModel(
319
- cross_attention_dim=self.pipe.unet.config.cross_attention_dim,
320
- clip_embeddings_dim=self.image_encoder.config.hidden_size,
321
- ).to(self.device, dtype=torch.float16)
322
- return image_proj_model
323
-
324
-
325
- class IPAdapterPlusXL(IPAdapter):
326
- """SDXL"""
327
-
328
- def init_proj(self):
329
- image_proj_model = Resampler(
330
- dim=1280,
331
- depth=4,
332
- dim_head=64,
333
- heads=20,
334
- num_queries=self.num_tokens,
335
- embedding_dim=self.image_encoder.config.hidden_size,
336
- output_dim=self.pipe.unet.config.cross_attention_dim,
337
- ff_mult=4,
338
- ).to(self.device, dtype=torch.float16)
339
- return image_proj_model
340
-
341
- @torch.inference_mode()
342
- def get_image_embeds(self, pil_image):
343
- if isinstance(pil_image, Image.Image):
344
- pil_image = [pil_image]
345
- clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
346
- clip_image = clip_image.to(self.device, dtype=torch.float16)
347
- clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]
348
- image_prompt_embeds = self.image_proj_model(clip_image_embeds)
349
- uncond_clip_image_embeds = self.image_encoder(
350
- torch.zeros_like(clip_image), output_hidden_states=True
351
- ).hidden_states[-2]
352
- uncond_image_prompt_embeds = self.image_proj_model(uncond_clip_image_embeds)
353
- return image_prompt_embeds, uncond_image_prompt_embeds
354
-
355
- def generate(
356
- self,
357
- pil_image,
358
- prompt=None,
359
- negative_prompt=None,
360
- scale=1.0,
361
- num_samples=4,
362
- seed=None,
363
- num_inference_steps=30,
364
- **kwargs,
365
- ):
366
- self.set_scale(scale)
367
-
368
- num_prompts = 1 if isinstance(pil_image, Image.Image) else len(pil_image)
369
-
370
- if prompt is None:
371
- prompt = "best quality, high quality"
372
- if negative_prompt is None:
373
- negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
374
-
375
- if not isinstance(prompt, List):
376
- prompt = [prompt] * num_prompts
377
- if not isinstance(negative_prompt, List):
378
- negative_prompt = [negative_prompt] * num_prompts
379
-
380
- image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(pil_image)
381
- bs_embed, seq_len, _ = image_prompt_embeds.shape
382
- image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
383
- image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
384
- uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
385
- uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
386
-
387
- with torch.inference_mode():
388
- (
389
- prompt_embeds,
390
- negative_prompt_embeds,
391
- pooled_prompt_embeds,
392
- negative_pooled_prompt_embeds,
393
- ) = self.pipe.encode_prompt(
394
- prompt,
395
- num_images_per_prompt=num_samples,
396
- do_classifier_free_guidance=True,
397
- negative_prompt=negative_prompt,
398
- )
399
- prompt_embeds = torch.cat([prompt_embeds, image_prompt_embeds], dim=1)
400
- negative_prompt_embeds = torch.cat([negative_prompt_embeds, uncond_image_prompt_embeds], dim=1)
401
-
402
- generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None
403
- images = self.pipe(
404
- prompt_embeds=prompt_embeds,
405
- negative_prompt_embeds=negative_prompt_embeds,
406
- pooled_prompt_embeds=pooled_prompt_embeds,
407
- negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
408
- num_inference_steps=num_inference_steps,
409
- generator=generator,
410
- **kwargs,
411
- ).images
412
-
413
- return images
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ip_adapter/ip_adapter_faceid.py DELETED
@@ -1,166 +0,0 @@
1
- import os
2
- from typing import List
3
-
4
- import torch
5
- from diffusers import StableDiffusionPipeline
6
- from diffusers.pipelines.controlnet import MultiControlNetModel
7
- from PIL import Image
8
- from safetensors import safe_open
9
- from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
10
-
11
- from .attention_processor_faceid import LoRAAttnProcessor, LoRAIPAttnProcessor
12
-
13
-
14
- class MLPProjModel(torch.nn.Module):
15
- """SD model with image prompt"""
16
- def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, num_tokens=4):
17
- super().__init__()
18
-
19
- self.cross_attention_dim = cross_attention_dim
20
- self.num_tokens = num_tokens
21
-
22
- self.proj = torch.nn.Sequential(
23
- torch.nn.Linear(id_embeddings_dim, id_embeddings_dim*2),
24
- torch.nn.GELU(),
25
- torch.nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens),
26
- )
27
- self.norm = torch.nn.LayerNorm(cross_attention_dim)
28
-
29
- def forward(self, id_embeds):
30
- x = self.proj(id_embeds)
31
- x = x.reshape(-1, self.num_tokens, self.cross_attention_dim)
32
- x = self.norm(x)
33
- return x
34
-
35
-
36
- class IPAdapterFaceID:
37
- def __init__(self, sd_pipe, ip_ckpt, device, lora_rank=128, num_tokens=4):
38
- self.device = device
39
- self.ip_ckpt = ip_ckpt
40
- self.lora_rank = lora_rank
41
- self.num_tokens = num_tokens
42
-
43
- self.pipe = sd_pipe.to(self.device)
44
- self.set_ip_adapter()
45
-
46
- # image proj model
47
- self.image_proj_model = self.init_proj()
48
-
49
- self.load_ip_adapter()
50
-
51
- def init_proj(self):
52
- image_proj_model = MLPProjModel(
53
- cross_attention_dim=self.pipe.unet.config.cross_attention_dim,
54
- id_embeddings_dim=512,
55
- num_tokens=self.num_tokens,
56
- ).to(self.device, dtype=torch.float16)
57
- return image_proj_model
58
-
59
- def set_ip_adapter(self):
60
- unet = self.pipe.unet
61
- attn_procs = {}
62
- for name in unet.attn_processors.keys():
63
- cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
64
- if name.startswith("mid_block"):
65
- hidden_size = unet.config.block_out_channels[-1]
66
- elif name.startswith("up_blocks"):
67
- block_id = int(name[len("up_blocks.")])
68
- hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
69
- elif name.startswith("down_blocks"):
70
- block_id = int(name[len("down_blocks.")])
71
- hidden_size = unet.config.block_out_channels[block_id]
72
- if cross_attention_dim is None:
73
- attn_procs[name] = LoRAAttnProcessor(
74
- hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=self.lora_rank,
75
- ).to(self.device, dtype=torch.float16)
76
- else:
77
- attn_procs[name] = LoRAIPAttnProcessor(
78
- hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, scale=1.0, rank=self.lora_rank, num_tokens=self.num_tokens,
79
- ).to(self.device, dtype=torch.float16)
80
- unet.set_attn_processor(attn_procs)
81
-
82
- def load_ip_adapter(self):
83
- if os.path.splitext(self.ip_ckpt)[-1] == ".safetensors":
84
- state_dict = {"image_proj": {}, "ip_adapter": {}}
85
- with safe_open(self.ip_ckpt, framework="pt", device="cpu") as f:
86
- for key in f.keys():
87
- if key.startswith("image_proj."):
88
- state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key)
89
- elif key.startswith("ip_adapter."):
90
- state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key)
91
- else:
92
- state_dict = torch.load(self.ip_ckpt, map_location="cpu")
93
- self.image_proj_model.load_state_dict(state_dict["image_proj"])
94
- ip_layers = torch.nn.ModuleList(self.pipe.unet.attn_processors.values())
95
- ip_layers.load_state_dict(state_dict["ip_adapter"])
96
-
97
- @torch.inference_mode()
98
- def get_image_embeds(self, faceid_embeds):
99
-
100
- faceid_embeds = faceid_embeds.to(self.device, dtype=torch.float16)
101
- image_prompt_embeds = self.image_proj_model(faceid_embeds)
102
- uncond_image_prompt_embeds = self.image_proj_model(torch.zeros_like(faceid_embeds))
103
- return image_prompt_embeds, uncond_image_prompt_embeds
104
-
105
- def set_scale(self, scale):
106
- for attn_processor in self.pipe.unet.attn_processors.values():
107
- if isinstance(attn_processor, LoRAIPAttnProcessor):
108
- attn_processor.scale = scale
109
-
110
- def generate(
111
- self,
112
- faceid_embeds=None,
113
- prompt=None,
114
- negative_prompt=None,
115
- scale=1.0,
116
- num_samples=4,
117
- seed=None,
118
- guidance_scale=7.5,
119
- num_inference_steps=30,
120
- **kwargs,
121
- ):
122
- self.set_scale(scale)
123
-
124
-
125
- num_prompts = faceid_embeds.size(0)
126
-
127
- if prompt is None:
128
- prompt = "best quality, high quality"
129
- if negative_prompt is None:
130
- negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
131
-
132
- if not isinstance(prompt, List):
133
- prompt = [prompt] * num_prompts
134
- if not isinstance(negative_prompt, List):
135
- negative_prompt = [negative_prompt] * num_prompts
136
-
137
- image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(faceid_embeds)
138
-
139
- bs_embed, seq_len, _ = image_prompt_embeds.shape
140
- image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
141
- image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
142
- uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
143
- uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
144
-
145
- with torch.inference_mode():
146
- prompt_embeds_, negative_prompt_embeds_ = self.pipe.encode_prompt(
147
- prompt,
148
- device=self.device,
149
- num_images_per_prompt=num_samples,
150
- do_classifier_free_guidance=True,
151
- negative_prompt=negative_prompt,
152
- )
153
- prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1)
154
- negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1)
155
-
156
- generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None
157
- images = self.pipe(
158
- prompt_embeds=prompt_embeds,
159
- negative_prompt_embeds=negative_prompt_embeds,
160
- guidance_scale=guidance_scale,
161
- num_inference_steps=num_inference_steps,
162
- generator=generator,
163
- **kwargs,
164
- ).images
165
-
166
- return images
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ip_adapter/resampler.py DELETED
@@ -1,158 +0,0 @@
1
- # modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py
2
- # and https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py
3
-
4
- import math
5
-
6
- import torch
7
- import torch.nn as nn
8
- from einops import rearrange
9
- from einops.layers.torch import Rearrange
10
-
11
-
12
- # FFN
13
- def FeedForward(dim, mult=4):
14
- inner_dim = int(dim * mult)
15
- return nn.Sequential(
16
- nn.LayerNorm(dim),
17
- nn.Linear(dim, inner_dim, bias=False),
18
- nn.GELU(),
19
- nn.Linear(inner_dim, dim, bias=False),
20
- )
21
-
22
-
23
- def reshape_tensor(x, heads):
24
- bs, length, width = x.shape
25
- # (bs, length, width) --> (bs, length, n_heads, dim_per_head)
26
- x = x.view(bs, length, heads, -1)
27
- # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
28
- x = x.transpose(1, 2)
29
- # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
30
- x = x.reshape(bs, heads, length, -1)
31
- return x
32
-
33
-
34
- class PerceiverAttention(nn.Module):
35
- def __init__(self, *, dim, dim_head=64, heads=8):
36
- super().__init__()
37
- self.scale = dim_head**-0.5
38
- self.dim_head = dim_head
39
- self.heads = heads
40
- inner_dim = dim_head * heads
41
-
42
- self.norm1 = nn.LayerNorm(dim)
43
- self.norm2 = nn.LayerNorm(dim)
44
-
45
- self.to_q = nn.Linear(dim, inner_dim, bias=False)
46
- self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
47
- self.to_out = nn.Linear(inner_dim, dim, bias=False)
48
-
49
- def forward(self, x, latents):
50
- """
51
- Args:
52
- x (torch.Tensor): image features
53
- shape (b, n1, D)
54
- latent (torch.Tensor): latent features
55
- shape (b, n2, D)
56
- """
57
- x = self.norm1(x)
58
- latents = self.norm2(latents)
59
-
60
- b, l, _ = latents.shape
61
-
62
- q = self.to_q(latents)
63
- kv_input = torch.cat((x, latents), dim=-2)
64
- k, v = self.to_kv(kv_input).chunk(2, dim=-1)
65
-
66
- q = reshape_tensor(q, self.heads)
67
- k = reshape_tensor(k, self.heads)
68
- v = reshape_tensor(v, self.heads)
69
-
70
- # attention
71
- scale = 1 / math.sqrt(math.sqrt(self.dim_head))
72
- weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
73
- weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
74
- out = weight @ v
75
-
76
- out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
77
-
78
- return self.to_out(out)
79
-
80
-
81
- class Resampler(nn.Module):
82
- def __init__(
83
- self,
84
- dim=1024,
85
- depth=8,
86
- dim_head=64,
87
- heads=16,
88
- num_queries=8,
89
- embedding_dim=768,
90
- output_dim=1024,
91
- ff_mult=4,
92
- max_seq_len: int = 257, # CLIP tokens + CLS token
93
- apply_pos_emb: bool = False,
94
- num_latents_mean_pooled: int = 0, # number of latents derived from mean pooled representation of the sequence
95
- ):
96
- super().__init__()
97
- self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None
98
-
99
- self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
100
-
101
- self.proj_in = nn.Linear(embedding_dim, dim)
102
-
103
- self.proj_out = nn.Linear(dim, output_dim)
104
- self.norm_out = nn.LayerNorm(output_dim)
105
-
106
- self.to_latents_from_mean_pooled_seq = (
107
- nn.Sequential(
108
- nn.LayerNorm(dim),
109
- nn.Linear(dim, dim * num_latents_mean_pooled),
110
- Rearrange("b (n d) -> b n d", n=num_latents_mean_pooled),
111
- )
112
- if num_latents_mean_pooled > 0
113
- else None
114
- )
115
-
116
- self.layers = nn.ModuleList([])
117
- for _ in range(depth):
118
- self.layers.append(
119
- nn.ModuleList(
120
- [
121
- PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
122
- FeedForward(dim=dim, mult=ff_mult),
123
- ]
124
- )
125
- )
126
-
127
- def forward(self, x):
128
- if self.pos_emb is not None:
129
- n, device = x.shape[1], x.device
130
- pos_emb = self.pos_emb(torch.arange(n, device=device))
131
- x = x + pos_emb
132
-
133
- latents = self.latents.repeat(x.size(0), 1, 1)
134
-
135
- x = self.proj_in(x)
136
-
137
- if self.to_latents_from_mean_pooled_seq:
138
- meanpooled_seq = masked_mean(x, dim=1, mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool))
139
- meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq)
140
- latents = torch.cat((meanpooled_latents, latents), dim=-2)
141
-
142
- for attn, ff in self.layers:
143
- latents = attn(x, latents) + latents
144
- latents = ff(latents) + latents
145
-
146
- latents = self.proj_out(latents)
147
- return self.norm_out(latents)
148
-
149
-
150
- def masked_mean(t, *, dim, mask=None):
151
- if mask is None:
152
- return t.mean(dim=dim)
153
-
154
- denom = mask.sum(dim=dim, keepdim=True)
155
- mask = rearrange(mask, "b n -> b n 1")
156
- masked_t = t.masked_fill(~mask, 0.0)
157
-
158
- return masked_t.sum(dim=dim) / denom.clamp(min=1e-5)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ip_adapter/test_resampler.py DELETED
@@ -1,44 +0,0 @@
1
- import torch
2
- from resampler import Resampler
3
- from transformers import CLIPVisionModel
4
-
5
- BATCH_SIZE = 2
6
- OUTPUT_DIM = 1280
7
- NUM_QUERIES = 8
8
- NUM_LATENTS_MEAN_POOLED = 4 # 0 for no mean pooling (previous behavior)
9
- APPLY_POS_EMB = True # False for no positional embeddings (previous behavior)
10
- IMAGE_ENCODER_NAME_OR_PATH = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
11
-
12
-
13
- def main():
14
- image_encoder = CLIPVisionModel.from_pretrained(IMAGE_ENCODER_NAME_OR_PATH)
15
- embedding_dim = image_encoder.config.hidden_size
16
- print(f"image_encoder hidden size: ", embedding_dim)
17
-
18
- image_proj_model = Resampler(
19
- dim=1024,
20
- depth=2,
21
- dim_head=64,
22
- heads=16,
23
- num_queries=NUM_QUERIES,
24
- embedding_dim=embedding_dim,
25
- output_dim=OUTPUT_DIM,
26
- ff_mult=2,
27
- max_seq_len=257,
28
- apply_pos_emb=APPLY_POS_EMB,
29
- num_latents_mean_pooled=NUM_LATENTS_MEAN_POOLED,
30
- )
31
-
32
- dummy_images = torch.randn(BATCH_SIZE, 3, 224, 224)
33
- with torch.no_grad():
34
- image_embeds = image_encoder(dummy_images, output_hidden_states=True).hidden_states[-2]
35
- print("image_embds shape: ", image_embeds.shape)
36
-
37
- with torch.no_grad():
38
- ip_tokens = image_proj_model(image_embeds)
39
- print("ip_tokens shape:", ip_tokens.shape)
40
- assert ip_tokens.shape == (BATCH_SIZE, NUM_QUERIES + NUM_LATENTS_MEAN_POOLED, OUTPUT_DIM)
41
-
42
-
43
- if __name__ == "__main__":
44
- main()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ip_adapter/utils.py DELETED
@@ -1,5 +0,0 @@
1
- import torch.nn.functional as F
2
-
3
-
4
- def is_torch2_available():
5
- return hasattr(F, "scaled_dot_product_attention")