multimodalart HF Staff commited on
Commit
7f21e40
·
verified ·
1 Parent(s): 7ce70f3

Upload 5 files

Browse files
controlnet_flux.py ADDED
@@ -0,0 +1,420 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass
2
+ from typing import Any, Dict, List, Optional, Tuple, Union
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+
7
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
8
+ from diffusers.loaders import PeftAdapterMixin
9
+ from diffusers.models.modeling_utils import ModelMixin
10
+ from diffusers.models.attention_processor import AttentionProcessor
11
+ from diffusers.utils import (
12
+ USE_PEFT_BACKEND,
13
+ is_torch_version,
14
+ logging,
15
+ scale_lora_layers,
16
+ unscale_lora_layers,
17
+ )
18
+ from diffusers.models.controlnet import BaseOutput, zero_module
19
+ from diffusers.models.embeddings import (
20
+ CombinedTimestepGuidanceTextProjEmbeddings,
21
+ CombinedTimestepTextProjEmbeddings,
22
+ FluxPosEmbed,
23
+ )
24
+ from diffusers.models.modeling_outputs import Transformer2DModelOutput
25
+ from transformer_flux import (
26
+ FluxSingleTransformerBlock,
27
+ FluxTransformerBlock,
28
+ )
29
+
30
+
31
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
32
+
33
+
34
+ @dataclass
35
+ class FluxControlNetOutput(BaseOutput):
36
+ controlnet_block_samples: Tuple[torch.Tensor]
37
+ controlnet_single_block_samples: Tuple[torch.Tensor]
38
+
39
+
40
+ class FluxControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
41
+ _supports_gradient_checkpointing = True
42
+
43
+ @register_to_config
44
+ def __init__(
45
+ self,
46
+ patch_size: int = 1,
47
+ in_channels: int = 64,
48
+ num_layers: int = 19,
49
+ num_single_layers: int = 38,
50
+ attention_head_dim: int = 128,
51
+ num_attention_heads: int = 24,
52
+ joint_attention_dim: int = 4096,
53
+ pooled_projection_dim: int = 768,
54
+ guidance_embeds: bool = False,
55
+ axes_dims_rope: List[int] = [16, 56, 56],
56
+ extra_condition_channels: int = 1 * 4,
57
+ ):
58
+ super().__init__()
59
+ self.out_channels = in_channels
60
+ self.inner_dim = num_attention_heads * attention_head_dim
61
+
62
+ # self.pos_embed = EmbedND(
63
+ # dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope
64
+ # )
65
+ self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
66
+ text_time_guidance_cls = (
67
+ CombinedTimestepGuidanceTextProjEmbeddings
68
+ if guidance_embeds
69
+ else CombinedTimestepTextProjEmbeddings
70
+ )
71
+ self.time_text_embed = text_time_guidance_cls(
72
+ embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
73
+ )
74
+
75
+ self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
76
+ self.x_embedder = nn.Linear(in_channels, self.inner_dim)
77
+
78
+ self.transformer_blocks = nn.ModuleList(
79
+ [
80
+ FluxTransformerBlock(
81
+ dim=self.inner_dim,
82
+ num_attention_heads=num_attention_heads,
83
+ attention_head_dim=attention_head_dim,
84
+ )
85
+ for _ in range(num_layers)
86
+ ]
87
+ )
88
+
89
+ self.single_transformer_blocks = nn.ModuleList(
90
+ [
91
+ FluxSingleTransformerBlock(
92
+ dim=self.inner_dim,
93
+ num_attention_heads=num_attention_heads,
94
+ attention_head_dim=attention_head_dim,
95
+ )
96
+ for _ in range(num_single_layers)
97
+ ]
98
+ )
99
+
100
+ # controlnet_blocks
101
+ self.controlnet_blocks = nn.ModuleList([])
102
+ for _ in range(len(self.transformer_blocks)):
103
+ self.controlnet_blocks.append(
104
+ zero_module(nn.Linear(self.inner_dim, self.inner_dim))
105
+ )
106
+
107
+ self.controlnet_single_blocks = nn.ModuleList([])
108
+ for _ in range(len(self.single_transformer_blocks)):
109
+ self.controlnet_single_blocks.append(
110
+ zero_module(nn.Linear(self.inner_dim, self.inner_dim))
111
+ )
112
+
113
+ self.controlnet_x_embedder = zero_module(
114
+ torch.nn.Linear(in_channels + extra_condition_channels, self.inner_dim)
115
+ )
116
+
117
+ self.gradient_checkpointing = False
118
+
119
+ @property
120
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
121
+ def attn_processors(self):
122
+ r"""
123
+ Returns:
124
+ `dict` of attention processors: A dictionary containing all attention processors used in the model with
125
+ indexed by its weight name.
126
+ """
127
+ # set recursively
128
+ processors = {}
129
+
130
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
131
+ if hasattr(module, "get_processor"):
132
+ processors[f"{name}.processor"] = module.get_processor()
133
+
134
+ for sub_name, child in module.named_children():
135
+ fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
136
+
137
+ return processors
138
+
139
+ for name, module in self.named_children():
140
+ fn_recursive_add_processors(name, module, processors)
141
+
142
+ return processors
143
+
144
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
145
+ def set_attn_processor(self, processor):
146
+ r"""
147
+ Sets the attention processor to use to compute attention.
148
+
149
+ Parameters:
150
+ processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
151
+ The instantiated processor class or a dictionary of processor classes that will be set as the processor
152
+ for **all** `Attention` layers.
153
+
154
+ If `processor` is a dict, the key needs to define the path to the corresponding cross attention
155
+ processor. This is strongly recommended when setting trainable attention processors.
156
+
157
+ """
158
+ count = len(self.attn_processors.keys())
159
+
160
+ if isinstance(processor, dict) and len(processor) != count:
161
+ raise ValueError(
162
+ f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
163
+ f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
164
+ )
165
+
166
+ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
167
+ if hasattr(module, "set_processor"):
168
+ if not isinstance(processor, dict):
169
+ module.set_processor(processor)
170
+ else:
171
+ module.set_processor(processor.pop(f"{name}.processor"))
172
+
173
+ for sub_name, child in module.named_children():
174
+ fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
175
+
176
+ for name, module in self.named_children():
177
+ fn_recursive_attn_processor(name, module, processor)
178
+
179
+ def _set_gradient_checkpointing(self, module, value=False):
180
+ if hasattr(module, "gradient_checkpointing"):
181
+ module.gradient_checkpointing = value
182
+
183
+ @classmethod
184
+ def from_transformer(
185
+ cls,
186
+ transformer,
187
+ num_layers: int = 4,
188
+ num_single_layers: int = 10,
189
+ attention_head_dim: int = 128,
190
+ num_attention_heads: int = 24,
191
+ load_weights_from_transformer=True,
192
+ ):
193
+ config = transformer.config
194
+ config["num_layers"] = num_layers
195
+ config["num_single_layers"] = num_single_layers
196
+ config["attention_head_dim"] = attention_head_dim
197
+ config["num_attention_heads"] = num_attention_heads
198
+
199
+ controlnet = cls(**config)
200
+
201
+ if load_weights_from_transformer:
202
+ controlnet.pos_embed.load_state_dict(transformer.pos_embed.state_dict())
203
+ controlnet.time_text_embed.load_state_dict(
204
+ transformer.time_text_embed.state_dict()
205
+ )
206
+ controlnet.context_embedder.load_state_dict(
207
+ transformer.context_embedder.state_dict()
208
+ )
209
+ controlnet.x_embedder.load_state_dict(transformer.x_embedder.state_dict())
210
+ controlnet.transformer_blocks.load_state_dict(
211
+ transformer.transformer_blocks.state_dict(), strict=False
212
+ )
213
+ controlnet.single_transformer_blocks.load_state_dict(
214
+ transformer.single_transformer_blocks.state_dict(), strict=False
215
+ )
216
+
217
+ controlnet.controlnet_x_embedder = zero_module(
218
+ controlnet.controlnet_x_embedder
219
+ )
220
+
221
+ return controlnet
222
+
223
+ def forward(
224
+ self,
225
+ hidden_states: torch.Tensor,
226
+ controlnet_cond: torch.Tensor,
227
+ conditioning_scale: float = 1.0,
228
+ encoder_hidden_states: torch.Tensor = None,
229
+ pooled_projections: torch.Tensor = None,
230
+ timestep: torch.LongTensor = None,
231
+ img_ids: torch.Tensor = None,
232
+ txt_ids: torch.Tensor = None,
233
+ guidance: torch.Tensor = None,
234
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
235
+ return_dict: bool = True,
236
+ ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
237
+ """
238
+ The [`FluxTransformer2DModel`] forward method.
239
+
240
+ Args:
241
+ hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
242
+ Input `hidden_states`.
243
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
244
+ Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
245
+ pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
246
+ from the embeddings of input conditions.
247
+ timestep ( `torch.LongTensor`):
248
+ Used to indicate denoising step.
249
+ block_controlnet_hidden_states: (`list` of `torch.Tensor`):
250
+ A list of tensors that if specified are added to the residuals of transformer blocks.
251
+ joint_attention_kwargs (`dict`, *optional*):
252
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
253
+ `self.processor` in
254
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
255
+ return_dict (`bool`, *optional*, defaults to `True`):
256
+ Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
257
+ tuple.
258
+
259
+ Returns:
260
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
261
+ `tuple` where the first element is the sample tensor.
262
+ """
263
+ if joint_attention_kwargs is not None:
264
+ joint_attention_kwargs = joint_attention_kwargs.copy()
265
+ lora_scale = joint_attention_kwargs.pop("scale", 1.0)
266
+ else:
267
+ lora_scale = 1.0
268
+
269
+ if USE_PEFT_BACKEND:
270
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
271
+ scale_lora_layers(self, lora_scale)
272
+ else:
273
+ if (
274
+ joint_attention_kwargs is not None
275
+ and joint_attention_kwargs.get("scale", None) is not None
276
+ ):
277
+ logger.warning(
278
+ "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
279
+ )
280
+ hidden_states = self.x_embedder(hidden_states)
281
+
282
+ # add condition
283
+ hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond)
284
+
285
+ timestep = timestep.to(hidden_states.dtype) * 1000
286
+ if guidance is not None:
287
+ guidance = guidance.to(hidden_states.dtype) * 1000
288
+ else:
289
+ guidance = None
290
+ temb = (
291
+ self.time_text_embed(timestep, pooled_projections)
292
+ if guidance is None
293
+ else self.time_text_embed(timestep, guidance, pooled_projections)
294
+ )
295
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
296
+
297
+ txt_ids = txt_ids.expand(img_ids.size(0), -1, -1)
298
+ ids = torch.cat((txt_ids, img_ids), dim=1)
299
+ image_rotary_emb = self.pos_embed(ids[0])
300
+ # image_rotary_emb = torch.stack([self.pos_embed(id) for id in ids])
301
+
302
+ block_samples = ()
303
+ for _, block in enumerate(self.transformer_blocks):
304
+ if self.training and self.gradient_checkpointing:
305
+
306
+ def create_custom_forward(module, return_dict=None):
307
+ def custom_forward(*inputs):
308
+ if return_dict is not None:
309
+ return module(*inputs, return_dict=return_dict)
310
+ else:
311
+ return module(*inputs)
312
+
313
+ return custom_forward
314
+
315
+ ckpt_kwargs: Dict[str, Any] = (
316
+ {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
317
+ )
318
+ (
319
+ encoder_hidden_states,
320
+ hidden_states,
321
+ ) = torch.utils.checkpoint.checkpoint(
322
+ create_custom_forward(block),
323
+ hidden_states,
324
+ encoder_hidden_states,
325
+ temb,
326
+ image_rotary_emb,
327
+ **ckpt_kwargs,
328
+ )
329
+
330
+ else:
331
+ encoder_hidden_states, hidden_states = block(
332
+ hidden_states=hidden_states,
333
+ encoder_hidden_states=encoder_hidden_states,
334
+ temb=temb,
335
+ image_rotary_emb=image_rotary_emb,
336
+ )
337
+ block_samples = block_samples + (hidden_states,)
338
+
339
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
340
+
341
+ single_block_samples = ()
342
+ for _, block in enumerate(self.single_transformer_blocks):
343
+ if self.training and self.gradient_checkpointing:
344
+
345
+ def create_custom_forward(module, return_dict=None):
346
+ def custom_forward(*inputs):
347
+ if return_dict is not None:
348
+ return module(*inputs, return_dict=return_dict)
349
+ else:
350
+ return module(*inputs)
351
+
352
+ return custom_forward
353
+
354
+ ckpt_kwargs: Dict[str, Any] = (
355
+ {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
356
+ )
357
+ hidden_states = torch.utils.checkpoint.checkpoint(
358
+ create_custom_forward(block),
359
+ hidden_states,
360
+ temb,
361
+ image_rotary_emb,
362
+ **ckpt_kwargs,
363
+ )
364
+
365
+ else:
366
+ hidden_states = block(
367
+ hidden_states=hidden_states,
368
+ temb=temb,
369
+ image_rotary_emb=image_rotary_emb,
370
+ )
371
+ single_block_samples = single_block_samples + (
372
+ hidden_states[:, encoder_hidden_states.shape[1] :],
373
+ )
374
+
375
+ # controlnet block
376
+ controlnet_block_samples = ()
377
+ for block_sample, controlnet_block in zip(
378
+ block_samples, self.controlnet_blocks
379
+ ):
380
+ block_sample = controlnet_block(block_sample)
381
+ controlnet_block_samples = controlnet_block_samples + (block_sample,)
382
+
383
+ controlnet_single_block_samples = ()
384
+ for single_block_sample, controlnet_block in zip(
385
+ single_block_samples, self.controlnet_single_blocks
386
+ ):
387
+ single_block_sample = controlnet_block(single_block_sample)
388
+ controlnet_single_block_samples = controlnet_single_block_samples + (
389
+ single_block_sample,
390
+ )
391
+
392
+ # scaling
393
+ controlnet_block_samples = [
394
+ sample * conditioning_scale for sample in controlnet_block_samples
395
+ ]
396
+ controlnet_single_block_samples = [
397
+ sample * conditioning_scale for sample in controlnet_single_block_samples
398
+ ]
399
+
400
+ #
401
+ controlnet_block_samples = (
402
+ None if len(controlnet_block_samples) == 0 else controlnet_block_samples
403
+ )
404
+ controlnet_single_block_samples = (
405
+ None
406
+ if len(controlnet_single_block_samples) == 0
407
+ else controlnet_single_block_samples
408
+ )
409
+
410
+ if USE_PEFT_BACKEND:
411
+ # remove `lora_scale` from each PEFT layer
412
+ unscale_lora_layers(self, lora_scale)
413
+
414
+ if not return_dict:
415
+ return (controlnet_block_samples, controlnet_single_block_samples)
416
+
417
+ return FluxControlNetOutput(
418
+ controlnet_block_samples=controlnet_block_samples,
419
+ controlnet_single_block_samples=controlnet_single_block_samples,
420
+ )
diptych_prompting_inference.py ADDED
@@ -0,0 +1,377 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn.functional as F
3
+
4
+ from diffusers.utils import load_image, check_min_version
5
+ from controlnet_flux import FluxControlNetModel
6
+ from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline
7
+ import os
8
+ import numpy as np
9
+ from PIL import Image
10
+ import argparse
11
+
12
+ from diffusers.models.attention_processor import Attention
13
+
14
+ from dataclasses import dataclass
15
+ from typing import Any, List, Dict, Optional, Union, Tuple
16
+ import cv2
17
+ from transformers import AutoProcessor, pipeline, AutoModelForMaskGeneration
18
+
19
+ @dataclass
20
+ class BoundingBox:
21
+ xmin: int
22
+ ymin: int
23
+ xmax: int
24
+ ymax: int
25
+
26
+ @property
27
+ def xyxy(self) -> List[float]:
28
+ return [self.xmin, self.ymin, self.xmax, self.ymax]
29
+
30
+ @dataclass
31
+ class DetectionResult:
32
+ score: float
33
+ label: str
34
+ box: BoundingBox
35
+ mask: Optional[np.array] = None
36
+
37
+ @classmethod
38
+ def from_dict(cls, detection_dict: Dict) -> 'DetectionResult':
39
+ return cls(score=detection_dict['score'],
40
+ label=detection_dict['label'],
41
+ box=BoundingBox(xmin=detection_dict['box']['xmin'],
42
+ ymin=detection_dict['box']['ymin'],
43
+ xmax=detection_dict['box']['xmax'],
44
+ ymax=detection_dict['box']['ymax']))
45
+
46
+
47
+ def mask_to_polygon(mask: np.ndarray) -> List[List[int]]:
48
+ # Find contours in the binary mask
49
+ contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
50
+
51
+ # Find the contour with the largest area
52
+ largest_contour = max(contours, key=cv2.contourArea)
53
+
54
+ # Extract the vertices of the contour
55
+ polygon = largest_contour.reshape(-1, 2).tolist()
56
+
57
+ return polygon
58
+
59
+ def polygon_to_mask(polygon: List[Tuple[int, int]], image_shape: Tuple[int, int]) -> np.ndarray:
60
+ """
61
+ Convert a polygon to a segmentation mask.
62
+
63
+ Args:
64
+ - polygon (list): List of (x, y) coordinates representing the vertices of the polygon.
65
+ - image_shape (tuple): Shape of the image (height, width) for the mask.
66
+
67
+ Returns:
68
+ - np.ndarray: Segmentation mask with the polygon filled.
69
+ """
70
+ # Create an empty mask
71
+ mask = np.zeros(image_shape, dtype=np.uint8)
72
+
73
+ # Convert polygon to an array of points
74
+ pts = np.array(polygon, dtype=np.int32)
75
+
76
+ # Fill the polygon with white color (255)
77
+ cv2.fillPoly(mask, [pts], color=(255,))
78
+
79
+ return mask
80
+
81
+ def get_boxes(results: DetectionResult) -> List[List[List[float]]]:
82
+ boxes = []
83
+ for result in results:
84
+ xyxy = result.box.xyxy
85
+ boxes.append(xyxy)
86
+
87
+ return [boxes]
88
+
89
+ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]:
90
+ masks = masks.cpu().float()
91
+ masks = masks.permute(0, 2, 3, 1)
92
+ masks = masks.mean(axis=-1)
93
+ masks = (masks > 0).int()
94
+ masks = masks.numpy().astype(np.uint8)
95
+ masks = list(masks)
96
+
97
+ if polygon_refinement:
98
+ for idx, mask in enumerate(masks):
99
+ shape = mask.shape
100
+ polygon = mask_to_polygon(mask)
101
+ mask = polygon_to_mask(polygon, shape)
102
+ masks[idx] = mask
103
+
104
+ return masks
105
+
106
+ def detect(
107
+ object_detector,
108
+ image: Image.Image,
109
+ labels: List[str],
110
+ threshold: float = 0.3,
111
+ detector_id: Optional[str] = None
112
+ ) -> List[Dict[str, Any]]:
113
+ """
114
+ Use Grounding DINO to detect a set of labels in an image in a zero-shot fashion.
115
+ """
116
+ device = "cuda" if torch.cuda.is_available() else "cpu"
117
+ detector_id = detector_id if detector_id is not None else "IDEA-Research/grounding-dino-tiny"
118
+ # object_detector = detect_pipeline(model=detector_id, task="zero-shot-object-detection", device=device)
119
+
120
+ labels = [label if label.endswith(".") else label+"." for label in labels]
121
+
122
+ results = object_detector(image, candidate_labels=labels, threshold=threshold)
123
+ results = [DetectionResult.from_dict(result) for result in results]
124
+
125
+ return results
126
+
127
+ def segment(
128
+ segmentator,
129
+ processor,
130
+ image: Image.Image,
131
+ detection_results: List[Dict[str, Any]],
132
+ polygon_refinement: bool = False,
133
+ ) -> List[DetectionResult]:
134
+ """
135
+ Use Segment Anything (SAM) to generate masks given an image + a set of bounding boxes.
136
+ """
137
+ device = "cuda" if torch.cuda.is_available() else "cpu"
138
+
139
+ boxes = get_boxes(detection_results)
140
+ inputs = processor(images=image, input_boxes=boxes, return_tensors="pt").to(device)
141
+
142
+ outputs = segmentator(**inputs)
143
+ masks = processor.post_process_masks(
144
+ masks=outputs.pred_masks,
145
+ original_sizes=inputs.original_sizes,
146
+ reshaped_input_sizes=inputs.reshaped_input_sizes
147
+ )[0]
148
+
149
+ masks = refine_masks(masks, polygon_refinement)
150
+
151
+ for detection_result, mask in zip(detection_results, masks):
152
+ detection_result.mask = mask
153
+
154
+ return detection_results
155
+
156
+ def grounded_segmentation(
157
+ detect_pipeline,
158
+ segmentator,
159
+ segment_processor,
160
+ image: Union[Image.Image, str],
161
+ labels: List[str],
162
+ threshold: float = 0.3,
163
+ polygon_refinement: bool = False,
164
+ detector_id: Optional[str] = None,
165
+ segmenter_id: Optional[str] = None
166
+ ) -> Tuple[np.ndarray, List[DetectionResult]]:
167
+ if isinstance(image, str):
168
+ image = load_image(image)
169
+
170
+ detections = detect(detect_pipeline, image, labels, threshold, detector_id)
171
+ detections = segment(segmentator, segment_processor, image, detections, polygon_refinement)
172
+
173
+ return np.array(image), detections
174
+
175
+
176
+ class CustomFluxAttnProcessor2_0:
177
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
178
+
179
+ def __init__(self, height=44, width=88, attn_enforce=1.0):
180
+ if not hasattr(F, "scaled_dot_product_attention"):
181
+ raise ImportError("FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
182
+
183
+ self.height = height
184
+ self.width = width
185
+ self.num_pixels = height * width
186
+ self.step = 0
187
+ self.attn_enforce = attn_enforce
188
+
189
+ def __call__(
190
+ self,
191
+ attn: Attention,
192
+ hidden_states: torch.FloatTensor,
193
+ encoder_hidden_states: torch.FloatTensor = None,
194
+ attention_mask: Optional[torch.FloatTensor] = None,
195
+ image_rotary_emb: Optional[torch.Tensor] = None,
196
+ ) -> torch.FloatTensor:
197
+ self.step += 1
198
+ batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
199
+
200
+ # `sample` projections.
201
+ query = attn.to_q(hidden_states)
202
+ key = attn.to_k(hidden_states)
203
+ value = attn.to_v(hidden_states)
204
+
205
+ inner_dim = key.shape[-1]
206
+ head_dim = inner_dim // attn.heads
207
+
208
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
209
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
210
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
211
+
212
+ if attn.norm_q is not None:
213
+ query = attn.norm_q(query)
214
+ if attn.norm_k is not None:
215
+ key = attn.norm_k(key)
216
+
217
+ # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
218
+ if encoder_hidden_states is not None:
219
+ # `context` projections.
220
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
221
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
222
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
223
+
224
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
225
+ batch_size, -1, attn.heads, head_dim
226
+ ).transpose(1, 2)
227
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
228
+ batch_size, -1, attn.heads, head_dim
229
+ ).transpose(1, 2)
230
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
231
+ batch_size, -1, attn.heads, head_dim
232
+ ).transpose(1, 2)
233
+
234
+ if attn.norm_added_q is not None:
235
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
236
+ if attn.norm_added_k is not None:
237
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
238
+
239
+ # attention
240
+ query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
241
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
242
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
243
+
244
+ if image_rotary_emb is not None:
245
+ from diffusers.models.embeddings import apply_rotary_emb
246
+
247
+ query = apply_rotary_emb(query, image_rotary_emb)
248
+ key = apply_rotary_emb(key, image_rotary_emb)
249
+
250
+
251
+ ######### attn_enforce
252
+ if self.attn_enforce != 1.0:
253
+ attn_probs = (torch.einsum('bhqd,bhkd->bhqk', query, key) * attn.scale).softmax(dim=-1)
254
+ img_attn_probs = attn_probs[:, :, -self.num_pixels:, -self.num_pixels:]
255
+ img_attn_probs = img_attn_probs.reshape((batch_size, attn.heads, self.height, self.width, self.height, self.width))
256
+ img_attn_probs[:, :, :, self.width//2:, :, :self.width//2] *= self.attn_enforce
257
+ img_attn_probs = img_attn_probs.reshape((batch_size, attn.heads, self.num_pixels, self.num_pixels))
258
+ attn_probs[:, :, -self.num_pixels:, -self.num_pixels:] = img_attn_probs
259
+ hidden_states = torch.einsum('bhqk,bhkd->bhqd', attn_probs, value)
260
+ else:
261
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
262
+
263
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
264
+ hidden_states = hidden_states.to(query.dtype)
265
+
266
+ if encoder_hidden_states is not None:
267
+ encoder_hidden_states, hidden_states = (
268
+ hidden_states[:, : encoder_hidden_states.shape[1]],
269
+ hidden_states[:, encoder_hidden_states.shape[1] :],
270
+ )
271
+
272
+ # linear proj
273
+ hidden_states = attn.to_out[0](hidden_states)
274
+ # dropout
275
+ hidden_states = attn.to_out[1](hidden_states)
276
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
277
+
278
+ return hidden_states, encoder_hidden_states
279
+ else:
280
+ return hidden_states
281
+
282
+ if __name__ == '__main__':
283
+ parser = argparse.ArgumentParser()
284
+ parser.add_argument('--attn_enforce', type=float, default=1.3)
285
+ parser.add_argument('--ctrl_scale', type=float, default=0.95)
286
+ parser.add_argument('--width', type=int, default=768)
287
+ parser.add_argument('--height', type=int, default=768)
288
+ parser.add_argument('--pixel_offset', type=int, default=8)
289
+ parser.add_argument('--input_image_path', type=str, default='./assets/bear_plushie.jpg')
290
+ parser.add_argument('--subject_name', type=str, default='bear plushie')
291
+ parser.add_argument('--target_prompt', type=str, default='a photo of a bear plushie surfing on the beach')
292
+
293
+ args = parser.parse_args()
294
+
295
+ # Build pipeline
296
+ controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16)
297
+ pipe = FluxControlNetInpaintingPipeline.from_pretrained(
298
+ "black-forest-labs/FLUX.1-dev",
299
+ controlnet=controlnet,
300
+ torch_dtype=torch.bfloat16
301
+ ).to("cuda")
302
+ pipe.transformer.to(torch.bfloat16)
303
+ pipe.controlnet.to(torch.bfloat16)
304
+ base_attn_procs = pipe.transformer.attn_processors.copy()
305
+
306
+ detector_id = "IDEA-Research/grounding-dino-tiny"
307
+ segmenter_id = "facebook/sam-vit-base"
308
+
309
+ segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).cuda()
310
+ segment_processor = AutoProcessor.from_pretrained(segmenter_id)
311
+ object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device=torch.device("cuda"))
312
+
313
+ def segment_image(image, object_name):
314
+ image_array, detections = grounded_segmentation(
315
+ object_detector,
316
+ segmentator,
317
+ segment_processor,
318
+ image=image,
319
+ labels=object_name,
320
+ threshold=0.3,
321
+ polygon_refinement=True,
322
+ )
323
+ segment_result = image_array * np.expand_dims(detections[0].mask / 255, axis=-1) + np.ones_like(image_array) * (
324
+ 1 - np.expand_dims(detections[0].mask / 255, axis=-1)) * 255
325
+ segmented_image = Image.fromarray(segment_result.astype(np.uint8))
326
+ return segmented_image
327
+
328
+
329
+ def make_diptych(image):
330
+ ref_image = np.array(image)
331
+ ref_image = np.concatenate([ref_image, np.zeros_like(ref_image)], axis=1)
332
+ ref_image = Image.fromarray(ref_image)
333
+ return ref_image
334
+
335
+
336
+ # Load image and mask
337
+ width = args.width + args.pixel_offset * 2
338
+ height = args.height + args.pixel_offset * 2
339
+ size = (width*2, height)
340
+
341
+ subject_name = args.subject_name
342
+ base_prompt = f"a photo of {subject_name}"
343
+ target_prompt = args.target_prompt
344
+ diptych_text_prompt = f"A diptych with two side-by-side images of same {subject_name}. On the left, {base_prompt}. On the right, replicate this {subject_name} exactly but as {target_prompt}"
345
+
346
+ reference_image = load_image(args.input_image_path).resize((width, height)).convert("RGB")
347
+
348
+ ctrl_scale=args.ctrl_scale
349
+ segmented_image = segment_image(reference_image, subject_name)
350
+ mask_image = np.concatenate([np.zeros((height, width, 3)), np.ones((height, width, 3))*255], axis=1)
351
+ mask_image = Image.fromarray(mask_image.astype(np.uint8))
352
+ diptych_image_prompt = make_diptych(segmented_image)
353
+
354
+ new_attn_procs = base_attn_procs.copy()
355
+ for i, (k, v) in enumerate(new_attn_procs.items()):
356
+ new_attn_procs[k] = CustomFluxAttnProcessor2_0(height=height // 16, width=width // 16 * 2, attn_enforce=args.attn_enforce)
357
+ pipe.transformer.set_attn_processor(new_attn_procs)
358
+
359
+ generator = torch.Generator(device="cuda").manual_seed(42)
360
+ # Inpaint
361
+ result = pipe(
362
+ prompt=diptych_text_prompt,
363
+ height=size[1],
364
+ width=size[0],
365
+ control_image=diptych_image_prompt,
366
+ control_mask=mask_image,
367
+ num_inference_steps=30,
368
+ generator=generator,
369
+ controlnet_conditioning_scale=ctrl_scale,
370
+ guidance_scale=3.5,
371
+ negative_prompt="",
372
+ true_guidance_scale=3.5
373
+ ).images[0]
374
+
375
+ result = result.crop((width, 0, width*2, height))
376
+ result = result.crop((args.pixel_offset, args.pixel_offset, width-args.pixel_offset, height-args.pixel_offset))
377
+ result.save('result.png')
pipeline_flux_controlnet_inpaint.py ADDED
@@ -0,0 +1,1149 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import inspect
2
+ from typing import Any, Callable, Dict, List, Optional, Union
3
+
4
+ import numpy as np
5
+ import torch
6
+ from transformers import (
7
+ CLIPTextModel,
8
+ CLIPTokenizer,
9
+ T5EncoderModel,
10
+ T5TokenizerFast,
11
+ )
12
+
13
+ from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
14
+ from diffusers.loaders import FluxLoraLoaderMixin
15
+ from diffusers.models.autoencoders import AutoencoderKL
16
+
17
+ from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
18
+ from diffusers.utils import (
19
+ USE_PEFT_BACKEND,
20
+ is_torch_xla_available,
21
+ logging,
22
+ replace_example_docstring,
23
+ scale_lora_layers,
24
+ unscale_lora_layers,
25
+ )
26
+ from diffusers.utils.torch_utils import randn_tensor
27
+ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
28
+ from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
29
+
30
+ from transformer_flux import FluxTransformer2DModel
31
+ from controlnet_flux import FluxControlNetModel
32
+
33
+ if is_torch_xla_available():
34
+ import torch_xla.core.xla_model as xm
35
+
36
+ XLA_AVAILABLE = True
37
+ else:
38
+ XLA_AVAILABLE = False
39
+
40
+
41
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
42
+
43
+ EXAMPLE_DOC_STRING = """
44
+ Examples:
45
+ ```py
46
+ >>> import torch
47
+ >>> from diffusers.utils import load_image
48
+ >>> from diffusers import FluxControlNetPipeline
49
+ >>> from diffusers import FluxControlNetModel
50
+
51
+ >>> controlnet_model = "InstantX/FLUX.1-dev-controlnet-canny-alpha"
52
+ >>> controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
53
+ >>> pipe = FluxControlNetPipeline.from_pretrained(
54
+ ... base_model, controlnet=controlnet, torch_dtype=torch.bfloat16
55
+ ... )
56
+ >>> pipe.to("cuda")
57
+ >>> control_image = load_image("https://huggingface.co/InstantX/SD3-Controlnet-Canny/resolve/main/canny.jpg")
58
+ >>> control_mask = load_image("https://huggingface.co/InstantX/SD3-Controlnet-Canny/resolve/main/canny.jpg")
59
+ >>> prompt = "A girl in city, 25 years old, cool, futuristic"
60
+ >>> image = pipe(
61
+ ... prompt,
62
+ ... control_image=control_image,
63
+ ... controlnet_conditioning_scale=0.6,
64
+ ... num_inference_steps=28,
65
+ ... guidance_scale=3.5,
66
+ ... ).images[0]
67
+ >>> image.save("flux.png")
68
+ ```
69
+ """
70
+
71
+
72
+ # Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
73
+ def calculate_shift(
74
+ image_seq_len,
75
+ base_seq_len: int = 256,
76
+ max_seq_len: int = 4096,
77
+ base_shift: float = 0.5,
78
+ max_shift: float = 1.16,
79
+ ):
80
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
81
+ b = base_shift - m * base_seq_len
82
+ mu = image_seq_len * m + b
83
+ return mu
84
+
85
+
86
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
87
+ def retrieve_timesteps(
88
+ scheduler,
89
+ num_inference_steps: Optional[int] = None,
90
+ device: Optional[Union[str, torch.device]] = None,
91
+ timesteps: Optional[List[int]] = None,
92
+ sigmas: Optional[List[float]] = None,
93
+ **kwargs,
94
+ ):
95
+ """
96
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
97
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
98
+
99
+ Args:
100
+ scheduler (`SchedulerMixin`):
101
+ The scheduler to get timesteps from.
102
+ num_inference_steps (`int`):
103
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
104
+ must be `None`.
105
+ device (`str` or `torch.device`, *optional*):
106
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
107
+ timesteps (`List[int]`, *optional*):
108
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
109
+ `num_inference_steps` and `sigmas` must be `None`.
110
+ sigmas (`List[float]`, *optional*):
111
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
112
+ `num_inference_steps` and `timesteps` must be `None`.
113
+
114
+ Returns:
115
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
116
+ second element is the number of inference steps.
117
+ """
118
+ if timesteps is not None and sigmas is not None:
119
+ raise ValueError(
120
+ "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
121
+ )
122
+ if timesteps is not None:
123
+ accepts_timesteps = "timesteps" in set(
124
+ inspect.signature(scheduler.set_timesteps).parameters.keys()
125
+ )
126
+ if not accepts_timesteps:
127
+ raise ValueError(
128
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
129
+ f" timestep schedules. Please check whether you are using the correct scheduler."
130
+ )
131
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
132
+ timesteps = scheduler.timesteps
133
+ num_inference_steps = len(timesteps)
134
+ elif sigmas is not None:
135
+ accept_sigmas = "sigmas" in set(
136
+ inspect.signature(scheduler.set_timesteps).parameters.keys()
137
+ )
138
+ if not accept_sigmas:
139
+ raise ValueError(
140
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
141
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
142
+ )
143
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
144
+ timesteps = scheduler.timesteps
145
+ num_inference_steps = len(timesteps)
146
+ else:
147
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
148
+ timesteps = scheduler.timesteps
149
+ return timesteps, num_inference_steps
150
+
151
+
152
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
153
+ def retrieve_latents(
154
+ encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
155
+ ):
156
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
157
+ return encoder_output.latent_dist.sample(generator)
158
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
159
+ return encoder_output.latent_dist.mode()
160
+ elif hasattr(encoder_output, "latents"):
161
+ return encoder_output.latents
162
+ else:
163
+ raise AttributeError("Could not access latents of provided encoder_output")
164
+
165
+
166
+ class FluxControlNetInpaintingPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
167
+ r"""
168
+ The Flux pipeline for text-to-image generation.
169
+
170
+ Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
171
+
172
+ Args:
173
+ transformer ([`FluxTransformer2DModel`]):
174
+ Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
175
+ scheduler ([`FlowMatchEulerDiscreteScheduler`]):
176
+ A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
177
+ vae ([`AutoencoderKL`]):
178
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
179
+ text_encoder ([`CLIPTextModel`]):
180
+ [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
181
+ the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
182
+ text_encoder_2 ([`T5EncoderModel`]):
183
+ [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
184
+ the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
185
+ tokenizer (`CLIPTokenizer`):
186
+ Tokenizer of class
187
+ [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
188
+ tokenizer_2 (`T5TokenizerFast`):
189
+ Second Tokenizer of class
190
+ [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
191
+ """
192
+
193
+ model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
194
+ _optional_components = []
195
+ _callback_tensor_inputs = ["latents", "prompt_embeds"]
196
+
197
+ def __init__(
198
+ self,
199
+ scheduler: FlowMatchEulerDiscreteScheduler,
200
+ vae: AutoencoderKL,
201
+ text_encoder: CLIPTextModel,
202
+ tokenizer: CLIPTokenizer,
203
+ text_encoder_2: T5EncoderModel,
204
+ tokenizer_2: T5TokenizerFast,
205
+ transformer: FluxTransformer2DModel,
206
+ controlnet: FluxControlNetModel,
207
+ ):
208
+ super().__init__()
209
+
210
+ self.register_modules(
211
+ vae=vae,
212
+ text_encoder=text_encoder,
213
+ text_encoder_2=text_encoder_2,
214
+ tokenizer=tokenizer,
215
+ tokenizer_2=tokenizer_2,
216
+ transformer=transformer,
217
+ scheduler=scheduler,
218
+ controlnet=controlnet,
219
+ )
220
+ self.vae_scale_factor = (
221
+ 2 ** (len(self.vae.config.block_out_channels))
222
+ if hasattr(self, "vae") and self.vae is not None
223
+ else 16
224
+ )
225
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_resize=True, do_convert_rgb=True, do_normalize=True)
226
+ self.mask_processor = VaeImageProcessor(
227
+ vae_scale_factor=self.vae_scale_factor,
228
+ do_resize=True,
229
+ do_convert_grayscale=True,
230
+ do_normalize=False,
231
+ do_binarize=True,
232
+ )
233
+ self.tokenizer_max_length = (
234
+ self.tokenizer.model_max_length
235
+ if hasattr(self, "tokenizer") and self.tokenizer is not None
236
+ else 77
237
+ )
238
+ self.default_sample_size = 64
239
+
240
+ @property
241
+ def do_classifier_free_guidance(self):
242
+ return self._guidance_scale > 1
243
+
244
+ def _get_t5_prompt_embeds(
245
+ self,
246
+ prompt: Union[str, List[str]] = None,
247
+ num_images_per_prompt: int = 1,
248
+ max_sequence_length: int = 512,
249
+ device: Optional[torch.device] = None,
250
+ dtype: Optional[torch.dtype] = None,
251
+ ):
252
+ device = device or self._execution_device
253
+ dtype = dtype or self.text_encoder.dtype
254
+
255
+ prompt = [prompt] if isinstance(prompt, str) else prompt
256
+ batch_size = len(prompt)
257
+
258
+ text_inputs = self.tokenizer_2(
259
+ prompt,
260
+ padding="max_length",
261
+ max_length=max_sequence_length,
262
+ truncation=True,
263
+ return_length=False,
264
+ return_overflowing_tokens=False,
265
+ return_tensors="pt",
266
+ )
267
+ text_input_ids = text_inputs.input_ids
268
+ untruncated_ids = self.tokenizer_2(
269
+ prompt, padding="longest", return_tensors="pt"
270
+ ).input_ids
271
+
272
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
273
+ text_input_ids, untruncated_ids
274
+ ):
275
+ removed_text = self.tokenizer_2.batch_decode(
276
+ untruncated_ids[:, self.tokenizer_max_length - 1 : -1]
277
+ )
278
+ logger.warning(
279
+ "The following part of your input was truncated because `max_sequence_length` is set to "
280
+ f" {max_sequence_length} tokens: {removed_text}"
281
+ )
282
+
283
+ prompt_embeds = self.text_encoder_2(
284
+ text_input_ids.to(device), output_hidden_states=False
285
+ )[0]
286
+
287
+ dtype = self.text_encoder_2.dtype
288
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
289
+
290
+ _, seq_len, _ = prompt_embeds.shape
291
+
292
+ # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
293
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
294
+ prompt_embeds = prompt_embeds.view(
295
+ batch_size * num_images_per_prompt, seq_len, -1
296
+ )
297
+
298
+ return prompt_embeds
299
+
300
+ def _get_clip_prompt_embeds(
301
+ self,
302
+ prompt: Union[str, List[str]],
303
+ num_images_per_prompt: int = 1,
304
+ device: Optional[torch.device] = None,
305
+ ):
306
+ device = device or self._execution_device
307
+
308
+ prompt = [prompt] if isinstance(prompt, str) else prompt
309
+ batch_size = len(prompt)
310
+
311
+ text_inputs = self.tokenizer(
312
+ prompt,
313
+ padding="max_length",
314
+ max_length=self.tokenizer_max_length,
315
+ truncation=True,
316
+ return_overflowing_tokens=False,
317
+ return_length=False,
318
+ return_tensors="pt",
319
+ )
320
+
321
+ text_input_ids = text_inputs.input_ids
322
+ untruncated_ids = self.tokenizer(
323
+ prompt, padding="longest", return_tensors="pt"
324
+ ).input_ids
325
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
326
+ text_input_ids, untruncated_ids
327
+ ):
328
+ removed_text = self.tokenizer.batch_decode(
329
+ untruncated_ids[:, self.tokenizer_max_length - 1 : -1]
330
+ )
331
+ logger.warning(
332
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
333
+ f" {self.tokenizer_max_length} tokens: {removed_text}"
334
+ )
335
+ prompt_embeds = self.text_encoder(
336
+ text_input_ids.to(device), output_hidden_states=False
337
+ )
338
+
339
+ # Use pooled output of CLIPTextModel
340
+ prompt_embeds = prompt_embeds.pooler_output
341
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
342
+
343
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
344
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
345
+ prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
346
+
347
+ return prompt_embeds
348
+
349
+ def encode_prompt(
350
+ self,
351
+ prompt: Union[str, List[str]],
352
+ prompt_2: Union[str, List[str]],
353
+ device: Optional[torch.device] = None,
354
+ num_images_per_prompt: int = 1,
355
+ do_classifier_free_guidance: bool = True,
356
+ negative_prompt: Optional[Union[str, List[str]]] = None,
357
+ negative_prompt_2: Optional[Union[str, List[str]]] = None,
358
+ prompt_embeds: Optional[torch.FloatTensor] = None,
359
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
360
+ max_sequence_length: int = 512,
361
+ lora_scale: Optional[float] = None,
362
+ ):
363
+ r"""
364
+
365
+ Args:
366
+ prompt (`str` or `List[str]`, *optional*):
367
+ prompt to be encoded
368
+ prompt_2 (`str` or `List[str]`, *optional*):
369
+ The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
370
+ used in all text-encoders
371
+ device: (`torch.device`):
372
+ torch device
373
+ num_images_per_prompt (`int`):
374
+ number of images that should be generated per prompt
375
+ do_classifier_free_guidance (`bool`):
376
+ whether to use classifier-free guidance or not
377
+ negative_prompt (`str` or `List[str]`, *optional*):
378
+ negative prompt to be encoded
379
+ negative_prompt_2 (`str` or `List[str]`, *optional*):
380
+ negative prompt to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is
381
+ used in all text-encoders
382
+ prompt_embeds (`torch.FloatTensor`, *optional*):
383
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
384
+ provided, text embeddings will be generated from `prompt` input argument.
385
+ pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
386
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
387
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
388
+ clip_skip (`int`, *optional*):
389
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
390
+ the output of the pre-final layer will be used for computing the prompt embeddings.
391
+ lora_scale (`float`, *optional*):
392
+ A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
393
+ """
394
+ device = device or self._execution_device
395
+
396
+ # set lora scale so that monkey patched LoRA
397
+ # function of text encoder can correctly access it
398
+ if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
399
+ self._lora_scale = lora_scale
400
+
401
+ # dynamically adjust the LoRA scale
402
+ if self.text_encoder is not None and USE_PEFT_BACKEND:
403
+ scale_lora_layers(self.text_encoder, lora_scale)
404
+ if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
405
+ scale_lora_layers(self.text_encoder_2, lora_scale)
406
+
407
+ prompt = [prompt] if isinstance(prompt, str) else prompt
408
+ if prompt is not None:
409
+ batch_size = len(prompt)
410
+ else:
411
+ batch_size = prompt_embeds.shape[0]
412
+
413
+ if prompt_embeds is None:
414
+ prompt_2 = prompt_2 or prompt
415
+ prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
416
+
417
+ # We only use the pooled prompt output from the CLIPTextModel
418
+ pooled_prompt_embeds = self._get_clip_prompt_embeds(
419
+ prompt=prompt,
420
+ device=device,
421
+ num_images_per_prompt=num_images_per_prompt,
422
+ )
423
+ prompt_embeds = self._get_t5_prompt_embeds(
424
+ prompt=prompt_2,
425
+ num_images_per_prompt=num_images_per_prompt,
426
+ max_sequence_length=max_sequence_length,
427
+ device=device,
428
+ )
429
+
430
+ if do_classifier_free_guidance:
431
+ # 处理 negative prompt
432
+ negative_prompt = negative_prompt or ""
433
+ negative_prompt_2 = negative_prompt_2 or negative_prompt
434
+
435
+ negative_pooled_prompt_embeds = self._get_clip_prompt_embeds(
436
+ negative_prompt,
437
+ device=device,
438
+ num_images_per_prompt=num_images_per_prompt,
439
+ ).repeat_interleave(batch_size, dim=0)
440
+ negative_prompt_embeds = self._get_t5_prompt_embeds(
441
+ negative_prompt_2,
442
+ num_images_per_prompt=num_images_per_prompt,
443
+ max_sequence_length=max_sequence_length,
444
+ device=device,
445
+ ).repeat_interleave(batch_size, dim=0)
446
+ else:
447
+ negative_pooled_prompt_embeds = None
448
+ negative_prompt_embeds = None
449
+
450
+ if self.text_encoder is not None:
451
+ if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
452
+ # Retrieve the original scale by scaling back the LoRA layers
453
+ unscale_lora_layers(self.text_encoder, lora_scale)
454
+
455
+ if self.text_encoder_2 is not None:
456
+ if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
457
+ # Retrieve the original scale by scaling back the LoRA layers
458
+ unscale_lora_layers(self.text_encoder_2, lora_scale)
459
+
460
+ text_ids = torch.zeros(batch_size, prompt_embeds.shape[1], 3).to(
461
+ device=device, dtype=self.text_encoder.dtype
462
+ )
463
+
464
+ return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds,text_ids
465
+
466
+ def check_inputs(
467
+ self,
468
+ prompt,
469
+ prompt_2,
470
+ height,
471
+ width,
472
+ prompt_embeds=None,
473
+ pooled_prompt_embeds=None,
474
+ callback_on_step_end_tensor_inputs=None,
475
+ max_sequence_length=None,
476
+ ):
477
+ if height % 8 != 0 or width % 8 != 0:
478
+ raise ValueError(
479
+ f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
480
+ )
481
+
482
+ if callback_on_step_end_tensor_inputs is not None and not all(
483
+ k in self._callback_tensor_inputs
484
+ for k in callback_on_step_end_tensor_inputs
485
+ ):
486
+ raise ValueError(
487
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
488
+ )
489
+
490
+ if prompt is not None and prompt_embeds is not None:
491
+ raise ValueError(
492
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
493
+ " only forward one of the two."
494
+ )
495
+ elif prompt_2 is not None and prompt_embeds is not None:
496
+ raise ValueError(
497
+ f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
498
+ " only forward one of the two."
499
+ )
500
+ elif prompt is None and prompt_embeds is None:
501
+ raise ValueError(
502
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
503
+ )
504
+ elif prompt is not None and (
505
+ not isinstance(prompt, str) and not isinstance(prompt, list)
506
+ ):
507
+ raise ValueError(
508
+ f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
509
+ )
510
+ elif prompt_2 is not None and (
511
+ not isinstance(prompt_2, str) and not isinstance(prompt_2, list)
512
+ ):
513
+ raise ValueError(
514
+ f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}"
515
+ )
516
+
517
+ if prompt_embeds is not None and pooled_prompt_embeds is None:
518
+ raise ValueError(
519
+ "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
520
+ )
521
+
522
+ if max_sequence_length is not None and max_sequence_length > 512:
523
+ raise ValueError(
524
+ f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}"
525
+ )
526
+
527
+ # Copied from diffusers.pipelines.flux.pipeline_flux._prepare_latent_image_ids
528
+ @staticmethod
529
+ def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
530
+ latent_image_ids = torch.zeros(height // 2, width // 2, 3)
531
+ latent_image_ids[..., 1] = (
532
+ latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
533
+ )
534
+ latent_image_ids[..., 2] = (
535
+ latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
536
+ )
537
+
538
+ (
539
+ latent_image_id_height,
540
+ latent_image_id_width,
541
+ latent_image_id_channels,
542
+ ) = latent_image_ids.shape
543
+
544
+ latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
545
+ latent_image_ids = latent_image_ids.reshape(
546
+ batch_size,
547
+ latent_image_id_height * latent_image_id_width,
548
+ latent_image_id_channels,
549
+ )
550
+
551
+ return latent_image_ids.to(device=device, dtype=dtype)
552
+
553
+ # Copied from diffusers.pipelines.flux.pipeline_flux._pack_latents
554
+ @staticmethod
555
+ def _pack_latents(latents, batch_size, num_channels_latents, height, width):
556
+ latents = latents.view(
557
+ batch_size, num_channels_latents, height // 2, 2, width // 2, 2
558
+ )
559
+ latents = latents.permute(0, 2, 4, 1, 3, 5)
560
+ latents = latents.reshape(
561
+ batch_size, (height // 2) * (width // 2), num_channels_latents * 4
562
+ )
563
+
564
+ return latents
565
+
566
+ # Copied from diffusers.pipelines.flux.pipeline_flux._unpack_latents
567
+ @staticmethod
568
+ def _unpack_latents(latents, height, width, vae_scale_factor):
569
+ batch_size, num_patches, channels = latents.shape
570
+
571
+ height = height // vae_scale_factor
572
+ width = width // vae_scale_factor
573
+
574
+ latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
575
+ latents = latents.permute(0, 3, 1, 4, 2, 5)
576
+
577
+ latents = latents.reshape(
578
+ batch_size, channels // (2 * 2), height * 2, width * 2
579
+ )
580
+
581
+ return latents
582
+
583
+ # Copied from diffusers.pipelines.flux.pipeline_flux.prepare_latents
584
+ def prepare_latents(
585
+ self,
586
+ batch_size,
587
+ num_channels_latents,
588
+ height,
589
+ width,
590
+ dtype,
591
+ device,
592
+ generator,
593
+ latents=None,
594
+ ):
595
+ height = 2 * (int(height) // self.vae_scale_factor)
596
+ width = 2 * (int(width) // self.vae_scale_factor)
597
+
598
+ shape = (batch_size, num_channels_latents, height, width)
599
+
600
+ if latents is not None:
601
+ latent_image_ids = self._prepare_latent_image_ids(
602
+ batch_size, height, width, device, dtype
603
+ )
604
+ return latents.to(device=device, dtype=dtype), latent_image_ids
605
+
606
+ if isinstance(generator, list) and len(generator) != batch_size:
607
+ raise ValueError(
608
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
609
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
610
+ )
611
+
612
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
613
+ latents = self._pack_latents(
614
+ latents, batch_size, num_channels_latents, height, width
615
+ )
616
+
617
+ latent_image_ids = self._prepare_latent_image_ids(
618
+ batch_size, height, width, device, dtype
619
+ )
620
+
621
+ return latents, latent_image_ids
622
+ # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_inpaint.StableDiffusion3InpaintPipeline._encode_vae_image
623
+
624
+ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
625
+ if isinstance(generator, list):
626
+ image_latents = [
627
+ retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
628
+ for i in range(image.shape[0])
629
+ ]
630
+ image_latents = torch.cat(image_latents, dim=0)
631
+ else:
632
+ image_latents = retrieve_latents(self.vae.encode(image), generator=generator)
633
+
634
+ image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
635
+
636
+ return image_latents
637
+
638
+ def prepare_latents_with_init_image(
639
+ self,
640
+ image,
641
+ timestep,
642
+ batch_size,
643
+ num_channels_latents,
644
+ height,
645
+ width,
646
+ dtype,
647
+ device,
648
+ generator,
649
+ latents=None,
650
+ ):
651
+ if isinstance(generator, list) and len(generator) != batch_size:
652
+ raise ValueError(
653
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
654
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
655
+ )
656
+
657
+ height = 2 * (int(height) // self.vae_scale_factor)
658
+ width = 2 * (int(width) // self.vae_scale_factor)
659
+
660
+ shape = (batch_size, num_channels_latents, height, width)
661
+ latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype)
662
+
663
+ if latents is not None:
664
+ return latents.to(device=device, dtype=dtype), latent_image_ids
665
+
666
+ image = image.to(device=device, dtype=dtype)
667
+ image_latents = self._encode_vae_image(image=image, generator=generator)
668
+ if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
669
+ # expand init_latents for batch_size
670
+ additional_image_per_prompt = batch_size // image_latents.shape[0]
671
+ image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
672
+ elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
673
+ raise ValueError(
674
+ f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
675
+ )
676
+ else:
677
+ image_latents = torch.cat([image_latents], dim=0)
678
+
679
+ noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
680
+ latents = self.scheduler.scale_noise(image_latents, timestep, noise)
681
+ latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
682
+ return latents, latent_image_ids
683
+
684
+ # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image
685
+ def prepare_image(
686
+ self,
687
+ image,
688
+ width,
689
+ height,
690
+ batch_size,
691
+ num_images_per_prompt,
692
+ device,
693
+ dtype,
694
+ ):
695
+ if isinstance(image, torch.Tensor):
696
+ pass
697
+ else:
698
+ image = self.image_processor.preprocess(image, height=height, width=width)
699
+
700
+ image_batch_size = image.shape[0]
701
+
702
+ if image_batch_size == 1:
703
+ repeat_by = batch_size
704
+ else:
705
+ # image batch size is the same as prompt batch size
706
+ repeat_by = num_images_per_prompt
707
+
708
+ image = image.repeat_interleave(repeat_by, dim=0)
709
+
710
+ image = image.to(device=device, dtype=dtype)
711
+
712
+ return image
713
+
714
+ def prepare_image_with_mask(
715
+ self,
716
+ image,
717
+ mask,
718
+ width,
719
+ height,
720
+ batch_size,
721
+ num_images_per_prompt,
722
+ device,
723
+ dtype,
724
+ do_classifier_free_guidance = False,
725
+ ):
726
+ # Prepare image
727
+ if isinstance(image, torch.Tensor):
728
+ pass
729
+ else:
730
+ image = self.image_processor.preprocess(image, height=height, width=width)
731
+
732
+ image_batch_size = image.shape[0]
733
+ if image_batch_size == 1:
734
+ repeat_by = batch_size
735
+ else:
736
+ # image batch size is the same as prompt batch size
737
+ repeat_by = num_images_per_prompt
738
+ image = image.repeat_interleave(repeat_by, dim=0)
739
+ image = image.to(device=device, dtype=dtype)
740
+
741
+ # Prepare mask
742
+ if isinstance(mask, torch.Tensor):
743
+ pass
744
+ else:
745
+ mask = self.mask_processor.preprocess(mask, height=height, width=width)
746
+ mask = mask.repeat_interleave(repeat_by, dim=0)
747
+ mask = mask.to(device=device, dtype=dtype)
748
+
749
+ # Get masked image
750
+ masked_image = image.clone()
751
+ masked_image[(mask > 0.5).repeat(1, 3, 1, 1)] = -1
752
+
753
+ # Encode to latents
754
+ image_latents = self.vae.encode(masked_image.to(self.vae.dtype)).latent_dist.sample()
755
+ image_latents = (
756
+ image_latents - self.vae.config.shift_factor
757
+ ) * self.vae.config.scaling_factor
758
+ image_latents = image_latents.to(dtype)
759
+
760
+ mask = torch.nn.functional.interpolate(
761
+ mask, size=(height // self.vae_scale_factor * 2, width // self.vae_scale_factor * 2)
762
+ )
763
+ mask = 1 - mask
764
+
765
+ control_image = torch.cat([image_latents, mask], dim=1)
766
+
767
+ # Pack cond latents
768
+ packed_control_image = self._pack_latents(
769
+ control_image,
770
+ batch_size * num_images_per_prompt,
771
+ control_image.shape[1],
772
+ control_image.shape[2],
773
+ control_image.shape[3],
774
+ )
775
+
776
+ packed_image_latents = self._pack_latents(
777
+ image_latents,
778
+ batch_size * num_images_per_prompt,
779
+ image_latents.shape[1],
780
+ image_latents.shape[2],
781
+ image_latents.shape[3],
782
+ )
783
+ packed_mask = self._pack_latents(
784
+ mask.repeat_interleave(image_latents.shape[1], dim=1),
785
+ batch_size * num_images_per_prompt,
786
+ image_latents.shape[1],
787
+ mask.shape[2],
788
+ mask.shape[3],
789
+ )
790
+ if do_classifier_free_guidance:
791
+ packed_control_image = torch.cat([packed_control_image] * 2)
792
+
793
+ return packed_control_image, height, width, packed_image_latents, packed_mask
794
+
795
+ @property
796
+ def guidance_scale(self):
797
+ return self._guidance_scale
798
+
799
+ @property
800
+ def joint_attention_kwargs(self):
801
+ return self._joint_attention_kwargs
802
+
803
+ @property
804
+ def num_timesteps(self):
805
+ return self._num_timesteps
806
+
807
+ @property
808
+ def interrupt(self):
809
+ return self._interrupt
810
+
811
+ @torch.no_grad()
812
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
813
+ def __call__(
814
+ self,
815
+ prompt: Union[str, List[str]] = None,
816
+ prompt_2: Optional[Union[str, List[str]]] = None,
817
+ height: Optional[int] = None,
818
+ width: Optional[int] = None,
819
+ num_inference_steps: int = 28,
820
+ timesteps: List[int] = None,
821
+ guidance_scale: float = 7.0,
822
+ true_guidance_scale: float = 3.5 ,
823
+ negative_prompt: Optional[Union[str, List[str]]] = None,
824
+ negative_prompt_2: Optional[Union[str, List[str]]] = None,
825
+ control_image: PipelineImageInput = None,
826
+ control_mask: PipelineImageInput = None,
827
+ controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
828
+ num_images_per_prompt: Optional[int] = 1,
829
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
830
+ latents: Optional[torch.FloatTensor] = None,
831
+ prompt_embeds: Optional[torch.FloatTensor] = None,
832
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
833
+ output_type: Optional[str] = "pil",
834
+ return_dict: bool = True,
835
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
836
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
837
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
838
+ max_sequence_length: int = 512,
839
+ ):
840
+ r"""
841
+ Function invoked when calling the pipeline for generation.
842
+
843
+ Args:
844
+ prompt (`str` or `List[str]`, *optional*):
845
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
846
+ instead.
847
+ prompt_2 (`str` or `List[str]`, *optional*):
848
+ The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
849
+ will be used instead
850
+ height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
851
+ The height in pixels of the generated image. This is set to 1024 by default for the best results.
852
+ width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
853
+ The width in pixels of the generated image. This is set to 1024 by default for the best results.
854
+ num_inference_steps (`int`, *optional*, defaults to 50):
855
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
856
+ expense of slower inference.
857
+ timesteps (`List[int]`, *optional*):
858
+ Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
859
+ in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
860
+ passed will be used. Must be in descending order.
861
+ guidance_scale (`float`, *optional*, defaults to 7.0):
862
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
863
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
864
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
865
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
866
+ usually at the expense of lower image quality.
867
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
868
+ The number of images to generate per prompt.
869
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
870
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
871
+ to make generation deterministic.
872
+ latents (`torch.FloatTensor`, *optional*):
873
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
874
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
875
+ tensor will ge generated by sampling using the supplied random `generator`.
876
+ prompt_embeds (`torch.FloatTensor`, *optional*):
877
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
878
+ provided, text embeddings will be generated from `prompt` input argument.
879
+ pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
880
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
881
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
882
+ output_type (`str`, *optional*, defaults to `"pil"`):
883
+ The output format of the generate image. Choose between
884
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
885
+ return_dict (`bool`, *optional*, defaults to `True`):
886
+ Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
887
+ joint_attention_kwargs (`dict`, *optional*):
888
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
889
+ `self.processor` in
890
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
891
+ callback_on_step_end (`Callable`, *optional*):
892
+ A function that calls at the end of each denoising steps during the inference. The function is called
893
+ with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
894
+ callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
895
+ `callback_on_step_end_tensor_inputs`.
896
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
897
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
898
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
899
+ `._callback_tensor_inputs` attribute of your pipeline class.
900
+ max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
901
+
902
+ Examples:
903
+
904
+ Returns:
905
+ [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
906
+ is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
907
+ images.
908
+ """
909
+
910
+ height = height or self.default_sample_size * self.vae_scale_factor
911
+ width = width or self.default_sample_size * self.vae_scale_factor
912
+
913
+ # 1. Check inputs. Raise error if not correct
914
+ self.check_inputs(
915
+ prompt,
916
+ prompt_2,
917
+ height,
918
+ width,
919
+ prompt_embeds=prompt_embeds,
920
+ pooled_prompt_embeds=pooled_prompt_embeds,
921
+ callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
922
+ max_sequence_length=max_sequence_length,
923
+ )
924
+
925
+ self._guidance_scale = true_guidance_scale
926
+ self._joint_attention_kwargs = joint_attention_kwargs
927
+ self._interrupt = False
928
+
929
+ # 2. Define call parameters
930
+ if prompt is not None and isinstance(prompt, str):
931
+ batch_size = 1
932
+ elif prompt is not None and isinstance(prompt, list):
933
+ batch_size = len(prompt)
934
+ else:
935
+ batch_size = prompt_embeds.shape[0]
936
+
937
+ device = self._execution_device
938
+ dtype = self.transformer.dtype
939
+
940
+ lora_scale = (
941
+ self.joint_attention_kwargs.get("scale", None)
942
+ if self.joint_attention_kwargs is not None
943
+ else None
944
+ )
945
+ (
946
+ prompt_embeds,
947
+ pooled_prompt_embeds,
948
+ negative_prompt_embeds,
949
+ negative_pooled_prompt_embeds,
950
+ text_ids
951
+ ) = self.encode_prompt(
952
+ prompt=prompt,
953
+ prompt_2=prompt_2,
954
+ prompt_embeds=prompt_embeds,
955
+ pooled_prompt_embeds=pooled_prompt_embeds,
956
+ do_classifier_free_guidance = self.do_classifier_free_guidance,
957
+ negative_prompt = negative_prompt,
958
+ negative_prompt_2 = negative_prompt_2,
959
+ device=device,
960
+ num_images_per_prompt=num_images_per_prompt,
961
+ max_sequence_length=max_sequence_length,
962
+ lora_scale=lora_scale,
963
+ )
964
+
965
+ # 在 encode_prompt 之后
966
+ if self.do_classifier_free_guidance:
967
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim = 0)
968
+ pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim = 0)
969
+ text_ids = torch.cat([text_ids, text_ids], dim = 0)
970
+
971
+ # 3. Prepare control image
972
+ num_channels_latents = self.transformer.config.in_channels // 4
973
+ if isinstance(self.controlnet, FluxControlNetModel):
974
+ control_image, height, width, packed_image_latents, packed_mask = self.prepare_image_with_mask(
975
+ image=control_image,
976
+ mask=control_mask,
977
+ width=width,
978
+ height=height,
979
+ batch_size=batch_size * num_images_per_prompt,
980
+ num_images_per_prompt=num_images_per_prompt,
981
+ device=device,
982
+ dtype=dtype,
983
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
984
+ )
985
+
986
+ # 4. Prepare latent variables
987
+ num_channels_latents = self.transformer.config.in_channels // 4
988
+ latents, latent_image_ids = self.prepare_latents(
989
+ batch_size * num_images_per_prompt,
990
+ num_channels_latents,
991
+ height,
992
+ width,
993
+ prompt_embeds.dtype,
994
+ device,
995
+ generator,
996
+ latents,
997
+ )
998
+ # noise = latents.clone()
999
+
1000
+ if self.do_classifier_free_guidance:
1001
+ latent_image_ids = torch.cat([latent_image_ids] * 2)
1002
+
1003
+ # 5. Prepare timesteps
1004
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
1005
+ image_seq_len = latents.shape[1]
1006
+ mu = calculate_shift(
1007
+ image_seq_len,
1008
+ self.scheduler.config.base_image_seq_len,
1009
+ self.scheduler.config.max_image_seq_len,
1010
+ self.scheduler.config.base_shift,
1011
+ self.scheduler.config.max_shift,
1012
+ )
1013
+ timesteps, num_inference_steps = retrieve_timesteps(
1014
+ self.scheduler,
1015
+ num_inference_steps,
1016
+ device,
1017
+ timesteps,
1018
+ sigmas,
1019
+ mu=mu,
1020
+ )
1021
+
1022
+ num_warmup_steps = max(
1023
+ len(timesteps) - num_inference_steps * self.scheduler.order, 0
1024
+ )
1025
+ self._num_timesteps = len(timesteps)
1026
+
1027
+ # 6. Denoising loop
1028
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
1029
+ for i, t in enumerate(timesteps):
1030
+ if self.interrupt:
1031
+ continue
1032
+
1033
+ latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
1034
+
1035
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
1036
+ timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
1037
+
1038
+ # handle guidance
1039
+ if self.transformer.config.guidance_embeds:
1040
+ guidance = torch.tensor([guidance_scale], device=device)
1041
+ guidance = guidance.expand(latent_model_input.shape[0])
1042
+ else:
1043
+ guidance = None
1044
+
1045
+ # controlnet
1046
+ (
1047
+ controlnet_block_samples,
1048
+ controlnet_single_block_samples,
1049
+ ) = self.controlnet(
1050
+ hidden_states=latent_model_input,
1051
+ controlnet_cond=control_image,
1052
+ conditioning_scale=controlnet_conditioning_scale,
1053
+ timestep=timestep / 1000,
1054
+ guidance=guidance,
1055
+ pooled_projections=pooled_prompt_embeds,
1056
+ encoder_hidden_states=prompt_embeds,
1057
+ txt_ids=text_ids,
1058
+ img_ids=latent_image_ids,
1059
+ joint_attention_kwargs=self.joint_attention_kwargs,
1060
+ return_dict=False,
1061
+ )
1062
+
1063
+ noise_pred = self.transformer(
1064
+ hidden_states=latent_model_input,
1065
+ # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
1066
+ timestep=timestep / 1000,
1067
+ guidance=guidance,
1068
+ pooled_projections=pooled_prompt_embeds,
1069
+ encoder_hidden_states=prompt_embeds,
1070
+ controlnet_block_samples=[
1071
+ sample.to(dtype=self.transformer.dtype)
1072
+ for sample in controlnet_block_samples
1073
+ ],
1074
+ controlnet_single_block_samples=[
1075
+ sample.to(dtype=self.transformer.dtype)
1076
+ for sample in controlnet_single_block_samples
1077
+ ] if controlnet_single_block_samples is not None else controlnet_single_block_samples,
1078
+ txt_ids=text_ids[0],
1079
+ img_ids=latent_image_ids[0],
1080
+ joint_attention_kwargs=self.joint_attention_kwargs,
1081
+ return_dict=False,
1082
+ )[0]
1083
+
1084
+ # 在生成循环中_org_replace
1085
+ if self.do_classifier_free_guidance:
1086
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
1087
+ noise_pred = noise_pred_uncond + true_guidance_scale * (noise_pred_text - noise_pred_uncond)
1088
+
1089
+ # compute the previous noisy sample x_t -> x_t-1
1090
+ latents_dtype = latents.dtype
1091
+ latents = self.scheduler.step(
1092
+ noise_pred, t, latents, return_dict=False
1093
+ )[0]
1094
+
1095
+ if latents.dtype != latents_dtype:
1096
+ if torch.backends.mps.is_available():
1097
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
1098
+ latents = latents.to(latents_dtype)
1099
+ #
1100
+ # init_latents_proper = packed_image_latents
1101
+ # if i < len(timesteps) - 1:
1102
+ # noise_timestep = timesteps[i + 1]
1103
+ # init_latents_proper = self.scheduler.scale_noise(
1104
+ # init_latents_proper, torch.tensor([noise_timestep]), noise
1105
+ # )
1106
+ #
1107
+ # latents = packed_mask * init_latents_proper + (1 - packed_mask) * latents
1108
+
1109
+
1110
+ if callback_on_step_end is not None:
1111
+ callback_kwargs = {}
1112
+ for k in callback_on_step_end_tensor_inputs:
1113
+ callback_kwargs[k] = locals()[k]
1114
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
1115
+
1116
+ latents = callback_outputs.pop("latents", latents)
1117
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
1118
+
1119
+ # call the callback, if provided
1120
+ if i == len(timesteps) - 1 or (
1121
+ (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
1122
+ ):
1123
+ progress_bar.update()
1124
+
1125
+ if XLA_AVAILABLE:
1126
+ xm.mark_step()
1127
+
1128
+ if output_type == "latent":
1129
+ image = latents
1130
+
1131
+ else:
1132
+ latents = self._unpack_latents(
1133
+ latents, height, width, self.vae_scale_factor
1134
+ )
1135
+ latents = (
1136
+ latents / self.vae.config.scaling_factor
1137
+ ) + self.vae.config.shift_factor
1138
+ latents = latents.to(self.vae.dtype)
1139
+
1140
+ image = self.vae.decode(latents, return_dict=False)[0]
1141
+ image = self.image_processor.postprocess(image, output_type=output_type)
1142
+
1143
+ # Offload all models
1144
+ self.maybe_free_model_hooks()
1145
+
1146
+ if not return_dict:
1147
+ return (image,)
1148
+
1149
+ return FluxPipelineOutput(images=image)
requirements.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ opencv-python
2
+ accelerate
3
+ protobuf
4
+ sentencepiece
5
+ diffusers==0.31.0
6
+ huggingface-hub==0.26.2
7
+ transformers==4.46.3
8
+ Pillow==9.5.0
transformer_flux.py ADDED
@@ -0,0 +1,525 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict, List, Optional, Union
2
+
3
+ import numpy as np
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+
8
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
9
+ from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
10
+ from diffusers.models.attention import FeedForward
11
+ from diffusers.models.attention_processor import (
12
+ Attention,
13
+ FluxAttnProcessor2_0,
14
+ FluxSingleAttnProcessor2_0,
15
+ )
16
+ from diffusers.models.modeling_utils import ModelMixin
17
+ from diffusers.models.normalization import (
18
+ AdaLayerNormContinuous,
19
+ AdaLayerNormZero,
20
+ AdaLayerNormZeroSingle,
21
+ )
22
+ from diffusers.utils import (
23
+ USE_PEFT_BACKEND,
24
+ is_torch_version,
25
+ logging,
26
+ scale_lora_layers,
27
+ unscale_lora_layers,
28
+ )
29
+ from diffusers.utils.torch_utils import maybe_allow_in_graph
30
+ from diffusers.models.embeddings import (
31
+ CombinedTimestepGuidanceTextProjEmbeddings,
32
+ CombinedTimestepTextProjEmbeddings,
33
+ )
34
+ from diffusers.models.modeling_outputs import Transformer2DModelOutput
35
+
36
+
37
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
38
+
39
+
40
+ # YiYi to-do: refactor rope related functions/classes
41
+ def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
42
+ assert dim % 2 == 0, "The dimension must be even."
43
+
44
+ scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
45
+ omega = 1.0 / (theta**scale)
46
+
47
+ batch_size, seq_length = pos.shape
48
+ out = torch.einsum("...n,d->...nd", pos, omega)
49
+ cos_out = torch.cos(out)
50
+ sin_out = torch.sin(out)
51
+
52
+ stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1)
53
+ out = stacked_out.view(batch_size, -1, dim // 2, 2, 2)
54
+ return out.float()
55
+
56
+
57
+ # YiYi to-do: refactor rope related functions/classes
58
+ class EmbedND(nn.Module):
59
+ def __init__(self, dim: int, theta: int, axes_dim: List[int]):
60
+ super().__init__()
61
+ self.dim = dim
62
+ self.theta = theta
63
+ self.axes_dim = axes_dim
64
+
65
+ def forward(self, ids: torch.Tensor) -> torch.Tensor:
66
+ n_axes = ids.shape[-1]
67
+ emb = torch.cat(
68
+ [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
69
+ dim=-3,
70
+ )
71
+ return emb.unsqueeze(1)
72
+
73
+
74
+ @maybe_allow_in_graph
75
+ class FluxSingleTransformerBlock(nn.Module):
76
+ r"""
77
+ A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
78
+
79
+ Reference: https://arxiv.org/abs/2403.03206
80
+
81
+ Parameters:
82
+ dim (`int`): The number of channels in the input and output.
83
+ num_attention_heads (`int`): The number of heads to use for multi-head attention.
84
+ attention_head_dim (`int`): The number of channels in each head.
85
+ context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
86
+ processing of `context` conditions.
87
+ """
88
+
89
+ def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
90
+ super().__init__()
91
+ self.mlp_hidden_dim = int(dim * mlp_ratio)
92
+
93
+ self.norm = AdaLayerNormZeroSingle(dim)
94
+ self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
95
+ self.act_mlp = nn.GELU(approximate="tanh")
96
+ self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
97
+
98
+ processor = FluxSingleAttnProcessor2_0()
99
+ self.attn = Attention(
100
+ query_dim=dim,
101
+ cross_attention_dim=None,
102
+ dim_head=attention_head_dim,
103
+ heads=num_attention_heads,
104
+ out_dim=dim,
105
+ bias=True,
106
+ processor=processor,
107
+ qk_norm="rms_norm",
108
+ eps=1e-6,
109
+ pre_only=True,
110
+ )
111
+
112
+ def forward(
113
+ self,
114
+ hidden_states: torch.FloatTensor,
115
+ temb: torch.FloatTensor,
116
+ image_rotary_emb=None,
117
+ ):
118
+ residual = hidden_states
119
+ norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
120
+ mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
121
+
122
+ attn_output = self.attn(
123
+ hidden_states=norm_hidden_states,
124
+ image_rotary_emb=image_rotary_emb,
125
+ )
126
+
127
+ hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
128
+ gate = gate.unsqueeze(1)
129
+ hidden_states = gate * self.proj_out(hidden_states)
130
+ hidden_states = residual + hidden_states
131
+ if hidden_states.dtype == torch.float16:
132
+ hidden_states = hidden_states.clip(-65504, 65504)
133
+
134
+ return hidden_states
135
+
136
+
137
+ @maybe_allow_in_graph
138
+ class FluxTransformerBlock(nn.Module):
139
+ r"""
140
+ A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
141
+
142
+ Reference: https://arxiv.org/abs/2403.03206
143
+
144
+ Parameters:
145
+ dim (`int`): The number of channels in the input and output.
146
+ num_attention_heads (`int`): The number of heads to use for multi-head attention.
147
+ attention_head_dim (`int`): The number of channels in each head.
148
+ context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
149
+ processing of `context` conditions.
150
+ """
151
+
152
+ def __init__(
153
+ self, dim, num_attention_heads, attention_head_dim, qk_norm="rms_norm", eps=1e-6
154
+ ):
155
+ super().__init__()
156
+
157
+ self.norm1 = AdaLayerNormZero(dim)
158
+
159
+ self.norm1_context = AdaLayerNormZero(dim)
160
+
161
+ if hasattr(F, "scaled_dot_product_attention"):
162
+ processor = FluxAttnProcessor2_0()
163
+ else:
164
+ raise ValueError(
165
+ "The current PyTorch version does not support the `scaled_dot_product_attention` function."
166
+ )
167
+ self.attn = Attention(
168
+ query_dim=dim,
169
+ cross_attention_dim=None,
170
+ added_kv_proj_dim=dim,
171
+ dim_head=attention_head_dim,
172
+ heads=num_attention_heads,
173
+ out_dim=dim,
174
+ context_pre_only=False,
175
+ bias=True,
176
+ processor=processor,
177
+ qk_norm=qk_norm,
178
+ eps=eps,
179
+ )
180
+
181
+ self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
182
+ self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
183
+
184
+ self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
185
+ self.ff_context = FeedForward(
186
+ dim=dim, dim_out=dim, activation_fn="gelu-approximate"
187
+ )
188
+
189
+ # let chunk size default to None
190
+ self._chunk_size = None
191
+ self._chunk_dim = 0
192
+
193
+ def forward(
194
+ self,
195
+ hidden_states: torch.FloatTensor,
196
+ encoder_hidden_states: torch.FloatTensor,
197
+ temb: torch.FloatTensor,
198
+ image_rotary_emb=None,
199
+ ):
200
+ norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
201
+ hidden_states, emb=temb
202
+ )
203
+
204
+ (
205
+ norm_encoder_hidden_states,
206
+ c_gate_msa,
207
+ c_shift_mlp,
208
+ c_scale_mlp,
209
+ c_gate_mlp,
210
+ ) = self.norm1_context(encoder_hidden_states, emb=temb)
211
+
212
+ # Attention.
213
+ attn_output, context_attn_output = self.attn(
214
+ hidden_states=norm_hidden_states,
215
+ encoder_hidden_states=norm_encoder_hidden_states,
216
+ image_rotary_emb=image_rotary_emb,
217
+ )
218
+
219
+ # Process attention outputs for the `hidden_states`.
220
+ attn_output = gate_msa.unsqueeze(1) * attn_output
221
+ hidden_states = hidden_states + attn_output
222
+
223
+ norm_hidden_states = self.norm2(hidden_states)
224
+ norm_hidden_states = (
225
+ norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
226
+ )
227
+
228
+ ff_output = self.ff(norm_hidden_states)
229
+ ff_output = gate_mlp.unsqueeze(1) * ff_output
230
+
231
+ hidden_states = hidden_states + ff_output
232
+
233
+ # Process attention outputs for the `encoder_hidden_states`.
234
+
235
+ context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
236
+ encoder_hidden_states = encoder_hidden_states + context_attn_output
237
+
238
+ norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
239
+ norm_encoder_hidden_states = (
240
+ norm_encoder_hidden_states * (1 + c_scale_mlp[:, None])
241
+ + c_shift_mlp[:, None]
242
+ )
243
+
244
+ context_ff_output = self.ff_context(norm_encoder_hidden_states)
245
+ encoder_hidden_states = (
246
+ encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
247
+ )
248
+ if encoder_hidden_states.dtype == torch.float16:
249
+ encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
250
+
251
+ return encoder_hidden_states, hidden_states
252
+
253
+
254
+ class FluxTransformer2DModel(
255
+ ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
256
+ ):
257
+ """
258
+ The Transformer model introduced in Flux.
259
+
260
+ Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
261
+
262
+ Parameters:
263
+ patch_size (`int`): Patch size to turn the input data into small patches.
264
+ in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
265
+ num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
266
+ num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
267
+ attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
268
+ num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
269
+ joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
270
+ pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
271
+ guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
272
+ """
273
+
274
+ _supports_gradient_checkpointing = True
275
+
276
+ @register_to_config
277
+ def __init__(
278
+ self,
279
+ patch_size: int = 1,
280
+ in_channels: int = 64,
281
+ num_layers: int = 19,
282
+ num_single_layers: int = 38,
283
+ attention_head_dim: int = 128,
284
+ num_attention_heads: int = 24,
285
+ joint_attention_dim: int = 4096,
286
+ pooled_projection_dim: int = 768,
287
+ guidance_embeds: bool = False,
288
+ axes_dims_rope: List[int] = [16, 56, 56],
289
+ ):
290
+ super().__init__()
291
+ self.out_channels = in_channels
292
+ self.inner_dim = (
293
+ self.config.num_attention_heads * self.config.attention_head_dim
294
+ )
295
+
296
+ self.pos_embed = EmbedND(
297
+ dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope
298
+ )
299
+ text_time_guidance_cls = (
300
+ CombinedTimestepGuidanceTextProjEmbeddings
301
+ if guidance_embeds
302
+ else CombinedTimestepTextProjEmbeddings
303
+ )
304
+ self.time_text_embed = text_time_guidance_cls(
305
+ embedding_dim=self.inner_dim,
306
+ pooled_projection_dim=self.config.pooled_projection_dim,
307
+ )
308
+
309
+ self.context_embedder = nn.Linear(
310
+ self.config.joint_attention_dim, self.inner_dim
311
+ )
312
+ self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim)
313
+
314
+ self.transformer_blocks = nn.ModuleList(
315
+ [
316
+ FluxTransformerBlock(
317
+ dim=self.inner_dim,
318
+ num_attention_heads=self.config.num_attention_heads,
319
+ attention_head_dim=self.config.attention_head_dim,
320
+ )
321
+ for i in range(self.config.num_layers)
322
+ ]
323
+ )
324
+
325
+ self.single_transformer_blocks = nn.ModuleList(
326
+ [
327
+ FluxSingleTransformerBlock(
328
+ dim=self.inner_dim,
329
+ num_attention_heads=self.config.num_attention_heads,
330
+ attention_head_dim=self.config.attention_head_dim,
331
+ )
332
+ for i in range(self.config.num_single_layers)
333
+ ]
334
+ )
335
+
336
+ self.norm_out = AdaLayerNormContinuous(
337
+ self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
338
+ )
339
+ self.proj_out = nn.Linear(
340
+ self.inner_dim, patch_size * patch_size * self.out_channels, bias=True
341
+ )
342
+
343
+ self.gradient_checkpointing = False
344
+
345
+ def _set_gradient_checkpointing(self, module, value=False):
346
+ if hasattr(module, "gradient_checkpointing"):
347
+ module.gradient_checkpointing = value
348
+
349
+ def forward(
350
+ self,
351
+ hidden_states: torch.Tensor,
352
+ encoder_hidden_states: torch.Tensor = None,
353
+ pooled_projections: torch.Tensor = None,
354
+ timestep: torch.LongTensor = None,
355
+ img_ids: torch.Tensor = None,
356
+ txt_ids: torch.Tensor = None,
357
+ guidance: torch.Tensor = None,
358
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
359
+ controlnet_block_samples=None,
360
+ controlnet_single_block_samples=None,
361
+ return_dict: bool = True,
362
+ ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
363
+ """
364
+ The [`FluxTransformer2DModel`] forward method.
365
+
366
+ Args:
367
+ hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
368
+ Input `hidden_states`.
369
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
370
+ Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
371
+ pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
372
+ from the embeddings of input conditions.
373
+ timestep ( `torch.LongTensor`):
374
+ Used to indicate denoising step.
375
+ block_controlnet_hidden_states: (`list` of `torch.Tensor`):
376
+ A list of tensors that if specified are added to the residuals of transformer blocks.
377
+ joint_attention_kwargs (`dict`, *optional*):
378
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
379
+ `self.processor` in
380
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
381
+ return_dict (`bool`, *optional*, defaults to `True`):
382
+ Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
383
+ tuple.
384
+
385
+ Returns:
386
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
387
+ `tuple` where the first element is the sample tensor.
388
+ """
389
+ if joint_attention_kwargs is not None:
390
+ joint_attention_kwargs = joint_attention_kwargs.copy()
391
+ lora_scale = joint_attention_kwargs.pop("scale", 1.0)
392
+ else:
393
+ lora_scale = 1.0
394
+
395
+ if USE_PEFT_BACKEND:
396
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
397
+ scale_lora_layers(self, lora_scale)
398
+ else:
399
+ if (
400
+ joint_attention_kwargs is not None
401
+ and joint_attention_kwargs.get("scale", None) is not None
402
+ ):
403
+ logger.warning(
404
+ "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
405
+ )
406
+ hidden_states = self.x_embedder(hidden_states)
407
+
408
+ timestep = timestep.to(hidden_states.dtype) * 1000
409
+ if guidance is not None:
410
+ guidance = guidance.to(hidden_states.dtype) * 1000
411
+ else:
412
+ guidance = None
413
+ temb = (
414
+ self.time_text_embed(timestep, pooled_projections)
415
+ if guidance is None
416
+ else self.time_text_embed(timestep, guidance, pooled_projections)
417
+ )
418
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
419
+
420
+ txt_ids = txt_ids.expand(img_ids.size(0), -1, -1)
421
+ ids = torch.cat((txt_ids, img_ids), dim=1)
422
+ image_rotary_emb = self.pos_embed(ids)
423
+
424
+ for index_block, block in enumerate(self.transformer_blocks):
425
+ if self.training and self.gradient_checkpointing:
426
+
427
+ def create_custom_forward(module, return_dict=None):
428
+ def custom_forward(*inputs):
429
+ if return_dict is not None:
430
+ return module(*inputs, return_dict=return_dict)
431
+ else:
432
+ return module(*inputs)
433
+
434
+ return custom_forward
435
+
436
+ ckpt_kwargs: Dict[str, Any] = (
437
+ {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
438
+ )
439
+ (
440
+ encoder_hidden_states,
441
+ hidden_states,
442
+ ) = torch.utils.checkpoint.checkpoint(
443
+ create_custom_forward(block),
444
+ hidden_states,
445
+ encoder_hidden_states,
446
+ temb,
447
+ image_rotary_emb,
448
+ **ckpt_kwargs,
449
+ )
450
+
451
+ else:
452
+ encoder_hidden_states, hidden_states = block(
453
+ hidden_states=hidden_states,
454
+ encoder_hidden_states=encoder_hidden_states,
455
+ temb=temb,
456
+ image_rotary_emb=image_rotary_emb,
457
+ )
458
+
459
+ # controlnet residual
460
+ if controlnet_block_samples is not None:
461
+ interval_control = len(self.transformer_blocks) / len(
462
+ controlnet_block_samples
463
+ )
464
+ interval_control = int(np.ceil(interval_control))
465
+ hidden_states = (
466
+ hidden_states
467
+ + controlnet_block_samples[index_block // interval_control]
468
+ )
469
+
470
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
471
+
472
+ for index_block, block in enumerate(self.single_transformer_blocks):
473
+ if self.training and self.gradient_checkpointing:
474
+
475
+ def create_custom_forward(module, return_dict=None):
476
+ def custom_forward(*inputs):
477
+ if return_dict is not None:
478
+ return module(*inputs, return_dict=return_dict)
479
+ else:
480
+ return module(*inputs)
481
+
482
+ return custom_forward
483
+
484
+ ckpt_kwargs: Dict[str, Any] = (
485
+ {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
486
+ )
487
+ hidden_states = torch.utils.checkpoint.checkpoint(
488
+ create_custom_forward(block),
489
+ hidden_states,
490
+ temb,
491
+ image_rotary_emb,
492
+ **ckpt_kwargs,
493
+ )
494
+
495
+ else:
496
+ hidden_states = block(
497
+ hidden_states=hidden_states,
498
+ temb=temb,
499
+ image_rotary_emb=image_rotary_emb,
500
+ )
501
+
502
+ # controlnet residual
503
+ if controlnet_single_block_samples is not None:
504
+ interval_control = len(self.single_transformer_blocks) / len(
505
+ controlnet_single_block_samples
506
+ )
507
+ interval_control = int(np.ceil(interval_control))
508
+ hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
509
+ hidden_states[:, encoder_hidden_states.shape[1] :, ...]
510
+ + controlnet_single_block_samples[index_block // interval_control]
511
+ )
512
+
513
+ hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
514
+
515
+ hidden_states = self.norm_out(hidden_states, temb)
516
+ output = self.proj_out(hidden_states)
517
+
518
+ if USE_PEFT_BACKEND:
519
+ # remove `lora_scale` from each PEFT layer
520
+ unscale_lora_layers(self, lora_scale)
521
+
522
+ if not return_dict:
523
+ return (output,)
524
+
525
+ return Transformer2DModelOutput(sample=output)