Spaces:
Running
on
Zero
Running
on
Zero
Commit
•
ca51647
1
Parent(s):
8ddce9c
Delete ip_adapter
Browse files- ip_adapter/__init__.py +0 -9
- ip_adapter/attention_processor.py +0 -554
- ip_adapter/attention_processor_faceid.py +0 -204
- ip_adapter/custom_pipelines.py +0 -394
- ip_adapter/ip_adapter.py +0 -413
- ip_adapter/ip_adapter_faceid.py +0 -166
- ip_adapter/resampler.py +0 -158
- ip_adapter/test_resampler.py +0 -44
- ip_adapter/utils.py +0 -5
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")
|
|
|
|
|
|
|
|
|
|
|
|