OmPrakashSingh1704 commited on
Commit
9bc5731
1 Parent(s): 4d52cb6
options/Banner.py CHANGED
@@ -1,4 +1,3 @@
1
- import torch
2
  from options.Banner_Model.Text2Banner import T2I
3
  from options.Banner_Model.Image2Image import I2I
4
  from options.Banner_Model.Image2Image_2 import I2I_2
 
 
1
  from options.Banner_Model.Text2Banner import T2I
2
  from options.Banner_Model.Image2Image import I2I
3
  from options.Banner_Model.Image2Image_2 import I2I_2
options/Banner_Model/Image2Image.py CHANGED
@@ -4,7 +4,7 @@ import requests
4
  import random,os
5
  import numpy as np
6
  import gradio as gr
7
- # import spaces
8
  import torch
9
  from PIL import Image
10
  from diffusers import FluxInpaintPipeline
@@ -39,7 +39,7 @@ def resize_image_dimensions(
39
  return new_width, new_height
40
 
41
 
42
- # @spaces.GPU(duration=100)
43
  def I2I(
44
  input_image_editor: dict,
45
  input_text: str,
 
4
  import random,os
5
  import numpy as np
6
  import gradio as gr
7
+ import spaces
8
  import torch
9
  from PIL import Image
10
  from diffusers import FluxInpaintPipeline
 
39
  return new_width, new_height
40
 
41
 
42
+ @spaces.GPU(duration=100)
43
  def I2I(
44
  input_image_editor: dict,
45
  input_text: str,
options/Banner_Model/Image2Image_2.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import torch
2
  from controlnet_aux import LineartDetector
3
  from diffusers import ControlNetModel,UniPCMultistepScheduler,StableDiffusionControlNetPipeline
@@ -6,7 +7,7 @@ from PIL import Image
6
  device= "cuda" if torch.cuda.is_available() else "cpu"
7
  print("Using device for I2I_2:", device)
8
 
9
-
10
  def I2I_2(image, prompt,size,num_inference_steps):
11
  processor = LineartDetector.from_pretrained("lllyasviel/Annotators")
12
 
 
1
+ import spaces
2
  import torch
3
  from controlnet_aux import LineartDetector
4
  from diffusers import ControlNetModel,UniPCMultistepScheduler,StableDiffusionControlNetPipeline
 
7
  device= "cuda" if torch.cuda.is_available() else "cpu"
8
  print("Using device for I2I_2:", device)
9
 
10
+ @spaces.GPU(duration=100)
11
  def I2I_2(image, prompt,size,num_inference_steps):
12
  processor = LineartDetector.from_pretrained("lllyasviel/Annotators")
13
 
options/Banner_Model/Text2Banner.py CHANGED
@@ -1,9 +1,8 @@
1
  from huggingface_hub import InferenceClient
2
- import torch
3
- device="cuda" if torch.cuda.is_available() else "cpu"
4
  def T2I(prompt, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28):
5
  # Initialize the model client
6
- model = InferenceClient(model="black-forest-labs/FLUX.1-dev").to(device)
7
 
8
  # Prepare the request parameters
9
  payload = {
 
1
  from huggingface_hub import InferenceClient
2
+
 
3
  def T2I(prompt, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28):
4
  # Initialize the model client
5
+ model = InferenceClient(model="black-forest-labs/FLUX.1-dev")
6
 
7
  # Prepare the request parameters
8
  payload = {
options/Banner_Model/__pycache__/Image2Image.cpython-310.pyc CHANGED
Binary files a/options/Banner_Model/__pycache__/Image2Image.cpython-310.pyc and b/options/Banner_Model/__pycache__/Image2Image.cpython-310.pyc differ
 
options/Banner_Model/__pycache__/Image2Image_2.cpython-310.pyc CHANGED
Binary files a/options/Banner_Model/__pycache__/Image2Image_2.cpython-310.pyc and b/options/Banner_Model/__pycache__/Image2Image_2.cpython-310.pyc differ
 
options/Banner_Model/__pycache__/Text2Banner.cpython-310.pyc CHANGED
Binary files a/options/Banner_Model/__pycache__/Text2Banner.cpython-310.pyc and b/options/Banner_Model/__pycache__/Text2Banner.cpython-310.pyc differ
 
options/Banner_Model/__pycache__/__init__.cpython-310.pyc CHANGED
Binary files a/options/Banner_Model/__pycache__/__init__.cpython-310.pyc and b/options/Banner_Model/__pycache__/__init__.cpython-310.pyc differ
 
options/Banner_Model/controlnet_flux.py DELETED
@@ -1,418 +0,0 @@
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
- )
23
- from diffusers.models.modeling_outputs import Transformer2DModelOutput
24
- from .transformer_flux import (
25
- EmbedND,
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
- text_time_guidance_cls = (
66
- CombinedTimestepGuidanceTextProjEmbeddings
67
- if guidance_embeds
68
- else CombinedTimestepTextProjEmbeddings
69
- )
70
- self.time_text_embed = text_time_guidance_cls(
71
- embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
72
- )
73
-
74
- self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
75
- self.x_embedder = nn.Linear(in_channels, self.inner_dim)
76
-
77
- self.transformer_blocks = nn.ModuleList(
78
- [
79
- FluxTransformerBlock(
80
- dim=self.inner_dim,
81
- num_attention_heads=num_attention_heads,
82
- attention_head_dim=attention_head_dim,
83
- )
84
- for _ in range(num_layers)
85
- ]
86
- )
87
-
88
- self.single_transformer_blocks = nn.ModuleList(
89
- [
90
- FluxSingleTransformerBlock(
91
- dim=self.inner_dim,
92
- num_attention_heads=num_attention_heads,
93
- attention_head_dim=attention_head_dim,
94
- )
95
- for _ in range(num_single_layers)
96
- ]
97
- )
98
-
99
- # controlnet_blocks
100
- self.controlnet_blocks = nn.ModuleList([])
101
- for _ in range(len(self.transformer_blocks)):
102
- self.controlnet_blocks.append(
103
- zero_module(nn.Linear(self.inner_dim, self.inner_dim))
104
- )
105
-
106
- self.controlnet_single_blocks = nn.ModuleList([])
107
- for _ in range(len(self.single_transformer_blocks)):
108
- self.controlnet_single_blocks.append(
109
- zero_module(nn.Linear(self.inner_dim, self.inner_dim))
110
- )
111
-
112
- self.controlnet_x_embedder = zero_module(
113
- torch.nn.Linear(in_channels + extra_condition_channels, self.inner_dim)
114
- )
115
-
116
- self.gradient_checkpointing = False
117
-
118
- @property
119
- # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
120
- def attn_processors(self):
121
- r"""
122
- Returns:
123
- `dict` of attention processors: A dictionary containing all attention processors used in the model with
124
- indexed by its weight name.
125
- """
126
- # set recursively
127
- processors = {}
128
-
129
- def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
130
- if hasattr(module, "get_processor"):
131
- processors[f"{name}.processor"] = module.get_processor()
132
-
133
- for sub_name, child in module.named_children():
134
- fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
135
-
136
- return processors
137
-
138
- for name, module in self.named_children():
139
- fn_recursive_add_processors(name, module, processors)
140
-
141
- return processors
142
-
143
- # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
144
- def set_attn_processor(self, processor):
145
- r"""
146
- Sets the attention processor to use to compute attention.
147
-
148
- Parameters:
149
- processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
150
- The instantiated processor class or a dictionary of processor classes that will be set as the processor
151
- for **all** `Attention` layers.
152
-
153
- If `processor` is a dict, the key needs to define the path to the corresponding cross attention
154
- processor. This is strongly recommended when setting trainable attention processors.
155
-
156
- """
157
- count = len(self.attn_processors.keys())
158
-
159
- if isinstance(processor, dict) and len(processor) != count:
160
- raise ValueError(
161
- f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
162
- f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
163
- )
164
-
165
- def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
166
- if hasattr(module, "set_processor"):
167
- if not isinstance(processor, dict):
168
- module.set_processor(processor)
169
- else:
170
- module.set_processor(processor.pop(f"{name}.processor"))
171
-
172
- for sub_name, child in module.named_children():
173
- fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
174
-
175
- for name, module in self.named_children():
176
- fn_recursive_attn_processor(name, module, processor)
177
-
178
- def _set_gradient_checkpointing(self, module, value=False):
179
- if hasattr(module, "gradient_checkpointing"):
180
- module.gradient_checkpointing = value
181
-
182
- @classmethod
183
- def from_transformer(
184
- cls,
185
- transformer,
186
- num_layers: int = 4,
187
- num_single_layers: int = 10,
188
- attention_head_dim: int = 128,
189
- num_attention_heads: int = 24,
190
- load_weights_from_transformer=True,
191
- ):
192
- config = transformer.config
193
- config["num_layers"] = num_layers
194
- config["num_single_layers"] = num_single_layers
195
- config["attention_head_dim"] = attention_head_dim
196
- config["num_attention_heads"] = num_attention_heads
197
-
198
- controlnet = cls(**config)
199
-
200
- if load_weights_from_transformer:
201
- controlnet.pos_embed.load_state_dict(transformer.pos_embed.state_dict())
202
- controlnet.time_text_embed.load_state_dict(
203
- transformer.time_text_embed.state_dict()
204
- )
205
- controlnet.context_embedder.load_state_dict(
206
- transformer.context_embedder.state_dict()
207
- )
208
- controlnet.x_embedder.load_state_dict(transformer.x_embedder.state_dict())
209
- controlnet.transformer_blocks.load_state_dict(
210
- transformer.transformer_blocks.state_dict(), strict=False
211
- )
212
- controlnet.single_transformer_blocks.load_state_dict(
213
- transformer.single_transformer_blocks.state_dict(), strict=False
214
- )
215
-
216
- controlnet.controlnet_x_embedder = zero_module(
217
- controlnet.controlnet_x_embedder
218
- )
219
-
220
- return controlnet
221
-
222
- def forward(
223
- self,
224
- hidden_states: torch.Tensor,
225
- controlnet_cond: torch.Tensor,
226
- conditioning_scale: float = 1.0,
227
- encoder_hidden_states: torch.Tensor = None,
228
- pooled_projections: torch.Tensor = None,
229
- timestep: torch.LongTensor = None,
230
- img_ids: torch.Tensor = None,
231
- txt_ids: torch.Tensor = None,
232
- guidance: torch.Tensor = None,
233
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
234
- return_dict: bool = True,
235
- ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
236
- """
237
- The [`FluxTransformer2DModel`] forward method.
238
-
239
- Args:
240
- hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
241
- Input `hidden_states`.
242
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
243
- Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
244
- pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
245
- from the embeddings of input conditions.
246
- timestep ( `torch.LongTensor`):
247
- Used to indicate denoising step.
248
- block_controlnet_hidden_states: (`list` of `torch.Tensor`):
249
- A list of tensors that if specified are added to the residuals of transformer blocks.
250
- joint_attention_kwargs (`dict`, *optional*):
251
- A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
252
- `self.processor` in
253
- [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
254
- return_dict (`bool`, *optional*, defaults to `True`):
255
- Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
256
- tuple.
257
-
258
- Returns:
259
- If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
260
- `tuple` where the first element is the sample tensor.
261
- """
262
- if joint_attention_kwargs is not None:
263
- joint_attention_kwargs = joint_attention_kwargs.copy()
264
- lora_scale = joint_attention_kwargs.pop("scale", 1.0)
265
- else:
266
- lora_scale = 1.0
267
-
268
- if USE_PEFT_BACKEND:
269
- # weight the lora layers by setting `lora_scale` for each PEFT layer
270
- scale_lora_layers(self, lora_scale)
271
- else:
272
- if (
273
- joint_attention_kwargs is not None
274
- and joint_attention_kwargs.get("scale", None) is not None
275
- ):
276
- logger.warning(
277
- "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
278
- )
279
- hidden_states = self.x_embedder(hidden_states)
280
-
281
- # add condition
282
- hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond)
283
-
284
- timestep = timestep.to(hidden_states.dtype) * 1000
285
- if guidance is not None:
286
- guidance = guidance.to(hidden_states.dtype) * 1000
287
- else:
288
- guidance = None
289
- temb = (
290
- self.time_text_embed(timestep, pooled_projections)
291
- if guidance is None
292
- else self.time_text_embed(timestep, guidance, pooled_projections)
293
- )
294
- encoder_hidden_states = self.context_embedder(encoder_hidden_states)
295
-
296
- txt_ids = txt_ids.expand(img_ids.size(0), -1, -1)
297
- ids = torch.cat((txt_ids, img_ids), dim=1)
298
- image_rotary_emb = self.pos_embed(ids)
299
-
300
- block_samples = ()
301
- for _, block in enumerate(self.transformer_blocks):
302
- if self.training and self.gradient_checkpointing:
303
-
304
- def create_custom_forward(module, return_dict=None):
305
- def custom_forward(*inputs):
306
- if return_dict is not None:
307
- return module(*inputs, return_dict=return_dict)
308
- else:
309
- return module(*inputs)
310
-
311
- return custom_forward
312
-
313
- ckpt_kwargs: Dict[str, Any] = (
314
- {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
315
- )
316
- (
317
- encoder_hidden_states,
318
- hidden_states,
319
- ) = torch.utils.checkpoint.checkpoint(
320
- create_custom_forward(block),
321
- hidden_states,
322
- encoder_hidden_states,
323
- temb,
324
- image_rotary_emb,
325
- **ckpt_kwargs,
326
- )
327
-
328
- else:
329
- encoder_hidden_states, hidden_states = block(
330
- hidden_states=hidden_states,
331
- encoder_hidden_states=encoder_hidden_states,
332
- temb=temb,
333
- image_rotary_emb=image_rotary_emb,
334
- )
335
- block_samples = block_samples + (hidden_states,)
336
-
337
- hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
338
-
339
- single_block_samples = ()
340
- for _, block in enumerate(self.single_transformer_blocks):
341
- if self.training and self.gradient_checkpointing:
342
-
343
- def create_custom_forward(module, return_dict=None):
344
- def custom_forward(*inputs):
345
- if return_dict is not None:
346
- return module(*inputs, return_dict=return_dict)
347
- else:
348
- return module(*inputs)
349
-
350
- return custom_forward
351
-
352
- ckpt_kwargs: Dict[str, Any] = (
353
- {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
354
- )
355
- hidden_states = torch.utils.checkpoint.checkpoint(
356
- create_custom_forward(block),
357
- hidden_states,
358
- temb,
359
- image_rotary_emb,
360
- **ckpt_kwargs,
361
- )
362
-
363
- else:
364
- hidden_states = block(
365
- hidden_states=hidden_states,
366
- temb=temb,
367
- image_rotary_emb=image_rotary_emb,
368
- )
369
- single_block_samples = single_block_samples + (
370
- hidden_states[:, encoder_hidden_states.shape[1] :],
371
- )
372
-
373
- # controlnet block
374
- controlnet_block_samples = ()
375
- for block_sample, controlnet_block in zip(
376
- block_samples, self.controlnet_blocks
377
- ):
378
- block_sample = controlnet_block(block_sample)
379
- controlnet_block_samples = controlnet_block_samples + (block_sample,)
380
-
381
- controlnet_single_block_samples = ()
382
- for single_block_sample, controlnet_block in zip(
383
- single_block_samples, self.controlnet_single_blocks
384
- ):
385
- single_block_sample = controlnet_block(single_block_sample)
386
- controlnet_single_block_samples = controlnet_single_block_samples + (
387
- single_block_sample,
388
- )
389
-
390
- # scaling
391
- controlnet_block_samples = [
392
- sample * conditioning_scale for sample in controlnet_block_samples
393
- ]
394
- controlnet_single_block_samples = [
395
- sample * conditioning_scale for sample in controlnet_single_block_samples
396
- ]
397
-
398
- #
399
- controlnet_block_samples = (
400
- None if len(controlnet_block_samples) == 0 else controlnet_block_samples
401
- )
402
- controlnet_single_block_samples = (
403
- None
404
- if len(controlnet_single_block_samples) == 0
405
- else controlnet_single_block_samples
406
- )
407
-
408
- if USE_PEFT_BACKEND:
409
- # remove `lora_scale` from each PEFT layer
410
- unscale_lora_layers(self, lora_scale)
411
-
412
- if not return_dict:
413
- return (controlnet_block_samples, controlnet_single_block_samples)
414
-
415
- return FluxControlNetOutput(
416
- controlnet_block_samples=controlnet_block_samples,
417
- controlnet_single_block_samples=controlnet_single_block_samples,
418
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
options/Banner_Model/pipeline_flux_controlnet_inpaint.py DELETED
@@ -1,1046 +0,0 @@
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
- class FluxControlNetInpaintingPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
153
- r"""
154
- The Flux pipeline for text-to-image generation.
155
-
156
- Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
157
-
158
- Args:
159
- transformer ([`FluxTransformer2DModel`]):
160
- Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
161
- scheduler ([`FlowMatchEulerDiscreteScheduler`]):
162
- A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
163
- vae ([`AutoencoderKL`]):
164
- Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
165
- text_encoder ([`CLIPTextModel`]):
166
- [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
167
- the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
168
- text_encoder_2 ([`T5EncoderModel`]):
169
- [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
170
- the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
171
- tokenizer (`CLIPTokenizer`):
172
- Tokenizer of class
173
- [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
174
- tokenizer_2 (`T5TokenizerFast`):
175
- Second Tokenizer of class
176
- [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
177
- """
178
-
179
- model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
180
- _optional_components = []
181
- _callback_tensor_inputs = ["latents", "prompt_embeds"]
182
-
183
- def __init__(
184
- self,
185
- scheduler: FlowMatchEulerDiscreteScheduler,
186
- vae: AutoencoderKL,
187
- text_encoder: CLIPTextModel,
188
- tokenizer: CLIPTokenizer,
189
- text_encoder_2: T5EncoderModel,
190
- tokenizer_2: T5TokenizerFast,
191
- transformer: FluxTransformer2DModel,
192
- controlnet: FluxControlNetModel,
193
- ):
194
- super().__init__()
195
-
196
- self.register_modules(
197
- vae=vae,
198
- text_encoder=text_encoder,
199
- text_encoder_2=text_encoder_2,
200
- tokenizer=tokenizer,
201
- tokenizer_2=tokenizer_2,
202
- transformer=transformer,
203
- scheduler=scheduler,
204
- controlnet=controlnet,
205
- )
206
- self.vae_scale_factor = (
207
- 2 ** (len(self.vae.config.block_out_channels))
208
- if hasattr(self, "vae") and self.vae is not None
209
- else 16
210
- )
211
- self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_resize=True, do_convert_rgb=True, do_normalize=True)
212
- self.mask_processor = VaeImageProcessor(
213
- vae_scale_factor=self.vae_scale_factor,
214
- do_resize=True,
215
- do_convert_grayscale=True,
216
- do_normalize=False,
217
- do_binarize=True,
218
- )
219
- self.tokenizer_max_length = (
220
- self.tokenizer.model_max_length
221
- if hasattr(self, "tokenizer") and self.tokenizer is not None
222
- else 77
223
- )
224
- self.default_sample_size = 64
225
-
226
- @property
227
- def do_classifier_free_guidance(self):
228
- return self._guidance_scale > 1
229
-
230
- def _get_t5_prompt_embeds(
231
- self,
232
- prompt: Union[str, List[str]] = None,
233
- num_images_per_prompt: int = 1,
234
- max_sequence_length: int = 512,
235
- device: Optional[torch.device] = None,
236
- dtype: Optional[torch.dtype] = None,
237
- ):
238
- device = device or self._execution_device
239
- dtype = dtype or self.text_encoder.dtype
240
-
241
- prompt = [prompt] if isinstance(prompt, str) else prompt
242
- batch_size = len(prompt)
243
-
244
- text_inputs = self.tokenizer_2(
245
- prompt,
246
- padding="max_length",
247
- max_length=max_sequence_length,
248
- truncation=True,
249
- return_length=False,
250
- return_overflowing_tokens=False,
251
- return_tensors="pt",
252
- )
253
- text_input_ids = text_inputs.input_ids
254
- untruncated_ids = self.tokenizer_2(
255
- prompt, padding="longest", return_tensors="pt"
256
- ).input_ids
257
-
258
- if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
259
- text_input_ids, untruncated_ids
260
- ):
261
- removed_text = self.tokenizer_2.batch_decode(
262
- untruncated_ids[:, self.tokenizer_max_length - 1 : -1]
263
- )
264
- logger.warning(
265
- "The following part of your input was truncated because `max_sequence_length` is set to "
266
- f" {max_sequence_length} tokens: {removed_text}"
267
- )
268
-
269
- prompt_embeds = self.text_encoder_2(
270
- text_input_ids.to(device), output_hidden_states=False
271
- )[0]
272
-
273
- dtype = self.text_encoder_2.dtype
274
- prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
275
-
276
- _, seq_len, _ = prompt_embeds.shape
277
-
278
- # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
279
- prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
280
- prompt_embeds = prompt_embeds.view(
281
- batch_size * num_images_per_prompt, seq_len, -1
282
- )
283
-
284
- return prompt_embeds
285
-
286
- def _get_clip_prompt_embeds(
287
- self,
288
- prompt: Union[str, List[str]],
289
- num_images_per_prompt: int = 1,
290
- device: Optional[torch.device] = None,
291
- ):
292
- device = device or self._execution_device
293
-
294
- prompt = [prompt] if isinstance(prompt, str) else prompt
295
- batch_size = len(prompt)
296
-
297
- text_inputs = self.tokenizer(
298
- prompt,
299
- padding="max_length",
300
- max_length=self.tokenizer_max_length,
301
- truncation=True,
302
- return_overflowing_tokens=False,
303
- return_length=False,
304
- return_tensors="pt",
305
- )
306
-
307
- text_input_ids = text_inputs.input_ids
308
- untruncated_ids = self.tokenizer(
309
- prompt, padding="longest", return_tensors="pt"
310
- ).input_ids
311
- if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
312
- text_input_ids, untruncated_ids
313
- ):
314
- removed_text = self.tokenizer.batch_decode(
315
- untruncated_ids[:, self.tokenizer_max_length - 1 : -1]
316
- )
317
- logger.warning(
318
- "The following part of your input was truncated because CLIP can only handle sequences up to"
319
- f" {self.tokenizer_max_length} tokens: {removed_text}"
320
- )
321
- prompt_embeds = self.text_encoder(
322
- text_input_ids.to(device), output_hidden_states=False
323
- )
324
-
325
- # Use pooled output of CLIPTextModel
326
- prompt_embeds = prompt_embeds.pooler_output
327
- prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
328
-
329
- # duplicate text embeddings for each generation per prompt, using mps friendly method
330
- prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
331
- prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
332
-
333
- return prompt_embeds
334
-
335
- def encode_prompt(
336
- self,
337
- prompt: Union[str, List[str]],
338
- prompt_2: Union[str, List[str]],
339
- device: Optional[torch.device] = None,
340
- num_images_per_prompt: int = 1,
341
- do_classifier_free_guidance: bool = True,
342
- negative_prompt: Optional[Union[str, List[str]]] = None,
343
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
344
- prompt_embeds: Optional[torch.FloatTensor] = None,
345
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
346
- max_sequence_length: int = 512,
347
- lora_scale: Optional[float] = None,
348
- ):
349
- r"""
350
-
351
- Args:
352
- prompt (`str` or `List[str]`, *optional*):
353
- prompt to be encoded
354
- prompt_2 (`str` or `List[str]`, *optional*):
355
- The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
356
- used in all text-encoders
357
- device: (`torch.device`):
358
- torch device
359
- num_images_per_prompt (`int`):
360
- number of images that should be generated per prompt
361
- do_classifier_free_guidance (`bool`):
362
- whether to use classifier-free guidance or not
363
- negative_prompt (`str` or `List[str]`, *optional*):
364
- negative prompt to be encoded
365
- negative_prompt_2 (`str` or `List[str]`, *optional*):
366
- negative prompt to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is
367
- used in all text-encoders
368
- prompt_embeds (`torch.FloatTensor`, *optional*):
369
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
370
- provided, text embeddings will be generated from `prompt` input argument.
371
- pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
372
- Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
373
- If not provided, pooled text embeddings will be generated from `prompt` input argument.
374
- clip_skip (`int`, *optional*):
375
- Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
376
- the output of the pre-final layer will be used for computing the prompt embeddings.
377
- lora_scale (`float`, *optional*):
378
- A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
379
- """
380
- device = device or self._execution_device
381
-
382
- # set lora scale so that monkey patched LoRA
383
- # function of text encoder can correctly access it
384
- if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
385
- self._lora_scale = lora_scale
386
-
387
- # dynamically adjust the LoRA scale
388
- if self.text_encoder is not None and USE_PEFT_BACKEND:
389
- scale_lora_layers(self.text_encoder, lora_scale)
390
- if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
391
- scale_lora_layers(self.text_encoder_2, lora_scale)
392
-
393
- prompt = [prompt] if isinstance(prompt, str) else prompt
394
- if prompt is not None:
395
- batch_size = len(prompt)
396
- else:
397
- batch_size = prompt_embeds.shape[0]
398
-
399
- if prompt_embeds is None:
400
- prompt_2 = prompt_2 or prompt
401
- prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
402
-
403
- # We only use the pooled prompt output from the CLIPTextModel
404
- pooled_prompt_embeds = self._get_clip_prompt_embeds(
405
- prompt=prompt,
406
- device=device,
407
- num_images_per_prompt=num_images_per_prompt,
408
- )
409
- prompt_embeds = self._get_t5_prompt_embeds(
410
- prompt=prompt_2,
411
- num_images_per_prompt=num_images_per_prompt,
412
- max_sequence_length=max_sequence_length,
413
- device=device,
414
- )
415
-
416
- if do_classifier_free_guidance:
417
- # 处理 negative prompt
418
- negative_prompt = negative_prompt or ""
419
- negative_prompt_2 = negative_prompt_2 or negative_prompt
420
-
421
- negative_pooled_prompt_embeds = self._get_clip_prompt_embeds(
422
- negative_prompt,
423
- device=device,
424
- num_images_per_prompt=num_images_per_prompt,
425
- )
426
- negative_prompt_embeds = self._get_t5_prompt_embeds(
427
- negative_prompt_2,
428
- num_images_per_prompt=num_images_per_prompt,
429
- max_sequence_length=max_sequence_length,
430
- device=device,
431
- )
432
-
433
- if self.text_encoder is not None:
434
- if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
435
- # Retrieve the original scale by scaling back the LoRA layers
436
- unscale_lora_layers(self.text_encoder, lora_scale)
437
-
438
- if self.text_encoder_2 is not None:
439
- if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
440
- # Retrieve the original scale by scaling back the LoRA layers
441
- unscale_lora_layers(self.text_encoder_2, lora_scale)
442
-
443
- text_ids = torch.zeros(batch_size, prompt_embeds.shape[1], 3).to(
444
- device=device, dtype=self.text_encoder.dtype
445
- )
446
-
447
- return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds,text_ids
448
-
449
- def check_inputs(
450
- self,
451
- prompt,
452
- prompt_2,
453
- height,
454
- width,
455
- prompt_embeds=None,
456
- pooled_prompt_embeds=None,
457
- callback_on_step_end_tensor_inputs=None,
458
- max_sequence_length=None,
459
- ):
460
- if height % 8 != 0 or width % 8 != 0:
461
- raise ValueError(
462
- f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
463
- )
464
-
465
- if callback_on_step_end_tensor_inputs is not None and not all(
466
- k in self._callback_tensor_inputs
467
- for k in callback_on_step_end_tensor_inputs
468
- ):
469
- raise ValueError(
470
- 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]}"
471
- )
472
-
473
- if prompt is not None and prompt_embeds is not None:
474
- raise ValueError(
475
- f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
476
- " only forward one of the two."
477
- )
478
- elif prompt_2 is not None and prompt_embeds is not None:
479
- raise ValueError(
480
- f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
481
- " only forward one of the two."
482
- )
483
- elif prompt is None and prompt_embeds is None:
484
- raise ValueError(
485
- "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
486
- )
487
- elif prompt is not None and (
488
- not isinstance(prompt, str) and not isinstance(prompt, list)
489
- ):
490
- raise ValueError(
491
- f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
492
- )
493
- elif prompt_2 is not None and (
494
- not isinstance(prompt_2, str) and not isinstance(prompt_2, list)
495
- ):
496
- raise ValueError(
497
- f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}"
498
- )
499
-
500
- if prompt_embeds is not None and pooled_prompt_embeds is None:
501
- raise ValueError(
502
- "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`."
503
- )
504
-
505
- if max_sequence_length is not None and max_sequence_length > 512:
506
- raise ValueError(
507
- f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}"
508
- )
509
-
510
- # Copied from diffusers.pipelines.flux.pipeline_flux._prepare_latent_image_ids
511
- @staticmethod
512
- def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
513
- latent_image_ids = torch.zeros(height // 2, width // 2, 3)
514
- latent_image_ids[..., 1] = (
515
- latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
516
- )
517
- latent_image_ids[..., 2] = (
518
- latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
519
- )
520
-
521
- (
522
- latent_image_id_height,
523
- latent_image_id_width,
524
- latent_image_id_channels,
525
- ) = latent_image_ids.shape
526
-
527
- latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
528
- latent_image_ids = latent_image_ids.reshape(
529
- batch_size,
530
- latent_image_id_height * latent_image_id_width,
531
- latent_image_id_channels,
532
- )
533
-
534
- return latent_image_ids.to(device=device, dtype=dtype)
535
-
536
- # Copied from diffusers.pipelines.flux.pipeline_flux._pack_latents
537
- @staticmethod
538
- def _pack_latents(latents, batch_size, num_channels_latents, height, width):
539
- latents = latents.view(
540
- batch_size, num_channels_latents, height // 2, 2, width // 2, 2
541
- )
542
- latents = latents.permute(0, 2, 4, 1, 3, 5)
543
- latents = latents.reshape(
544
- batch_size, (height // 2) * (width // 2), num_channels_latents * 4
545
- )
546
-
547
- return latents
548
-
549
- # Copied from diffusers.pipelines.flux.pipeline_flux._unpack_latents
550
- @staticmethod
551
- def _unpack_latents(latents, height, width, vae_scale_factor):
552
- batch_size, num_patches, channels = latents.shape
553
-
554
- height = height // vae_scale_factor
555
- width = width // vae_scale_factor
556
-
557
- latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
558
- latents = latents.permute(0, 3, 1, 4, 2, 5)
559
-
560
- latents = latents.reshape(
561
- batch_size, channels // (2 * 2), height * 2, width * 2
562
- )
563
-
564
- return latents
565
-
566
- # Copied from diffusers.pipelines.flux.pipeline_flux.prepare_latents
567
- def prepare_latents(
568
- self,
569
- batch_size,
570
- num_channels_latents,
571
- height,
572
- width,
573
- dtype,
574
- device,
575
- generator,
576
- latents=None,
577
- ):
578
- height = 2 * (int(height) // self.vae_scale_factor)
579
- width = 2 * (int(width) // self.vae_scale_factor)
580
-
581
- shape = (batch_size, num_channels_latents, height, width)
582
-
583
- if latents is not None:
584
- latent_image_ids = self._prepare_latent_image_ids(
585
- batch_size, height, width, device, dtype
586
- )
587
- return latents.to(device=device, dtype=dtype), latent_image_ids
588
-
589
- if isinstance(generator, list) and len(generator) != batch_size:
590
- raise ValueError(
591
- f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
592
- f" size of {batch_size}. Make sure the batch size matches the length of the generators."
593
- )
594
-
595
- latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
596
- latents = self._pack_latents(
597
- latents, batch_size, num_channels_latents, height, width
598
- )
599
-
600
- latent_image_ids = self._prepare_latent_image_ids(
601
- batch_size, height, width, device, dtype
602
- )
603
-
604
- return latents, latent_image_ids
605
-
606
- # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image
607
- def prepare_image(
608
- self,
609
- image,
610
- width,
611
- height,
612
- batch_size,
613
- num_images_per_prompt,
614
- device,
615
- dtype,
616
- ):
617
- if isinstance(image, torch.Tensor):
618
- pass
619
- else:
620
- image = self.image_processor.preprocess(image, height=height, width=width)
621
-
622
- image_batch_size = image.shape[0]
623
-
624
- if image_batch_size == 1:
625
- repeat_by = batch_size
626
- else:
627
- # image batch size is the same as prompt batch size
628
- repeat_by = num_images_per_prompt
629
-
630
- image = image.repeat_interleave(repeat_by, dim=0)
631
-
632
- image = image.to(device=device, dtype=dtype)
633
-
634
- return image
635
-
636
- def prepare_image_with_mask(
637
- self,
638
- image,
639
- mask,
640
- width,
641
- height,
642
- batch_size,
643
- num_images_per_prompt,
644
- device,
645
- dtype,
646
- do_classifier_free_guidance = False,
647
- ):
648
- # Prepare image
649
- if isinstance(image, torch.Tensor):
650
- pass
651
- else:
652
- image = self.image_processor.preprocess(image, height=height, width=width)
653
-
654
- image_batch_size = image.shape[0]
655
- if image_batch_size == 1:
656
- repeat_by = batch_size
657
- else:
658
- # image batch size is the same as prompt batch size
659
- repeat_by = num_images_per_prompt
660
- image = image.repeat_interleave(repeat_by, dim=0)
661
- image = image.to(device=device, dtype=dtype)
662
-
663
- # Prepare mask
664
- if isinstance(mask, torch.Tensor):
665
- pass
666
- else:
667
- mask = self.mask_processor.preprocess(mask, height=height, width=width)
668
- mask = mask.repeat_interleave(repeat_by, dim=0)
669
- mask = mask.to(device=device, dtype=dtype)
670
-
671
- # Get masked image
672
- masked_image = image.clone()
673
- masked_image[(mask > 0.5).repeat(1, 3, 1, 1)] = -1
674
-
675
- # Encode to latents
676
- image_latents = self.vae.encode(masked_image.to(self.vae.dtype)).latent_dist.sample()
677
- image_latents = (
678
- image_latents - self.vae.config.shift_factor
679
- ) * self.vae.config.scaling_factor
680
- image_latents = image_latents.to(dtype)
681
-
682
- mask = torch.nn.functional.interpolate(
683
- mask, size=(height // self.vae_scale_factor * 2, width // self.vae_scale_factor * 2)
684
- )
685
- mask = 1 - mask
686
-
687
- control_image = torch.cat([image_latents, mask], dim=1)
688
-
689
- # Pack cond latents
690
- packed_control_image = self._pack_latents(
691
- control_image,
692
- batch_size * num_images_per_prompt,
693
- control_image.shape[1],
694
- control_image.shape[2],
695
- control_image.shape[3],
696
- )
697
-
698
- if do_classifier_free_guidance:
699
- packed_control_image = torch.cat([packed_control_image] * 2)
700
-
701
- return packed_control_image, height, width
702
-
703
- @property
704
- def guidance_scale(self):
705
- return self._guidance_scale
706
-
707
- @property
708
- def joint_attention_kwargs(self):
709
- return self._joint_attention_kwargs
710
-
711
- @property
712
- def num_timesteps(self):
713
- return self._num_timesteps
714
-
715
- @property
716
- def interrupt(self):
717
- return self._interrupt
718
-
719
- @torch.no_grad()
720
- @replace_example_docstring(EXAMPLE_DOC_STRING)
721
- def __call__(
722
- self,
723
- prompt: Union[str, List[str]] = None,
724
- prompt_2: Optional[Union[str, List[str]]] = None,
725
- height: Optional[int] = None,
726
- width: Optional[int] = None,
727
- num_inference_steps: int = 28,
728
- timesteps: List[int] = None,
729
- guidance_scale: float = 7.0,
730
- true_guidance_scale: float = 3.5 ,
731
- negative_prompt: Optional[Union[str, List[str]]] = None,
732
- negative_prompt_2: Optional[Union[str, List[str]]] = None,
733
- control_image: PipelineImageInput = None,
734
- control_mask: PipelineImageInput = None,
735
- controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
736
- num_images_per_prompt: Optional[int] = 1,
737
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
738
- latents: Optional[torch.FloatTensor] = None,
739
- prompt_embeds: Optional[torch.FloatTensor] = None,
740
- pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
741
- output_type: Optional[str] = "pil",
742
- return_dict: bool = True,
743
- joint_attention_kwargs: Optional[Dict[str, Any]] = None,
744
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
745
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
746
- max_sequence_length: int = 512,
747
- ):
748
- r"""
749
- Function invoked when calling the pipeline for generation.
750
-
751
- Args:
752
- prompt (`str` or `List[str]`, *optional*):
753
- The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
754
- instead.
755
- prompt_2 (`str` or `List[str]`, *optional*):
756
- The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
757
- will be used instead
758
- height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
759
- The height in pixels of the generated image. This is set to 1024 by default for the best results.
760
- width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
761
- The width in pixels of the generated image. This is set to 1024 by default for the best results.
762
- num_inference_steps (`int`, *optional*, defaults to 50):
763
- The number of denoising steps. More denoising steps usually lead to a higher quality image at the
764
- expense of slower inference.
765
- timesteps (`List[int]`, *optional*):
766
- Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
767
- in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
768
- passed will be used. Must be in descending order.
769
- guidance_scale (`float`, *optional*, defaults to 7.0):
770
- Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
771
- `guidance_scale` is defined as `w` of equation 2. of [Imagen
772
- Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
773
- 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
774
- usually at the expense of lower image quality.
775
- num_images_per_prompt (`int`, *optional*, defaults to 1):
776
- The number of images to generate per prompt.
777
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
778
- One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
779
- to make generation deterministic.
780
- latents (`torch.FloatTensor`, *optional*):
781
- Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
782
- generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
783
- tensor will ge generated by sampling using the supplied random `generator`.
784
- prompt_embeds (`torch.FloatTensor`, *optional*):
785
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
786
- provided, text embeddings will be generated from `prompt` input argument.
787
- pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
788
- Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
789
- If not provided, pooled text embeddings will be generated from `prompt` input argument.
790
- output_type (`str`, *optional*, defaults to `"pil"`):
791
- The output format of the generate image. Choose between
792
- [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
793
- return_dict (`bool`, *optional*, defaults to `True`):
794
- Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
795
- joint_attention_kwargs (`dict`, *optional*):
796
- A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
797
- `self.processor` in
798
- [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
799
- callback_on_step_end (`Callable`, *optional*):
800
- A function that calls at the end of each denoising steps during the inference. The function is called
801
- with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
802
- callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
803
- `callback_on_step_end_tensor_inputs`.
804
- callback_on_step_end_tensor_inputs (`List`, *optional*):
805
- The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
806
- will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
807
- `._callback_tensor_inputs` attribute of your pipeline class.
808
- max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
809
-
810
- Examples:
811
-
812
- Returns:
813
- [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
814
- is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
815
- images.
816
- """
817
-
818
- height = height or self.default_sample_size * self.vae_scale_factor
819
- width = width or self.default_sample_size * self.vae_scale_factor
820
-
821
- # 1. Check inputs. Raise error if not correct
822
- self.check_inputs(
823
- prompt,
824
- prompt_2,
825
- height,
826
- width,
827
- prompt_embeds=prompt_embeds,
828
- pooled_prompt_embeds=pooled_prompt_embeds,
829
- callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
830
- max_sequence_length=max_sequence_length,
831
- )
832
-
833
- self._guidance_scale = true_guidance_scale
834
- self._joint_attention_kwargs = joint_attention_kwargs
835
- self._interrupt = False
836
-
837
- # 2. Define call parameters
838
- if prompt is not None and isinstance(prompt, str):
839
- batch_size = 1
840
- elif prompt is not None and isinstance(prompt, list):
841
- batch_size = len(prompt)
842
- else:
843
- batch_size = prompt_embeds.shape[0]
844
-
845
- device = self._execution_device
846
- dtype = self.transformer.dtype
847
-
848
- lora_scale = (
849
- self.joint_attention_kwargs.get("scale", None)
850
- if self.joint_attention_kwargs is not None
851
- else None
852
- )
853
- (
854
- prompt_embeds,
855
- pooled_prompt_embeds,
856
- negative_prompt_embeds,
857
- negative_pooled_prompt_embeds,
858
- text_ids
859
- ) = self.encode_prompt(
860
- prompt=prompt,
861
- prompt_2=prompt_2,
862
- prompt_embeds=prompt_embeds,
863
- pooled_prompt_embeds=pooled_prompt_embeds,
864
- do_classifier_free_guidance = self.do_classifier_free_guidance,
865
- negative_prompt = negative_prompt,
866
- negative_prompt_2 = negative_prompt_2,
867
- device=device,
868
- num_images_per_prompt=num_images_per_prompt,
869
- max_sequence_length=max_sequence_length,
870
- lora_scale=lora_scale,
871
- )
872
-
873
- # 在 encode_prompt 之后
874
- if self.do_classifier_free_guidance:
875
- prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim = 0)
876
- pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim = 0)
877
- text_ids = torch.cat([text_ids, text_ids], dim = 0)
878
-
879
- # 3. Prepare control image
880
- num_channels_latents = self.transformer.config.in_channels // 4
881
- if isinstance(self.controlnet, FluxControlNetModel):
882
- control_image, height, width = self.prepare_image_with_mask(
883
- image=control_image,
884
- mask=control_mask,
885
- width=width,
886
- height=height,
887
- batch_size=batch_size * num_images_per_prompt,
888
- num_images_per_prompt=num_images_per_prompt,
889
- device=device,
890
- dtype=dtype,
891
- do_classifier_free_guidance=self.do_classifier_free_guidance,
892
- )
893
-
894
- # 4. Prepare latent variables
895
- num_channels_latents = self.transformer.config.in_channels // 4
896
- latents, latent_image_ids = self.prepare_latents(
897
- batch_size * num_images_per_prompt,
898
- num_channels_latents,
899
- height,
900
- width,
901
- prompt_embeds.dtype,
902
- device,
903
- generator,
904
- latents,
905
- )
906
-
907
- if self.do_classifier_free_guidance:
908
- latent_image_ids = torch.cat([latent_image_ids] * 2)
909
-
910
- # 5. Prepare timesteps
911
- sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
912
- image_seq_len = latents.shape[1]
913
- mu = calculate_shift(
914
- image_seq_len,
915
- self.scheduler.config.base_image_seq_len,
916
- self.scheduler.config.max_image_seq_len,
917
- self.scheduler.config.base_shift,
918
- self.scheduler.config.max_shift,
919
- )
920
- timesteps, num_inference_steps = retrieve_timesteps(
921
- self.scheduler,
922
- num_inference_steps,
923
- device,
924
- timesteps,
925
- sigmas,
926
- mu=mu,
927
- )
928
-
929
- num_warmup_steps = max(
930
- len(timesteps) - num_inference_steps * self.scheduler.order, 0
931
- )
932
- self._num_timesteps = len(timesteps)
933
-
934
- # 6. Denoising loop
935
- with self.progress_bar(total=num_inference_steps) as progress_bar:
936
- for i, t in enumerate(timesteps):
937
- if self.interrupt:
938
- continue
939
-
940
- latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
941
-
942
- # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
943
- timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
944
-
945
- # handle guidance
946
- if self.transformer.config.guidance_embeds:
947
- guidance = torch.tensor([guidance_scale], device=device)
948
- guidance = guidance.expand(latent_model_input.shape[0])
949
- else:
950
- guidance = None
951
-
952
- # controlnet
953
- (
954
- controlnet_block_samples,
955
- controlnet_single_block_samples,
956
- ) = self.controlnet(
957
- hidden_states=latent_model_input,
958
- controlnet_cond=control_image,
959
- conditioning_scale=controlnet_conditioning_scale,
960
- timestep=timestep / 1000,
961
- guidance=guidance,
962
- pooled_projections=pooled_prompt_embeds,
963
- encoder_hidden_states=prompt_embeds,
964
- txt_ids=text_ids,
965
- img_ids=latent_image_ids,
966
- joint_attention_kwargs=self.joint_attention_kwargs,
967
- return_dict=False,
968
- )
969
-
970
- noise_pred = self.transformer(
971
- hidden_states=latent_model_input,
972
- # 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)
973
- timestep=timestep / 1000,
974
- guidance=guidance,
975
- pooled_projections=pooled_prompt_embeds,
976
- encoder_hidden_states=prompt_embeds,
977
- controlnet_block_samples=[
978
- sample.to(dtype=self.transformer.dtype)
979
- for sample in controlnet_block_samples
980
- ],
981
- controlnet_single_block_samples=[
982
- sample.to(dtype=self.transformer.dtype)
983
- for sample in controlnet_single_block_samples
984
- ] if controlnet_single_block_samples is not None else controlnet_single_block_samples,
985
- txt_ids=text_ids,
986
- img_ids=latent_image_ids,
987
- joint_attention_kwargs=self.joint_attention_kwargs,
988
- return_dict=False,
989
- )[0]
990
-
991
- # 在生成循环中
992
- if self.do_classifier_free_guidance:
993
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
994
- noise_pred = noise_pred_uncond + true_guidance_scale * (noise_pred_text - noise_pred_uncond)
995
-
996
- # compute the previous noisy sample x_t -> x_t-1
997
- latents_dtype = latents.dtype
998
- latents = self.scheduler.step(
999
- noise_pred, t, latents, return_dict=False
1000
- )[0]
1001
-
1002
- if latents.dtype != latents_dtype:
1003
- if torch.backends.mps.is_available():
1004
- # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
1005
- latents = latents.to(latents_dtype)
1006
-
1007
- if callback_on_step_end is not None:
1008
- callback_kwargs = {}
1009
- for k in callback_on_step_end_tensor_inputs:
1010
- callback_kwargs[k] = locals()[k]
1011
- callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
1012
-
1013
- latents = callback_outputs.pop("latents", latents)
1014
- prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
1015
-
1016
- # call the callback, if provided
1017
- if i == len(timesteps) - 1 or (
1018
- (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
1019
- ):
1020
- progress_bar.update()
1021
-
1022
- if XLA_AVAILABLE:
1023
- xm.mark_step()
1024
-
1025
- if output_type == "latent":
1026
- image = latents
1027
-
1028
- else:
1029
- latents = self._unpack_latents(
1030
- latents, height, width, self.vae_scale_factor
1031
- )
1032
- latents = (
1033
- latents / self.vae.config.scaling_factor
1034
- ) + self.vae.config.shift_factor
1035
- latents = latents.to(self.vae.dtype)
1036
-
1037
- image = self.vae.decode(latents, return_dict=False)[0]
1038
- image = self.image_processor.postprocess(image, output_type=output_type)
1039
-
1040
- # Offload all models
1041
- self.maybe_free_model_hooks()
1042
-
1043
- if not return_dict:
1044
- return (image,)
1045
-
1046
- return FluxPipelineOutput(images=image)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
options/Banner_Model/transformer_flux.py DELETED
@@ -1,525 +0,0 @@
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)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
options/Video_model/__pycache__/Model.cpython-310.pyc CHANGED
Binary files a/options/Video_model/__pycache__/Model.cpython-310.pyc and b/options/Video_model/__pycache__/Model.cpython-310.pyc differ
 
options/Video_model/__pycache__/__init__.cpython-310.pyc CHANGED
Binary files a/options/Video_model/__pycache__/__init__.cpython-310.pyc and b/options/Video_model/__pycache__/__init__.cpython-310.pyc differ
 
options/__pycache__/Banner.cpython-310.pyc CHANGED
Binary files a/options/__pycache__/Banner.cpython-310.pyc and b/options/__pycache__/Banner.cpython-310.pyc differ
 
options/__pycache__/Video.cpython-310.pyc CHANGED
Binary files a/options/__pycache__/Video.cpython-310.pyc and b/options/__pycache__/Video.cpython-310.pyc differ
 
requirements.txt CHANGED
@@ -2,7 +2,7 @@ imageio[ffmpeg]
2
  controlnet_aux
3
  sentencepiece
4
  mediapipe
5
- # git+https://github.com/huggingface/diffusers.git
6
  torch==2.1.0
7
  torchvision==0.16.0
8
  xformers==0.0.22.post7
 
2
  controlnet_aux
3
  sentencepiece
4
  mediapipe
5
+ spaces
6
  torch==2.1.0
7
  torchvision==0.16.0
8
  xformers==0.0.22.post7