Spaces:
Running
Running
OmPrakashSingh1704
commited on
Commit
•
9bc5731
1
Parent(s):
4d52cb6
- options/Banner.py +0 -1
- options/Banner_Model/Image2Image.py +2 -2
- options/Banner_Model/Image2Image_2.py +2 -1
- options/Banner_Model/Text2Banner.py +2 -3
- options/Banner_Model/__pycache__/Image2Image.cpython-310.pyc +0 -0
- options/Banner_Model/__pycache__/Image2Image_2.cpython-310.pyc +0 -0
- options/Banner_Model/__pycache__/Text2Banner.cpython-310.pyc +0 -0
- options/Banner_Model/__pycache__/__init__.cpython-310.pyc +0 -0
- options/Banner_Model/controlnet_flux.py +0 -418
- options/Banner_Model/pipeline_flux_controlnet_inpaint.py +0 -1046
- options/Banner_Model/transformer_flux.py +0 -525
- options/Video_model/__pycache__/Model.cpython-310.pyc +0 -0
- options/Video_model/__pycache__/__init__.cpython-310.pyc +0 -0
- options/__pycache__/Banner.cpython-310.pyc +0 -0
- options/__pycache__/Video.cpython-310.pyc +0 -0
- requirements.txt +1 -1
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 |
-
|
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 |
-
|
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 |
-
|
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")
|
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 |
-
|
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
|