怀墨 commited on
Commit
5db2d1c
·
1 Parent(s): cd10722

add pipelines code

Browse files
pipelines/__init__.py ADDED
File without changes
pipelines/sd3_model.py ADDED
@@ -0,0 +1,470 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Modified by AIDC-AI, 2025
2
+
3
+ # This project is licensed under the Attribution-NonCommercial 4.0 International
4
+ # License (SPDX-License-Identifier: CC-BY-NC-4.0).
5
+
6
+ # Copyright 2025 Stability AI, The HuggingFace Team and The InstantX Team. All rights reserved.
7
+ #
8
+ # Unless required by applicable law or agreed to in writing, software
9
+ # distributed under the License is distributed on an "AS IS" BASIS,
10
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11
+ # See the License for the specific language governing permissions and
12
+ # limitations under the License.
13
+
14
+ import torch
15
+ import torch.nn as nn
16
+ import torch.nn.functional as F
17
+
18
+
19
+ from typing import Any, Dict, List, Optional, Tuple, Union
20
+ from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin, SD3Transformer2DLoadersMixin
21
+ from diffusers.models.attention import FeedForward, JointTransformerBlock
22
+ from diffusers.models.attention_processor import (
23
+ Attention,
24
+ AttentionProcessor,
25
+ FusedJointAttnProcessor2_0,
26
+ JointAttnProcessor2_0,
27
+ )
28
+ from diffusers.models.modeling_utils import ModelMixin
29
+ from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero
30
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
31
+ from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
32
+ from diffusers.utils.torch_utils import maybe_allow_in_graph
33
+ from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed, TimestepEmbedding
34
+ from diffusers.models.modeling_outputs import Transformer2DModelOutput
35
+
36
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
37
+
38
+ @maybe_allow_in_graph
39
+ class SD3SingleTransformerBlock(nn.Module):
40
+ r"""
41
+ A Single Transformer block as part of the MMDiT architecture, used in Stable Diffusion 3 ControlNet.
42
+
43
+ Reference: https://arxiv.org/abs/2403.03206
44
+
45
+ Parameters:
46
+ dim (`int`): The number of channels in the input and output.
47
+ num_attention_heads (`int`): The number of heads to use for multi-head attention.
48
+ attention_head_dim (`int`): The number of channels in each head.
49
+ """
50
+
51
+ def __init__(
52
+ self,
53
+ dim: int,
54
+ num_attention_heads: int,
55
+ attention_head_dim: int,
56
+ ):
57
+ super().__init__()
58
+
59
+ self.norm1 = AdaLayerNormZero(dim)
60
+
61
+ if hasattr(F, "scaled_dot_product_attention"):
62
+ processor = JointAttnProcessor2_0()
63
+ else:
64
+ raise ValueError(
65
+ "The current PyTorch version does not support the `scaled_dot_product_attention` function."
66
+ )
67
+
68
+ self.attn = Attention(
69
+ query_dim=dim,
70
+ dim_head=attention_head_dim,
71
+ heads=num_attention_heads,
72
+ out_dim=dim,
73
+ bias=True,
74
+ processor=processor,
75
+ eps=1e-6,
76
+ )
77
+
78
+ self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
79
+ self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
80
+
81
+ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor):
82
+ norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
83
+ # Attention.
84
+ attn_output = self.attn(
85
+ hidden_states=norm_hidden_states,
86
+ encoder_hidden_states=None,
87
+ )
88
+
89
+ # Process attention outputs for the `hidden_states`.
90
+ attn_output = gate_msa.unsqueeze(1) * attn_output
91
+ hidden_states = hidden_states + attn_output
92
+
93
+ norm_hidden_states = self.norm2(hidden_states)
94
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
95
+
96
+ ff_output = self.ff(norm_hidden_states)
97
+ ff_output = gate_mlp.unsqueeze(1) * ff_output
98
+
99
+ hidden_states = hidden_states + ff_output
100
+
101
+ return hidden_states
102
+
103
+
104
+ class SD3Transformer2DModel(
105
+ ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, SD3Transformer2DLoadersMixin
106
+ ):
107
+ """
108
+ The Transformer model introduced in Stable Diffusion 3.
109
+
110
+ Reference: https://arxiv.org/abs/2403.03206
111
+
112
+ Parameters:
113
+ sample_size (`int`): The width of the latent images. This is fixed during training since
114
+ it is used to learn a number of position embeddings.
115
+ patch_size (`int`): Patch size to turn the input data into small patches.
116
+ in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
117
+ num_layers (`int`, *optional*, defaults to 18): The number of layers of Transformer blocks to use.
118
+ attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
119
+ num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
120
+ cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
121
+ caption_projection_dim (`int`): Number of dimensions to use when projecting the `encoder_hidden_states`.
122
+ pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
123
+ out_channels (`int`, defaults to 16): Number of output channels.
124
+
125
+ """
126
+
127
+ _supports_gradient_checkpointing = True
128
+ _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
129
+
130
+ @register_to_config
131
+ def __init__(
132
+ self,
133
+ sample_size: int = 128,
134
+ patch_size: int = 2,
135
+ in_channels: int = 16,
136
+ num_layers: int = 18,
137
+ attention_head_dim: int = 64,
138
+ num_attention_heads: int = 18,
139
+ joint_attention_dim: int = 4096,
140
+ caption_projection_dim: int = 1152,
141
+ pooled_projection_dim: int = 2048,
142
+ out_channels: int = 16,
143
+ pos_embed_max_size: int = 96,
144
+ dual_attention_layers: Tuple[
145
+ int, ...
146
+ ] = (), # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5
147
+ qk_norm: Optional[str] = None,
148
+ guidance=False,
149
+ txt_align_guidance=False
150
+ ):
151
+ super().__init__()
152
+ default_out_channels = in_channels
153
+ self.out_channels = out_channels if out_channels is not None else default_out_channels
154
+ self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
155
+
156
+ self.pos_embed = PatchEmbed(
157
+ height=self.config.sample_size,
158
+ width=self.config.sample_size,
159
+ patch_size=self.config.patch_size,
160
+ in_channels=self.config.in_channels,
161
+ embed_dim=self.inner_dim,
162
+ pos_embed_max_size=pos_embed_max_size, # hard-code for now.
163
+ )
164
+ self.time_text_embed = CombinedTimestepTextProjEmbeddings(
165
+ embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
166
+ )
167
+ self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.config.caption_projection_dim)
168
+
169
+ # `attention_head_dim` is doubled to account for the mixing.
170
+ # It needs to crafted when we get the actual checkpoints.
171
+ self.transformer_blocks = nn.ModuleList(
172
+ [
173
+ JointTransformerBlock(
174
+ dim=self.inner_dim,
175
+ num_attention_heads=self.config.num_attention_heads,
176
+ attention_head_dim=self.config.attention_head_dim,
177
+ context_pre_only=i == num_layers - 1,
178
+ qk_norm=qk_norm,
179
+ use_dual_attention=True if i in dual_attention_layers else False,
180
+ )
181
+ for i in range(self.config.num_layers)
182
+ ]
183
+ )
184
+
185
+ self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
186
+ self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
187
+
188
+ self.gradient_checkpointing = False
189
+
190
+ self.guidance = guidance
191
+ if self.guidance:
192
+ self.guidance_in = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim)
193
+
194
+ self.txt_align_guidance = txt_align_guidance
195
+
196
+ assert not (self.guidance and self.txt_align_guidance)
197
+
198
+ # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
199
+ def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
200
+ """
201
+ Sets the attention processor to use [feed forward
202
+ chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
203
+
204
+ Parameters:
205
+ chunk_size (`int`, *optional*):
206
+ The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
207
+ over each tensor of dim=`dim`.
208
+ dim (`int`, *optional*, defaults to `0`):
209
+ The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
210
+ or dim=1 (sequence length).
211
+ """
212
+ if dim not in [0, 1]:
213
+ raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
214
+
215
+ # By default chunk size is 1
216
+ chunk_size = chunk_size or 1
217
+
218
+ def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
219
+ if hasattr(module, "set_chunk_feed_forward"):
220
+ module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
221
+
222
+ for child in module.children():
223
+ fn_recursive_feed_forward(child, chunk_size, dim)
224
+
225
+ for module in self.children():
226
+ fn_recursive_feed_forward(module, chunk_size, dim)
227
+
228
+ # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking
229
+ def disable_forward_chunking(self):
230
+ def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
231
+ if hasattr(module, "set_chunk_feed_forward"):
232
+ module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
233
+
234
+ for child in module.children():
235
+ fn_recursive_feed_forward(child, chunk_size, dim)
236
+
237
+ for module in self.children():
238
+ fn_recursive_feed_forward(module, None, 0)
239
+
240
+ @property
241
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
242
+ def attn_processors(self) -> Dict[str, AttentionProcessor]:
243
+ r"""
244
+ Returns:
245
+ `dict` of attention processors: A dictionary containing all attention processors used in the model with
246
+ indexed by its weight name.
247
+ """
248
+ # set recursively
249
+ processors = {}
250
+
251
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
252
+ if hasattr(module, "get_processor"):
253
+ processors[f"{name}.processor"] = module.get_processor()
254
+
255
+ for sub_name, child in module.named_children():
256
+ fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
257
+
258
+ return processors
259
+
260
+ for name, module in self.named_children():
261
+ fn_recursive_add_processors(name, module, processors)
262
+
263
+ return processors
264
+
265
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
266
+ def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
267
+ r"""
268
+ Sets the attention processor to use to compute attention.
269
+
270
+ Parameters:
271
+ processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
272
+ The instantiated processor class or a dictionary of processor classes that will be set as the processor
273
+ for **all** `Attention` layers.
274
+
275
+ If `processor` is a dict, the key needs to define the path to the corresponding cross attention
276
+ processor. This is strongly recommended when setting trainable attention processors.
277
+
278
+ """
279
+ count = len(self.attn_processors.keys())
280
+
281
+ if isinstance(processor, dict) and len(processor) != count:
282
+ raise ValueError(
283
+ f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
284
+ f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
285
+ )
286
+
287
+ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
288
+ if hasattr(module, "set_processor"):
289
+ if not isinstance(processor, dict):
290
+ module.set_processor(processor)
291
+ else:
292
+ module.set_processor(processor.pop(f"{name}.processor"))
293
+
294
+ for sub_name, child in module.named_children():
295
+ fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
296
+
297
+ for name, module in self.named_children():
298
+ fn_recursive_attn_processor(name, module, processor)
299
+
300
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedJointAttnProcessor2_0
301
+ def fuse_qkv_projections(self):
302
+ """
303
+ Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
304
+ are fused. For cross-attention modules, key and value projection matrices are fused.
305
+
306
+ <Tip warning={true}>
307
+
308
+ This API is 🧪 experimental.
309
+
310
+ </Tip>
311
+ """
312
+ self.original_attn_processors = None
313
+
314
+ for _, attn_processor in self.attn_processors.items():
315
+ if "Added" in str(attn_processor.__class__.__name__):
316
+ raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
317
+
318
+ self.original_attn_processors = self.attn_processors
319
+
320
+ for module in self.modules():
321
+ if isinstance(module, Attention):
322
+ module.fuse_projections(fuse=True)
323
+
324
+ self.set_attn_processor(FusedJointAttnProcessor2_0())
325
+
326
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
327
+ def unfuse_qkv_projections(self):
328
+ """Disables the fused QKV projection if enabled.
329
+
330
+ <Tip warning={true}>
331
+
332
+ This API is 🧪 experimental.
333
+
334
+ </Tip>
335
+
336
+ """
337
+ if self.original_attn_processors is not None:
338
+ self.set_attn_processor(self.original_attn_processors)
339
+
340
+ def forward(
341
+ self,
342
+ hidden_states: torch.FloatTensor,
343
+ encoder_hidden_states: torch.FloatTensor = None,
344
+ pooled_projections: torch.FloatTensor = None,
345
+ timestep: torch.LongTensor = None,
346
+ block_controlnet_hidden_states: List = None,
347
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
348
+ return_dict: bool = True,
349
+ skip_layers: Optional[List[int]] = None,
350
+ guidance = None,
351
+ txt_align_guidance = None,
352
+ txt_align_vec = None
353
+ ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
354
+ """
355
+ The [`SD3Transformer2DModel`] forward method.
356
+
357
+ Args:
358
+ hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
359
+ Input `hidden_states`.
360
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
361
+ Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
362
+ pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`):
363
+ Embeddings projected from the embeddings of input conditions.
364
+ timestep (`torch.LongTensor`):
365
+ Used to indicate denoising step.
366
+ block_controlnet_hidden_states (`list` of `torch.Tensor`):
367
+ A list of tensors that if specified are added to the residuals of transformer blocks.
368
+ joint_attention_kwargs (`dict`, *optional*):
369
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
370
+ `self.processor` in
371
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
372
+ return_dict (`bool`, *optional*, defaults to `True`):
373
+ Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
374
+ tuple.
375
+ skip_layers (`list` of `int`, *optional*):
376
+ A list of layer indices to skip during the forward pass.
377
+
378
+ Returns:
379
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
380
+ `tuple` where the first element is the sample tensor.
381
+ """
382
+ if joint_attention_kwargs is not None:
383
+ joint_attention_kwargs = joint_attention_kwargs.copy()
384
+ lora_scale = joint_attention_kwargs.pop("scale", 1.0)
385
+ else:
386
+ lora_scale = 1.0
387
+
388
+ if USE_PEFT_BACKEND:
389
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
390
+ scale_lora_layers(self, lora_scale)
391
+ else:
392
+ if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
393
+ logger.warning(
394
+ "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
395
+ )
396
+
397
+ height, width = hidden_states.shape[-2:]
398
+
399
+ hidden_states = self.pos_embed(hidden_states) # takes care of adding positional embeddings too.
400
+ temb = self.time_text_embed(timestep, pooled_projections)
401
+
402
+ if hasattr(self, 'guidance_in'):
403
+ assert guidance is not None
404
+ temb = temb + self.guidance_in(self.time_text_embed.time_proj(guidance).to(dtype=pooled_projections.dtype))
405
+
406
+ ###############
407
+ if self.txt_align_guidance:
408
+ assert (txt_align_guidance is not None) and (txt_align_vec is not None)
409
+ temb = temb + (self.time_text_embed.timestep_embedder(self.time_text_embed.time_proj(txt_align_guidance).to(dtype=pooled_projections.dtype)) * self.time_text_embed.text_embedder(txt_align_vec)) / 256.
410
+ ###############
411
+
412
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
413
+
414
+ if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
415
+ ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
416
+ ip_hidden_states, ip_temb = self.image_proj(ip_adapter_image_embeds, timestep)
417
+
418
+ joint_attention_kwargs.update(ip_hidden_states=ip_hidden_states, temb=ip_temb)
419
+
420
+ for index_block, block in enumerate(self.transformer_blocks):
421
+ # Skip specified layers
422
+ is_skip = True if skip_layers is not None and index_block in skip_layers else False
423
+
424
+ if torch.is_grad_enabled() and self.gradient_checkpointing and not is_skip:
425
+ encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
426
+ block,
427
+ hidden_states,
428
+ encoder_hidden_states,
429
+ temb,
430
+ joint_attention_kwargs,
431
+ )
432
+ elif not is_skip:
433
+ encoder_hidden_states, hidden_states = block(
434
+ hidden_states=hidden_states,
435
+ encoder_hidden_states=encoder_hidden_states,
436
+ temb=temb,
437
+ joint_attention_kwargs=joint_attention_kwargs,
438
+ )
439
+
440
+ # controlnet residual
441
+ if block_controlnet_hidden_states is not None and block.context_pre_only is False:
442
+ interval_control = len(self.transformer_blocks) / len(block_controlnet_hidden_states)
443
+ hidden_states = hidden_states + block_controlnet_hidden_states[int(index_block / interval_control)]
444
+
445
+ hidden_states = self.norm_out(hidden_states, temb)
446
+ hidden_states = self.proj_out(hidden_states)
447
+
448
+ # unpatchify
449
+ patch_size = self.config.patch_size
450
+ height = height // patch_size
451
+ width = width // patch_size
452
+
453
+ hidden_states = hidden_states.reshape(
454
+ shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels)
455
+ )
456
+ hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
457
+ output = hidden_states.reshape(
458
+ shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
459
+ )
460
+
461
+ if USE_PEFT_BACKEND:
462
+ # remove `lora_scale` from each PEFT layer
463
+ unscale_lora_layers(self, lora_scale)
464
+
465
+ if not return_dict:
466
+ return (output,)
467
+
468
+ return Transformer2DModelOutput(sample=output)
469
+
470
+
pipelines/sd3_teefusion_pipeline.py ADDED
@@ -0,0 +1,264 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (C) 2025 AIDC-AI
2
+ # This project is licensed under the Attribution-NonCommercial 4.0 International
3
+ # License (SPDX-License-Identifier: CC-BY-NC-4.0).
4
+
5
+ # Unless required by applicable law or agreed to in writing, software
6
+ # distributed under the License is distributed on an "AS IS" BASIS,
7
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
8
+ # See the License for the specific language governing permissions and
9
+ # limitations under the License.
10
+
11
+ import os
12
+ import torch
13
+ import torch.nn as nn
14
+
15
+ from typing import Union, List, Any, Optional
16
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
17
+ from PIL import Image
18
+ from diffusers import DiffusionPipeline, AutoencoderKL
19
+ from transformers import CLIPTextModelWithProjection, T5EncoderModel, CLIPTokenizer, T5Tokenizer
20
+
21
+ def get_noise(
22
+ num_samples: int,
23
+ channel: int,
24
+ height: int,
25
+ width: int,
26
+ device: torch.device,
27
+ dtype: torch.dtype,
28
+ seed: int,
29
+ ):
30
+ return torch.randn(
31
+ num_samples,
32
+ channel,
33
+ height // 8,
34
+ width // 8,
35
+ device=device,
36
+ dtype=dtype,
37
+ generator=torch.Generator(device=device).manual_seed(seed),
38
+ )
39
+
40
+ def get_clip_prompt_embeds(
41
+ clip_tokenizers,
42
+ clip_text_encoders,
43
+ prompt: Union[str, List[str]],
44
+ num_images_per_prompt: int = 1,
45
+ device: Optional[torch.device] = None,
46
+ clip_skip: Optional[int] = None,
47
+ clip_model_index: int = 0,
48
+ ):
49
+
50
+ tokenizer_max_length = 77
51
+ tokenizer = clip_tokenizers[clip_model_index]
52
+ text_encoder = clip_text_encoders[clip_model_index]
53
+
54
+ batch_size = len(prompt)
55
+
56
+ text_inputs = tokenizer(
57
+ prompt,
58
+ padding="max_length",
59
+ max_length=tokenizer_max_length,
60
+ truncation=True,
61
+ return_tensors="pt",
62
+ )
63
+
64
+ text_input_ids = text_inputs.input_ids
65
+ untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
66
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
67
+ removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer_max_length - 1 : -1])
68
+
69
+ prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
70
+ pooled_prompt_embeds = prompt_embeds[0]
71
+
72
+ if clip_skip is None:
73
+ prompt_embeds = prompt_embeds.hidden_states[-2]
74
+ else:
75
+ prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
76
+
77
+ prompt_embeds = prompt_embeds.to(dtype=text_encoder.dtype, device=device)
78
+
79
+ _, seq_len, _ = prompt_embeds.shape
80
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
81
+ prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
82
+
83
+ pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
84
+ pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
85
+
86
+ return prompt_embeds, pooled_prompt_embeds
87
+
88
+ def get_t5_prompt_embeds(
89
+ tokenizer_3,
90
+ text_encoder_3,
91
+ prompt: Union[str, List[str]] = None,
92
+ num_images_per_prompt: int = 1,
93
+ max_sequence_length: int = 256,
94
+ device: Optional[torch.device] = None,
95
+ dtype: Optional[torch.dtype] = None,
96
+ ):
97
+
98
+ tokenizer_max_length = 77
99
+ batch_size = len(prompt)
100
+
101
+ text_inputs = tokenizer_3(
102
+ prompt,
103
+ padding="max_length",
104
+ max_length=max_sequence_length,
105
+ truncation=True,
106
+ add_special_tokens=True,
107
+ return_tensors="pt",
108
+ )
109
+ text_input_ids = text_inputs.input_ids
110
+ untruncated_ids = tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
111
+
112
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
113
+ removed_text = tokenizer_3.batch_decode(untruncated_ids[:, tokenizer_max_length - 1 : -1])
114
+
115
+ prompt_embeds = text_encoder_3(text_input_ids.to(device))[0]
116
+
117
+ dtype = text_encoder_3.dtype
118
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
119
+
120
+ _, seq_len, _ = prompt_embeds.shape
121
+
122
+ # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
123
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
124
+ prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
125
+
126
+ return prompt_embeds
127
+
128
+
129
+ @torch.no_grad()
130
+ def encode_text(clip_tokenizers, clip_text_encoders, tokenizer_3, text_encoder_3, prompt, device, max_sequence_length=256):
131
+
132
+ prompt_embed, pooled_prompt_embed = get_clip_prompt_embeds(clip_tokenizers, clip_text_encoders, prompt=prompt, device=device, clip_model_index=0)
133
+ prompt_2_embed, pooled_prompt_2_embed = get_clip_prompt_embeds(clip_tokenizers, clip_text_encoders, prompt=prompt, device=device, clip_model_index=1)
134
+ clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)
135
+
136
+ t5_prompt_embed = get_t5_prompt_embeds(tokenizer_3, text_encoder_3, prompt=prompt, max_sequence_length=max_sequence_length, device=device)
137
+
138
+ clip_prompt_embeds = torch.nn.functional.pad(clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]))
139
+
140
+ prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
141
+ pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
142
+
143
+ return prompt_embeds, pooled_prompt_embeds
144
+
145
+
146
+ class TeEFusionSD3Pipeline(DiffusionPipeline, ConfigMixin):
147
+
148
+ @register_to_config
149
+ def __init__(
150
+ self,
151
+ transformer: nn.Module,
152
+ text_encoder: CLIPTextModelWithProjection,
153
+ text_encoder_2: CLIPTextModelWithProjection,
154
+ text_encoder_3: T5EncoderModel,
155
+ tokenizer: CLIPTokenizer,
156
+ tokenizer_2: CLIPTokenizer,
157
+ tokenizer_3: T5Tokenizer,
158
+ vae: AutoencoderKL,
159
+ scheduler: Any
160
+ ):
161
+ super().__init__()
162
+
163
+ self.register_modules(
164
+ transformer=transformer,
165
+ text_encoder=text_encoder,
166
+ text_encoder_2=text_encoder_2,
167
+ text_encoder_3=text_encoder_3,
168
+ tokenizer=tokenizer,
169
+ tokenizer_2=tokenizer_2,
170
+ tokenizer_3=tokenizer_3,
171
+ vae=vae,
172
+ scheduler=scheduler
173
+ )
174
+
175
+
176
+ @classmethod
177
+ def from_pretrained(
178
+ cls,
179
+ pretrained_model_name_or_path: Union[str, os.PathLike],
180
+ **kwargs,
181
+ ) -> "TeEFusionSD3Pipeline":
182
+
183
+ return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
184
+
185
+ def save_pretrained(self, save_directory: Union[str, os.PathLike]):
186
+ super().save_pretrained(save_directory)
187
+
188
+ @torch.no_grad()
189
+ def __call__(
190
+ self,
191
+ prompt: Union[str, List[str]],
192
+ num_inference_steps: int = 50,
193
+ guidance_scale: float = 7.5,
194
+ latents: torch.FloatTensor = None,
195
+ height: int = 1024,
196
+ width: int = 1024,
197
+ seed: int = 0,
198
+ ):
199
+ if isinstance(prompt, str):
200
+ prompt = [prompt]
201
+
202
+ device = self.transformer.device
203
+
204
+ clip_tokenizers = [self.tokenizer, self.tokenizer_2]
205
+ clip_text_encoders = [self.text_encoder, self.text_encoder_2]
206
+
207
+ prompt_embeds, pooled_prompt_embeds = encode_text(clip_tokenizers, clip_text_encoders, self.tokenizer_3, self.text_encoder_3, prompt, device)
208
+
209
+ _, negative_pooled_prompt_embeds = encode_text(clip_tokenizers, clip_text_encoders, self.tokenizer_3, self.text_encoder_3, [''], device)
210
+
211
+
212
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
213
+ timesteps = self.scheduler.timesteps
214
+
215
+ bs = len(prompt)
216
+ channels = self.transformer.config.in_channels
217
+ height = 16 * (height // 16)
218
+ width = 16 * (width // 16)
219
+
220
+ # prepare input
221
+ if latents is None:
222
+ latents = get_noise(
223
+ bs,
224
+ channels,
225
+ height,
226
+ width,
227
+ device=device,
228
+ dtype=self.transformer.dtype,
229
+ seed=seed,
230
+ )
231
+
232
+ for i, t in enumerate(timesteps):
233
+ noise_pred = self.transformer(
234
+ hidden_states=latents,
235
+ timestep=t.reshape(1),
236
+ encoder_hidden_states=prompt_embeds,
237
+ pooled_projections=pooled_prompt_embeds,
238
+ return_dict=False,
239
+ txt_align_guidance=torch.tensor(data=(guidance_scale,), dtype=self.transformer.dtype, device=self.transformer.device) * 1000.,
240
+ txt_align_vec=pooled_prompt_embeds - negative_pooled_prompt_embeds
241
+ )[0]
242
+
243
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
244
+
245
+ x = latents.float()
246
+
247
+ with torch.no_grad():
248
+ with torch.autocast(device_type=device.type, dtype=torch.float32):
249
+ if hasattr(self.vae.config, 'scaling_factor') and self.vae.config.scaling_factor is not None:
250
+ x = x / self.vae.config.scaling_factor
251
+ if hasattr(self.vae.config, 'shift_factor') and self.vae.config.shift_factor is not None:
252
+ x = x + self.vae.config.shift_factor
253
+ x = self.vae.decode(x, return_dict=False)[0]
254
+
255
+ # bring into PIL format and save
256
+ x = (x / 2 + 0.5).clamp(0, 1)
257
+ x = x.cpu().permute(0, 2, 3, 1).float().numpy()
258
+ images = (x * 255).round().astype("uint8")
259
+ pil_images = [Image.fromarray(image) for image in images]
260
+
261
+ return pil_images
262
+
263
+
264
+