|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
__all__ = ['FluxTransformer2DModelWithMasking', 'CustomPipeline'] |
|
|
|
from typing import Any, Dict, List, Optional, Union |
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
from diffusers.configuration_utils import ConfigMixin, register_to_config |
|
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin |
|
from diffusers.models.attention import FeedForward |
|
from diffusers.models.attention_processor import ( |
|
Attention, |
|
apply_rope, |
|
) |
|
from diffusers.models.modeling_utils import ModelMixin |
|
from diffusers.models.normalization import ( |
|
AdaLayerNormContinuous, |
|
AdaLayerNormZero, |
|
AdaLayerNormZeroSingle, |
|
) |
|
from diffusers.utils import ( |
|
USE_PEFT_BACKEND, |
|
is_torch_version, |
|
logging, |
|
scale_lora_layers, |
|
unscale_lora_layers, |
|
) |
|
from diffusers.utils.torch_utils import maybe_allow_in_graph |
|
from diffusers.models.embeddings import ( |
|
CombinedTimestepGuidanceTextProjEmbeddings, |
|
CombinedTimestepTextProjEmbeddings, |
|
) |
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput |
|
|
|
from dataclasses import dataclass |
|
from typing import List, Union |
|
import PIL.Image |
|
from diffusers.utils import BaseOutput |
|
|
|
import inspect |
|
from functools import lru_cache |
|
from typing import Any, Callable, Dict, List, Optional, Union |
|
|
|
import numpy as np |
|
import torch |
|
from transformers import ( |
|
CLIPTextModel, |
|
CLIPTokenizer, |
|
T5EncoderModel, |
|
T5TokenizerFast, |
|
) |
|
|
|
from diffusers.image_processor import VaeImageProcessor |
|
from diffusers.loaders import SD3LoraLoaderMixin |
|
from diffusers.models.autoencoders import AutoencoderKL |
|
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler |
|
from diffusers.utils import ( |
|
USE_PEFT_BACKEND, |
|
is_torch_xla_available, |
|
logging, |
|
replace_example_docstring, |
|
scale_lora_layers, |
|
unscale_lora_layers, |
|
) |
|
from diffusers.utils.torch_utils import randn_tensor |
|
from diffusers.pipelines.pipeline_utils import DiffusionPipeline |
|
|
|
if is_torch_xla_available(): |
|
import torch_xla.core.xla_model as xm |
|
|
|
XLA_AVAILABLE = True |
|
else: |
|
XLA_AVAILABLE = False |
|
|
|
|
|
@dataclass |
|
class FluxPipelineOutput(BaseOutput): |
|
""" |
|
Output class for Stable Diffusion pipelines. |
|
|
|
Args: |
|
images (`List[PIL.Image.Image]` or `np.ndarray`) |
|
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, |
|
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. |
|
""" |
|
|
|
images: Union[List[PIL.Image.Image], np.ndarray] |
|
|
|
|
|
logger = logging.get_logger(__name__) |
|
|
|
|
|
class FluxSingleAttnProcessor2_0: |
|
r""" |
|
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). |
|
""" |
|
|
|
def __init__(self): |
|
if not hasattr(F, "scaled_dot_product_attention"): |
|
raise ImportError( |
|
"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." |
|
) |
|
|
|
def __call__( |
|
self, |
|
attn: Attention, |
|
hidden_states: torch.Tensor, |
|
encoder_hidden_states: Optional[torch.Tensor] = None, |
|
attention_mask: Optional[torch.FloatTensor] = None, |
|
image_rotary_emb: Optional[torch.Tensor] = None, |
|
) -> torch.Tensor: |
|
input_ndim = hidden_states.ndim |
|
|
|
if input_ndim == 4: |
|
batch_size, channel, height, width = hidden_states.shape |
|
hidden_states = hidden_states.view( |
|
batch_size, channel, height * width |
|
).transpose(1, 2) |
|
|
|
batch_size, _, _ = hidden_states.shape |
|
query = attn.to_q(hidden_states) |
|
key = attn.to_k(hidden_states) |
|
value = attn.to_v(hidden_states) |
|
|
|
inner_dim = key.shape[-1] |
|
head_dim = inner_dim // attn.heads |
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
|
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
|
|
|
if attn.norm_q is not None: |
|
query = attn.norm_q(query) |
|
if attn.norm_k is not None: |
|
key = attn.norm_k(key) |
|
|
|
|
|
if image_rotary_emb is not None: |
|
|
|
|
|
|
|
|
|
query, key = apply_rope(query, key, image_rotary_emb) |
|
|
|
if attention_mask is not None: |
|
attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) |
|
attention_mask = (attention_mask > 0).bool() |
|
attention_mask = attention_mask.to( |
|
device=hidden_states.device, dtype=hidden_states.dtype |
|
) |
|
|
|
|
|
|
|
hidden_states = F.scaled_dot_product_attention( |
|
query, |
|
key, |
|
value, |
|
dropout_p=0.0, |
|
is_causal=False, |
|
attn_mask=attention_mask, |
|
) |
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape( |
|
batch_size, -1, attn.heads * head_dim |
|
) |
|
hidden_states = hidden_states.to(query.dtype) |
|
|
|
if input_ndim == 4: |
|
hidden_states = hidden_states.transpose(-1, -2).reshape( |
|
batch_size, channel, height, width |
|
) |
|
|
|
return hidden_states |
|
|
|
|
|
class FluxAttnProcessor2_0: |
|
"""Attention processor used typically in processing the SD3-like self-attention projections.""" |
|
|
|
def __init__(self): |
|
if not hasattr(F, "scaled_dot_product_attention"): |
|
raise ImportError( |
|
"FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." |
|
) |
|
|
|
def __call__( |
|
self, |
|
attn: Attention, |
|
hidden_states: torch.FloatTensor, |
|
encoder_hidden_states: torch.FloatTensor = None, |
|
attention_mask: Optional[torch.FloatTensor] = None, |
|
image_rotary_emb: Optional[torch.Tensor] = None, |
|
) -> torch.FloatTensor: |
|
input_ndim = hidden_states.ndim |
|
if input_ndim == 4: |
|
batch_size, channel, height, width = hidden_states.shape |
|
hidden_states = hidden_states.view( |
|
batch_size, channel, height * width |
|
).transpose(1, 2) |
|
context_input_ndim = encoder_hidden_states.ndim |
|
if context_input_ndim == 4: |
|
batch_size, channel, height, width = encoder_hidden_states.shape |
|
encoder_hidden_states = encoder_hidden_states.view( |
|
batch_size, channel, height * width |
|
).transpose(1, 2) |
|
|
|
batch_size = encoder_hidden_states.shape[0] |
|
|
|
|
|
query = attn.to_q(hidden_states) |
|
key = attn.to_k(hidden_states) |
|
value = attn.to_v(hidden_states) |
|
|
|
inner_dim = key.shape[-1] |
|
head_dim = inner_dim // attn.heads |
|
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
|
|
|
if attn.norm_q is not None: |
|
query = attn.norm_q(query) |
|
if attn.norm_k is not None: |
|
key = attn.norm_k(key) |
|
|
|
|
|
encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) |
|
encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) |
|
encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) |
|
|
|
encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( |
|
batch_size, -1, attn.heads, head_dim |
|
).transpose(1, 2) |
|
encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( |
|
batch_size, -1, attn.heads, head_dim |
|
).transpose(1, 2) |
|
encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( |
|
batch_size, -1, attn.heads, head_dim |
|
).transpose(1, 2) |
|
|
|
if attn.norm_added_q is not None: |
|
encoder_hidden_states_query_proj = attn.norm_added_q( |
|
encoder_hidden_states_query_proj |
|
) |
|
if attn.norm_added_k is not None: |
|
encoder_hidden_states_key_proj = attn.norm_added_k( |
|
encoder_hidden_states_key_proj |
|
) |
|
|
|
|
|
query = torch.cat([encoder_hidden_states_query_proj, query], dim=2) |
|
key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) |
|
value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) |
|
|
|
if image_rotary_emb is not None: |
|
|
|
|
|
|
|
|
|
query, key = apply_rope(query, key, image_rotary_emb) |
|
|
|
if attention_mask is not None: |
|
attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) |
|
attention_mask = (attention_mask > 0).bool() |
|
attention_mask = attention_mask.to( |
|
device=hidden_states.device, dtype=hidden_states.dtype |
|
) |
|
|
|
hidden_states = F.scaled_dot_product_attention( |
|
query, |
|
key, |
|
value, |
|
dropout_p=0.0, |
|
is_causal=False, |
|
attn_mask=attention_mask, |
|
) |
|
hidden_states = hidden_states.transpose(1, 2).reshape( |
|
batch_size, -1, attn.heads * head_dim |
|
) |
|
hidden_states = hidden_states.to(query.dtype) |
|
|
|
encoder_hidden_states, hidden_states = ( |
|
hidden_states[:, : encoder_hidden_states.shape[1]], |
|
hidden_states[:, encoder_hidden_states.shape[1] :], |
|
) |
|
|
|
|
|
hidden_states = attn.to_out[0](hidden_states) |
|
|
|
hidden_states = attn.to_out[1](hidden_states) |
|
encoder_hidden_states = attn.to_add_out(encoder_hidden_states) |
|
|
|
if input_ndim == 4: |
|
hidden_states = hidden_states.transpose(-1, -2).reshape( |
|
batch_size, channel, height, width |
|
) |
|
if context_input_ndim == 4: |
|
encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape( |
|
batch_size, channel, height, width |
|
) |
|
|
|
return hidden_states, encoder_hidden_states |
|
|
|
|
|
|
|
def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: |
|
assert dim % 2 == 0, "The dimension must be even." |
|
|
|
scale = ( |
|
torch.arange( |
|
0, |
|
dim, |
|
2, |
|
dtype=torch.float64, |
|
device=pos.device, |
|
) |
|
/ dim |
|
) |
|
omega = 1.0 / (theta**scale) |
|
|
|
batch_size, seq_length = pos.shape |
|
out = torch.einsum("...n,d->...nd", pos, omega) |
|
cos_out = torch.cos(out) |
|
sin_out = torch.sin(out) |
|
|
|
stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1) |
|
out = stacked_out.view(batch_size, -1, dim // 2, 2, 2) |
|
return out.float() |
|
|
|
|
|
|
|
class EmbedND(nn.Module): |
|
def __init__(self, dim: int, theta: int, axes_dim: List[int]): |
|
super().__init__() |
|
self.dim = dim |
|
self.theta = theta |
|
self.axes_dim = axes_dim |
|
|
|
def forward(self, ids: torch.Tensor) -> torch.Tensor: |
|
n_axes = ids.shape[-1] |
|
emb = torch.cat( |
|
[rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], |
|
dim=-3, |
|
) |
|
|
|
return emb.unsqueeze(1) |
|
|
|
|
|
def expand_flux_attention_mask( |
|
hidden_states: torch.Tensor, |
|
attn_mask: torch.Tensor, |
|
) -> torch.Tensor: |
|
""" |
|
Expand a mask so that the image is included. |
|
""" |
|
bsz = attn_mask.shape[0] |
|
assert bsz == hidden_states.shape[0] |
|
residual_seq_len = hidden_states.shape[1] |
|
mask_seq_len = attn_mask.shape[1] |
|
|
|
expanded_mask = torch.ones(bsz, residual_seq_len) |
|
expanded_mask[:, :mask_seq_len] = attn_mask |
|
|
|
return expanded_mask |
|
|
|
|
|
@maybe_allow_in_graph |
|
class FluxSingleTransformerBlock(nn.Module): |
|
r""" |
|
A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3. |
|
|
|
Reference: https://arxiv.org/abs/2403.03206 |
|
|
|
Parameters: |
|
dim (`int`): The number of channels in the input and output. |
|
num_attention_heads (`int`): The number of heads to use for multi-head attention. |
|
attention_head_dim (`int`): The number of channels in each head. |
|
context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the |
|
processing of `context` conditions. |
|
""" |
|
|
|
def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0): |
|
super().__init__() |
|
self.mlp_hidden_dim = int(dim * mlp_ratio) |
|
|
|
self.norm = AdaLayerNormZeroSingle(dim) |
|
self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim) |
|
self.act_mlp = nn.GELU(approximate="tanh") |
|
self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim) |
|
|
|
processor = FluxSingleAttnProcessor2_0() |
|
self.attn = Attention( |
|
query_dim=dim, |
|
cross_attention_dim=None, |
|
dim_head=attention_head_dim, |
|
heads=num_attention_heads, |
|
out_dim=dim, |
|
bias=True, |
|
processor=processor, |
|
qk_norm="rms_norm", |
|
eps=1e-6, |
|
pre_only=True, |
|
) |
|
|
|
def forward( |
|
self, |
|
hidden_states: torch.FloatTensor, |
|
temb: torch.FloatTensor, |
|
image_rotary_emb=None, |
|
attention_mask: Optional[torch.Tensor] = None, |
|
): |
|
residual = hidden_states |
|
norm_hidden_states, gate = self.norm(hidden_states, emb=temb) |
|
mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) |
|
|
|
if attention_mask is not None: |
|
attention_mask = expand_flux_attention_mask( |
|
hidden_states, |
|
attention_mask, |
|
) |
|
|
|
attn_output = self.attn( |
|
hidden_states=norm_hidden_states, |
|
image_rotary_emb=image_rotary_emb, |
|
attention_mask=attention_mask, |
|
) |
|
|
|
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) |
|
gate = gate.unsqueeze(1) |
|
hidden_states = gate * self.proj_out(hidden_states) |
|
hidden_states = residual + hidden_states |
|
|
|
return hidden_states |
|
|
|
|
|
@maybe_allow_in_graph |
|
class FluxTransformerBlock(nn.Module): |
|
r""" |
|
A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3. |
|
|
|
Reference: https://arxiv.org/abs/2403.03206 |
|
|
|
Parameters: |
|
dim (`int`): The number of channels in the input and output. |
|
num_attention_heads (`int`): The number of heads to use for multi-head attention. |
|
attention_head_dim (`int`): The number of channels in each head. |
|
context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the |
|
processing of `context` conditions. |
|
""" |
|
|
|
def __init__( |
|
self, dim, num_attention_heads, attention_head_dim, qk_norm="rms_norm", eps=1e-6 |
|
): |
|
super().__init__() |
|
|
|
self.norm1 = AdaLayerNormZero(dim) |
|
|
|
self.norm1_context = AdaLayerNormZero(dim) |
|
|
|
if hasattr(F, "scaled_dot_product_attention"): |
|
processor = FluxAttnProcessor2_0() |
|
else: |
|
raise ValueError( |
|
"The current PyTorch version does not support the `scaled_dot_product_attention` function." |
|
) |
|
self.attn = Attention( |
|
query_dim=dim, |
|
cross_attention_dim=None, |
|
added_kv_proj_dim=dim, |
|
dim_head=attention_head_dim, |
|
heads=num_attention_heads, |
|
out_dim=dim, |
|
context_pre_only=False, |
|
bias=True, |
|
processor=processor, |
|
qk_norm=qk_norm, |
|
eps=eps, |
|
) |
|
|
|
self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
|
self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") |
|
|
|
self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
|
self.ff_context = FeedForward( |
|
dim=dim, dim_out=dim, activation_fn="gelu-approximate" |
|
) |
|
|
|
|
|
self._chunk_size = None |
|
self._chunk_dim = 0 |
|
|
|
def forward( |
|
self, |
|
hidden_states: torch.FloatTensor, |
|
encoder_hidden_states: torch.FloatTensor, |
|
temb: torch.FloatTensor, |
|
image_rotary_emb=None, |
|
attention_mask: Optional[torch.Tensor] = None, |
|
): |
|
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( |
|
hidden_states, emb=temb |
|
) |
|
|
|
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = ( |
|
self.norm1_context(encoder_hidden_states, emb=temb) |
|
) |
|
|
|
if attention_mask is not None: |
|
attention_mask = expand_flux_attention_mask( |
|
torch.cat([encoder_hidden_states, hidden_states], dim=1), |
|
attention_mask, |
|
) |
|
|
|
|
|
attn_output, context_attn_output = self.attn( |
|
hidden_states=norm_hidden_states, |
|
encoder_hidden_states=norm_encoder_hidden_states, |
|
image_rotary_emb=image_rotary_emb, |
|
attention_mask=attention_mask, |
|
) |
|
|
|
|
|
attn_output = gate_msa.unsqueeze(1) * attn_output |
|
hidden_states = hidden_states + attn_output |
|
|
|
norm_hidden_states = self.norm2(hidden_states) |
|
norm_hidden_states = ( |
|
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] |
|
) |
|
|
|
ff_output = self.ff(norm_hidden_states) |
|
ff_output = gate_mlp.unsqueeze(1) * ff_output |
|
|
|
hidden_states = hidden_states + ff_output |
|
|
|
|
|
|
|
context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output |
|
encoder_hidden_states = encoder_hidden_states + context_attn_output |
|
|
|
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) |
|
norm_encoder_hidden_states = ( |
|
norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) |
|
+ c_shift_mlp[:, None] |
|
) |
|
|
|
context_ff_output = self.ff_context(norm_encoder_hidden_states) |
|
encoder_hidden_states = ( |
|
encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output |
|
) |
|
|
|
return encoder_hidden_states, hidden_states |
|
|
|
|
|
class FluxTransformer2DModelWithMasking( |
|
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin |
|
): |
|
""" |
|
The Transformer model introduced in Flux. |
|
|
|
Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ |
|
|
|
Parameters: |
|
patch_size (`int`): Patch size to turn the input data into small patches. |
|
in_channels (`int`, *optional*, defaults to 16): The number of channels in the input. |
|
num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use. |
|
num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use. |
|
attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head. |
|
num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention. |
|
joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. |
|
pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`. |
|
guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings. |
|
""" |
|
|
|
_supports_gradient_checkpointing = True |
|
|
|
@register_to_config |
|
def __init__( |
|
self, |
|
patch_size: int = 1, |
|
in_channels: int = 64, |
|
num_layers: int = 19, |
|
num_single_layers: int = 38, |
|
attention_head_dim: int = 128, |
|
num_attention_heads: int = 24, |
|
joint_attention_dim: int = 4096, |
|
pooled_projection_dim: int = 768, |
|
guidance_embeds: bool = False, |
|
axes_dims_rope: List[int] = [16, 56, 56], |
|
): |
|
super().__init__() |
|
self.out_channels = in_channels |
|
self.inner_dim = ( |
|
self.config.num_attention_heads * self.config.attention_head_dim |
|
) |
|
|
|
self.pos_embed = EmbedND( |
|
dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope |
|
) |
|
text_time_guidance_cls = ( |
|
CombinedTimestepGuidanceTextProjEmbeddings |
|
if guidance_embeds |
|
else CombinedTimestepTextProjEmbeddings |
|
) |
|
self.time_text_embed = text_time_guidance_cls( |
|
embedding_dim=self.inner_dim, |
|
pooled_projection_dim=self.config.pooled_projection_dim, |
|
) |
|
|
|
self.context_embedder = nn.Linear( |
|
self.config.joint_attention_dim, self.inner_dim |
|
) |
|
self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim) |
|
|
|
self.transformer_blocks = nn.ModuleList( |
|
[ |
|
FluxTransformerBlock( |
|
dim=self.inner_dim, |
|
num_attention_heads=self.config.num_attention_heads, |
|
attention_head_dim=self.config.attention_head_dim, |
|
) |
|
for i in range(self.config.num_layers) |
|
] |
|
) |
|
|
|
self.single_transformer_blocks = nn.ModuleList( |
|
[ |
|
FluxSingleTransformerBlock( |
|
dim=self.inner_dim, |
|
num_attention_heads=self.config.num_attention_heads, |
|
attention_head_dim=self.config.attention_head_dim, |
|
) |
|
for i in range(self.config.num_single_layers) |
|
] |
|
) |
|
|
|
self.norm_out = AdaLayerNormContinuous( |
|
self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6 |
|
) |
|
self.proj_out = nn.Linear( |
|
self.inner_dim, patch_size * patch_size * self.out_channels, bias=True |
|
) |
|
|
|
self.gradient_checkpointing = False |
|
|
|
def _set_gradient_checkpointing(self, module, value=False): |
|
if hasattr(module, "gradient_checkpointing"): |
|
module.gradient_checkpointing = value |
|
|
|
def forward( |
|
self, |
|
hidden_states: torch.Tensor, |
|
encoder_hidden_states: torch.Tensor = None, |
|
pooled_projections: torch.Tensor = None, |
|
timestep: torch.LongTensor = None, |
|
img_ids: torch.Tensor = None, |
|
txt_ids: torch.Tensor = None, |
|
guidance: torch.Tensor = None, |
|
joint_attention_kwargs: Optional[Dict[str, Any]] = None, |
|
return_dict: bool = True, |
|
attention_mask: Optional[torch.Tensor] = None, |
|
) -> Union[torch.FloatTensor, Transformer2DModelOutput]: |
|
""" |
|
The [`FluxTransformer2DModelWithMasking`] forward method. |
|
|
|
Args: |
|
hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`): |
|
Input `hidden_states`. |
|
encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`): |
|
Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. |
|
pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected |
|
from the embeddings of input conditions. |
|
timestep ( `torch.LongTensor`): |
|
Used to indicate denoising step. |
|
block_controlnet_hidden_states: (`list` of `torch.Tensor`): |
|
A list of tensors that if specified are added to the residuals of transformer blocks. |
|
joint_attention_kwargs (`dict`, *optional*): |
|
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under |
|
`self.processor` in |
|
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). |
|
return_dict (`bool`, *optional*, defaults to `True`): |
|
Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain |
|
tuple. |
|
|
|
Returns: |
|
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a |
|
`tuple` where the first element is the sample tensor. |
|
""" |
|
if joint_attention_kwargs is not None: |
|
joint_attention_kwargs = joint_attention_kwargs.copy() |
|
lora_scale = joint_attention_kwargs.pop("scale", 1.0) |
|
else: |
|
lora_scale = 1.0 |
|
|
|
if USE_PEFT_BACKEND: |
|
|
|
scale_lora_layers(self, lora_scale) |
|
else: |
|
if ( |
|
joint_attention_kwargs is not None |
|
and joint_attention_kwargs.get("scale", None) is not None |
|
): |
|
logger.warning( |
|
"Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." |
|
) |
|
hidden_states = self.x_embedder(hidden_states) |
|
|
|
timestep = timestep.to(hidden_states.dtype) * 1000 |
|
if guidance is not None: |
|
guidance = guidance.to(hidden_states.dtype) * 1000 |
|
else: |
|
guidance = None |
|
temb = ( |
|
self.time_text_embed(timestep, pooled_projections) |
|
if guidance is None |
|
else self.time_text_embed(timestep, guidance, pooled_projections) |
|
) |
|
encoder_hidden_states = self.context_embedder(encoder_hidden_states) |
|
|
|
ids = torch.cat((txt_ids, img_ids), dim=1) |
|
image_rotary_emb = self.pos_embed(ids) |
|
|
|
for index_block, block in enumerate(self.transformer_blocks): |
|
if self.training and self.gradient_checkpointing: |
|
|
|
def create_custom_forward(module, return_dict=None): |
|
def custom_forward(*inputs): |
|
if return_dict is not None: |
|
return module(*inputs, return_dict=return_dict) |
|
else: |
|
return module(*inputs) |
|
|
|
return custom_forward |
|
|
|
ckpt_kwargs: Dict[str, Any] = ( |
|
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} |
|
) |
|
encoder_hidden_states, hidden_states = ( |
|
torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(block), |
|
hidden_states, |
|
encoder_hidden_states, |
|
temb, |
|
image_rotary_emb, |
|
attention_mask, |
|
**ckpt_kwargs, |
|
) |
|
) |
|
|
|
else: |
|
encoder_hidden_states, hidden_states = block( |
|
hidden_states=hidden_states, |
|
encoder_hidden_states=encoder_hidden_states, |
|
temb=temb, |
|
image_rotary_emb=image_rotary_emb, |
|
attention_mask=attention_mask, |
|
) |
|
|
|
|
|
|
|
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) |
|
|
|
for index_block, block in enumerate(self.single_transformer_blocks): |
|
if self.training and self.gradient_checkpointing: |
|
|
|
def create_custom_forward(module, return_dict=None): |
|
def custom_forward(*inputs): |
|
if return_dict is not None: |
|
return module(*inputs, return_dict=return_dict) |
|
else: |
|
return module(*inputs) |
|
|
|
return custom_forward |
|
|
|
ckpt_kwargs: Dict[str, Any] = ( |
|
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} |
|
) |
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(block), |
|
hidden_states, |
|
temb, |
|
image_rotary_emb, |
|
attention_mask, |
|
**ckpt_kwargs, |
|
) |
|
|
|
else: |
|
hidden_states = block( |
|
hidden_states=hidden_states, |
|
temb=temb, |
|
image_rotary_emb=image_rotary_emb, |
|
attention_mask=attention_mask, |
|
) |
|
|
|
hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] |
|
|
|
hidden_states = self.norm_out(hidden_states, temb) |
|
output = self.proj_out(hidden_states) |
|
|
|
if USE_PEFT_BACKEND: |
|
|
|
unscale_lora_layers(self, lora_scale) |
|
|
|
if not return_dict: |
|
return (output,) |
|
|
|
return Transformer2DModelOutput(sample=output) |
|
|
|
|
|
if __name__ == "__main__": |
|
dtype = torch.bfloat16 |
|
bsz = 2 |
|
img = torch.rand((bsz, 16, 64, 64)).to("cuda", dtype=dtype) |
|
timestep = torch.tensor([0.5, 0.5]).to("cuda", dtype=torch.float32) |
|
pooled = torch.rand(bsz, 768).to("cuda", dtype=dtype) |
|
text = torch.rand((bsz, 512, 4096)).to("cuda", dtype=dtype) |
|
attn_mask = torch.tensor([[1.0] * 384 + [0.0] * 128] * bsz).to( |
|
"cuda", dtype=dtype |
|
) |
|
|
|
def _pack_latents(latents, batch_size, num_channels_latents, height, width): |
|
latents = latents.view( |
|
batch_size, num_channels_latents, height // 2, 2, width // 2, 2 |
|
) |
|
latents = latents.permute(0, 2, 4, 1, 3, 5) |
|
latents = latents.reshape( |
|
batch_size, (height // 2) * (width // 2), num_channels_latents * 4 |
|
) |
|
|
|
return latents |
|
|
|
def _prepare_latent_image_ids( |
|
batch_size, height, width, device="cuda", dtype=dtype |
|
): |
|
latent_image_ids = torch.zeros(height // 2, width // 2, 3) |
|
latent_image_ids[..., 1] = ( |
|
latent_image_ids[..., 1] + torch.arange(height // 2)[:, None] |
|
) |
|
latent_image_ids[..., 2] = ( |
|
latent_image_ids[..., 2] + torch.arange(width // 2)[None, :] |
|
) |
|
|
|
latent_image_id_height, latent_image_id_width, latent_image_id_channels = ( |
|
latent_image_ids.shape |
|
) |
|
|
|
latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1) |
|
latent_image_ids = latent_image_ids.reshape( |
|
batch_size, |
|
latent_image_id_height * latent_image_id_width, |
|
latent_image_id_channels, |
|
) |
|
|
|
return latent_image_ids.to(device=device, dtype=dtype) |
|
|
|
txt_ids = torch.zeros(bsz, text.shape[1], 3).to(device="cuda", dtype=dtype) |
|
|
|
vae_scale_factor = 16 |
|
height = 2 * (int(512) // vae_scale_factor) |
|
width = 2 * (int(512) // vae_scale_factor) |
|
img_ids = _prepare_latent_image_ids(bsz, height, width) |
|
img = _pack_latents(img, img.shape[0], 16, height, width) |
|
|
|
|
|
transformer = FluxTransformer2DModelWithMasking.from_config( |
|
{ |
|
"attention_head_dim": 128, |
|
"guidance_embeds": True, |
|
"in_channels": 64, |
|
"joint_attention_dim": 4096, |
|
"num_attention_heads": 24, |
|
"num_layers": 4, |
|
"num_single_layers": 8, |
|
"patch_size": 1, |
|
"pooled_projection_dim": 768, |
|
} |
|
).to("cuda", dtype=dtype) |
|
|
|
guidance = torch.tensor([2.0], device="cuda") |
|
guidance = guidance.expand(bsz) |
|
|
|
with torch.no_grad(): |
|
no_mask = transformer( |
|
img, |
|
encoder_hidden_states=text, |
|
pooled_projections=pooled, |
|
timestep=timestep, |
|
img_ids=img_ids, |
|
txt_ids=txt_ids, |
|
guidance=guidance, |
|
) |
|
mask = transformer( |
|
img, |
|
encoder_hidden_states=text, |
|
pooled_projections=pooled, |
|
timestep=timestep, |
|
img_ids=img_ids, |
|
txt_ids=txt_ids, |
|
guidance=guidance, |
|
attention_mask=attn_mask, |
|
) |
|
|
|
assert torch.allclose(no_mask.sample, mask.sample) is False |
|
print("Attention masking test ran OK. Differences in output were detected.") |
|
|
|
logger = logging.get_logger(__name__) |
|
|
|
EXAMPLE_DOC_STRING = """ |
|
Examples: |
|
```py |
|
>>> import torch |
|
>>> from diffusers import FluxPipeline |
|
|
|
>>> pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) |
|
>>> pipe.to("cuda") |
|
>>> prompt = "A cat holding a sign that says hello world" |
|
>>> # Depending on the variant being used, the pipeline call will slightly vary. |
|
>>> # Refer to the pipeline documentation for more details. |
|
>>> image = pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0] |
|
>>> image.save("flux.png") |
|
``` |
|
""" |
|
|
|
|
|
def calculate_shift( |
|
image_seq_len, |
|
base_seq_len: int = 256, |
|
max_seq_len: int = 4096, |
|
base_shift: float = 0.5, |
|
max_shift: float = 1.16, |
|
): |
|
m = (max_shift - base_shift) / (max_seq_len - base_seq_len) |
|
b = base_shift - m * base_seq_len |
|
mu = image_seq_len * m + b |
|
return mu |
|
|
|
|
|
|
|
def retrieve_timesteps( |
|
scheduler, |
|
num_inference_steps: Optional[int] = None, |
|
device: Optional[Union[str, torch.device]] = None, |
|
timesteps: Optional[List[int]] = None, |
|
sigmas: Optional[List[float]] = None, |
|
**kwargs, |
|
): |
|
""" |
|
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles |
|
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. |
|
|
|
Args: |
|
scheduler (`SchedulerMixin`): |
|
The scheduler to get timesteps from. |
|
num_inference_steps (`int`): |
|
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` |
|
must be `None`. |
|
device (`str` or `torch.device`, *optional*): |
|
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. |
|
timesteps (`List[int]`, *optional*): |
|
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, |
|
`num_inference_steps` and `sigmas` must be `None`. |
|
sigmas (`List[float]`, *optional*): |
|
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, |
|
`num_inference_steps` and `timesteps` must be `None`. |
|
|
|
Returns: |
|
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the |
|
second element is the number of inference steps. |
|
""" |
|
if timesteps is not None and sigmas is not None: |
|
raise ValueError( |
|
"Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values" |
|
) |
|
if timesteps is not None: |
|
accepts_timesteps = "timesteps" in set( |
|
inspect.signature(scheduler.set_timesteps).parameters.keys() |
|
) |
|
if not accepts_timesteps: |
|
raise ValueError( |
|
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" |
|
f" timestep schedules. Please check whether you are using the correct scheduler." |
|
) |
|
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) |
|
timesteps = scheduler.timesteps |
|
num_inference_steps = len(timesteps) |
|
elif sigmas is not None: |
|
accept_sigmas = "sigmas" in set( |
|
inspect.signature(scheduler.set_timesteps).parameters.keys() |
|
) |
|
if not accept_sigmas: |
|
raise ValueError( |
|
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" |
|
f" sigmas schedules. Please check whether you are using the correct scheduler." |
|
) |
|
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) |
|
timesteps = scheduler.timesteps |
|
num_inference_steps = len(timesteps) |
|
else: |
|
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) |
|
timesteps = scheduler.timesteps |
|
return timesteps, num_inference_steps |
|
|
|
|
|
class CustomPipeline(DiffusionPipeline, SD3LoraLoaderMixin): |
|
r""" |
|
The Flux pipeline for text-to-image generation. |
|
|
|
Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ |
|
|
|
Args: |
|
transformer ([`FluxTransformer2DModelWithMasking`]): |
|
Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. |
|
scheduler ([`FlowMatchEulerDiscreteScheduler`]): |
|
A scheduler to be used in combination with `transformer` to denoise the encoded image latents. |
|
vae ([`AutoencoderKL`]): |
|
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. |
|
text_encoder ([`CLIPTextModelWithProjection`]): |
|
[CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), |
|
specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, |
|
with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` |
|
as its dimension. |
|
text_encoder_2 ([`CLIPTextModelWithProjection`]): |
|
[CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), |
|
specifically the |
|
[laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k) |
|
variant. |
|
tokenizer (`CLIPTokenizer`): |
|
Tokenizer of class |
|
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). |
|
tokenizer_2 (`CLIPTokenizer`): |
|
Second Tokenizer of class |
|
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). |
|
""" |
|
|
|
model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" |
|
_optional_components = [] |
|
_callback_tensor_inputs = ["latents", "prompt_embeds"] |
|
|
|
def __init__( |
|
self, |
|
scheduler: FlowMatchEulerDiscreteScheduler, |
|
vae: AutoencoderKL, |
|
text_encoder: CLIPTextModel, |
|
tokenizer: CLIPTokenizer, |
|
text_encoder_2: T5EncoderModel, |
|
tokenizer_2: T5TokenizerFast, |
|
transformer: FluxTransformer2DModelWithMasking, |
|
): |
|
super().__init__() |
|
|
|
self.register_modules( |
|
vae=vae, |
|
text_encoder=text_encoder, |
|
text_encoder_2=text_encoder_2, |
|
tokenizer=tokenizer, |
|
tokenizer_2=tokenizer_2, |
|
transformer=transformer, |
|
scheduler=scheduler, |
|
) |
|
self.vae_scale_factor = ( |
|
2 ** (len(self.vae.config.block_out_channels)) |
|
if hasattr(self, "vae") and self.vae is not None |
|
else 16 |
|
) |
|
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) |
|
self.tokenizer_max_length = ( |
|
self.tokenizer.model_max_length |
|
if hasattr(self, "tokenizer") and self.tokenizer is not None |
|
else 77 |
|
) |
|
self.default_sample_size = 64 |
|
|
|
def _get_t5_prompt_embeds( |
|
self, |
|
prompt: Union[str, List[str]] = None, |
|
num_images_per_prompt: int = 1, |
|
max_sequence_length: int = 512, |
|
device: Optional[torch.device] = None, |
|
dtype: Optional[torch.dtype] = None, |
|
): |
|
device = device or self._execution_device |
|
dtype = dtype or self.text_encoder.dtype |
|
|
|
prompt = [prompt] if isinstance(prompt, str) else prompt |
|
batch_size = len(prompt) |
|
|
|
text_inputs = self.tokenizer_2( |
|
prompt, |
|
padding="max_length", |
|
max_length=max_sequence_length, |
|
truncation=True, |
|
return_length=False, |
|
return_overflowing_tokens=False, |
|
return_tensors="pt", |
|
) |
|
prompt_attention_mask = text_inputs.attention_mask |
|
text_input_ids = text_inputs.input_ids |
|
untruncated_ids = self.tokenizer_2( |
|
prompt, padding="longest", return_tensors="pt" |
|
).input_ids |
|
|
|
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): |
|
removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1]) |
|
logger.warning( |
|
"The following part of your input was truncated because `max_sequence_length` is set to " |
|
f" {max_sequence_length} tokens: {removed_text}" |
|
) |
|
|
|
prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] |
|
|
|
dtype = self.text_encoder_2.dtype |
|
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) |
|
|
|
_, seq_len, _ = prompt_embeds.shape |
|
|
|
|
|
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) |
|
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) |
|
|
|
return prompt_embeds, prompt_attention_mask |
|
|
|
def _get_clip_prompt_embeds( |
|
self, |
|
prompt: Union[str, List[str]], |
|
num_images_per_prompt: int = 1, |
|
device: Optional[torch.device] = None, |
|
): |
|
device = device or self._execution_device |
|
|
|
prompt = [prompt] if isinstance(prompt, str) else prompt |
|
batch_size = len(prompt) |
|
|
|
text_inputs = self.tokenizer( |
|
prompt, |
|
padding="max_length", |
|
max_length=self.tokenizer_max_length, |
|
truncation=True, |
|
return_overflowing_tokens=False, |
|
return_length=False, |
|
return_tensors="pt", |
|
) |
|
|
|
text_input_ids = text_inputs.input_ids |
|
untruncated_ids = self.tokenizer( |
|
prompt, padding="longest", return_tensors="pt" |
|
).input_ids |
|
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( |
|
text_input_ids, untruncated_ids |
|
): |
|
removed_text = self.tokenizer.batch_decode( |
|
untruncated_ids[:, self.tokenizer_max_length - 1 : -1] |
|
) |
|
logger.warning( |
|
"The following part of your input was truncated because CLIP can only handle sequences up to" |
|
f" {self.tokenizer_max_length} tokens: {removed_text}" |
|
) |
|
prompt_embeds = self.text_encoder( |
|
text_input_ids.to(device), output_hidden_states=False |
|
) |
|
|
|
|
|
prompt_embeds = prompt_embeds.pooler_output |
|
prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) |
|
|
|
|
|
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) |
|
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) |
|
|
|
return prompt_embeds |
|
|
|
@lru_cache(maxsize=128) |
|
def encode_prompt( |
|
self, |
|
prompt: Union[str, List[str]], |
|
prompt_2: Union[str, List[str]], |
|
device: Optional[torch.device] = None, |
|
num_images_per_prompt: int = 1, |
|
prompt_embeds: Optional[torch.FloatTensor] = None, |
|
pooled_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
max_sequence_length: int = 512, |
|
lora_scale: Optional[float] = None, |
|
): |
|
r""" |
|
|
|
Args: |
|
prompt (`str` or `List[str]`, *optional*): |
|
prompt to be encoded |
|
prompt_2 (`str` or `List[str]`, *optional*): |
|
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is |
|
used in all text-encoders |
|
device: (`torch.device`): |
|
torch device |
|
num_images_per_prompt (`int`): |
|
number of images that should be generated per prompt |
|
prompt_embeds (`torch.FloatTensor`, *optional*): |
|
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not |
|
provided, text embeddings will be generated from `prompt` input argument. |
|
pooled_prompt_embeds (`torch.FloatTensor`, *optional*): |
|
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. |
|
If not provided, pooled text embeddings will be generated from `prompt` input argument. |
|
clip_skip (`int`, *optional*): |
|
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that |
|
the output of the pre-final layer will be used for computing the prompt embeddings. |
|
lora_scale (`float`, *optional*): |
|
A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. |
|
""" |
|
device = device or self._execution_device |
|
|
|
|
|
|
|
if lora_scale is not None and isinstance(self, SD3LoraLoaderMixin): |
|
self._lora_scale = lora_scale |
|
|
|
|
|
if self.text_encoder is not None and USE_PEFT_BACKEND: |
|
scale_lora_layers(self.text_encoder, lora_scale) |
|
if self.text_encoder_2 is not None and USE_PEFT_BACKEND: |
|
scale_lora_layers(self.text_encoder_2, lora_scale) |
|
|
|
prompt = [prompt] if isinstance(prompt, str) else prompt |
|
if prompt is not None: |
|
batch_size = len(prompt) |
|
else: |
|
batch_size = prompt_embeds.shape[0] |
|
|
|
prompt_attention_mask = None |
|
if prompt_embeds is None: |
|
prompt_2 = prompt_2 or prompt |
|
prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 |
|
|
|
|
|
pooled_prompt_embeds = self._get_clip_prompt_embeds( |
|
prompt=prompt, |
|
device=device, |
|
num_images_per_prompt=num_images_per_prompt, |
|
) |
|
prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds( |
|
prompt=prompt_2, |
|
num_images_per_prompt=num_images_per_prompt, |
|
max_sequence_length=max_sequence_length, |
|
device=device, |
|
) |
|
|
|
if self.text_encoder is not None: |
|
if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND: |
|
|
|
unscale_lora_layers(self.text_encoder, lora_scale) |
|
|
|
if self.text_encoder_2 is not None: |
|
if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND: |
|
|
|
unscale_lora_layers(self.text_encoder_2, lora_scale) |
|
|
|
dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype |
|
text_ids = torch.zeros(batch_size, prompt_embeds.shape[1], 3).to(device=device, dtype=dtype) |
|
text_ids = text_ids.repeat(num_images_per_prompt, 1, 1) |
|
|
|
return prompt_embeds, pooled_prompt_embeds, text_ids, prompt_attention_mask |
|
|
|
def check_inputs( |
|
self, |
|
prompt, |
|
prompt_2, |
|
height, |
|
width, |
|
prompt_embeds=None, |
|
pooled_prompt_embeds=None, |
|
callback_on_step_end_tensor_inputs=None, |
|
max_sequence_length=None, |
|
): |
|
if height % 8 != 0 or width % 8 != 0: |
|
raise ValueError( |
|
f"`height` and `width` have to be divisible by 8 but are {height} and {width}." |
|
) |
|
|
|
if callback_on_step_end_tensor_inputs is not None and not all( |
|
k in self._callback_tensor_inputs |
|
for k in callback_on_step_end_tensor_inputs |
|
): |
|
raise ValueError( |
|
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]}" |
|
) |
|
|
|
if prompt is not None and prompt_embeds is not None: |
|
raise ValueError( |
|
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" |
|
" only forward one of the two." |
|
) |
|
elif prompt_2 is not None and prompt_embeds is not None: |
|
raise ValueError( |
|
f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" |
|
" only forward one of the two." |
|
) |
|
elif prompt is None and prompt_embeds is None: |
|
raise ValueError( |
|
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." |
|
) |
|
elif prompt is not None and ( |
|
not isinstance(prompt, str) and not isinstance(prompt, list) |
|
): |
|
raise ValueError( |
|
f"`prompt` has to be of type `str` or `list` but is {type(prompt)}" |
|
) |
|
elif prompt_2 is not None and ( |
|
not isinstance(prompt_2, str) and not isinstance(prompt_2, list) |
|
): |
|
raise ValueError( |
|
f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}" |
|
) |
|
|
|
if prompt_embeds is not None and pooled_prompt_embeds is None: |
|
raise ValueError( |
|
"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`." |
|
) |
|
|
|
if max_sequence_length is not None and max_sequence_length > 512: |
|
raise ValueError( |
|
f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}" |
|
) |
|
|
|
@staticmethod |
|
def _prepare_latent_image_ids(batch_size, height, width, device, dtype): |
|
latent_image_ids = torch.zeros(height // 2, width // 2, 3) |
|
latent_image_ids[..., 1] = ( |
|
latent_image_ids[..., 1] + torch.arange(height // 2)[:, None] |
|
) |
|
latent_image_ids[..., 2] = ( |
|
latent_image_ids[..., 2] + torch.arange(width // 2)[None, :] |
|
) |
|
|
|
latent_image_id_height, latent_image_id_width, latent_image_id_channels = ( |
|
latent_image_ids.shape |
|
) |
|
|
|
latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1) |
|
latent_image_ids = latent_image_ids.reshape( |
|
batch_size, |
|
latent_image_id_height * latent_image_id_width, |
|
latent_image_id_channels, |
|
) |
|
|
|
return latent_image_ids |
|
|
|
@staticmethod |
|
def _pack_latents(latents, batch_size, num_channels_latents, height, width): |
|
latents = latents.view( |
|
batch_size, num_channels_latents, height // 2, 2, width // 2, 2 |
|
) |
|
latents = latents.permute(0, 2, 4, 1, 3, 5) |
|
latents = latents.reshape( |
|
batch_size, (height // 2) * (width // 2), num_channels_latents * 4 |
|
) |
|
|
|
return latents |
|
|
|
@staticmethod |
|
def _unpack_latents(latents, height, width, vae_scale_factor): |
|
batch_size, num_patches, channels = latents.shape |
|
|
|
height = height // vae_scale_factor |
|
width = width // vae_scale_factor |
|
|
|
latents = latents.view(batch_size, height, width, channels // 4, 2, 2) |
|
latents = latents.permute(0, 3, 1, 4, 2, 5) |
|
|
|
latents = latents.reshape( |
|
batch_size, channels // (2 * 2), height * 2, width * 2 |
|
) |
|
|
|
return latents |
|
|
|
def prepare_latents( |
|
self, |
|
batch_size, |
|
num_channels_latents, |
|
height, |
|
width, |
|
dtype, |
|
device, |
|
generator, |
|
latents=None, |
|
): |
|
height = 2 * (int(height) // self.vae_scale_factor) |
|
width = 2 * (int(width) // self.vae_scale_factor) |
|
|
|
shape = (batch_size, num_channels_latents, height, width) |
|
|
|
if latents is not None: |
|
latent_image_ids = self._prepare_latent_image_ids( |
|
batch_size, height, width, device, dtype |
|
) |
|
return latents, latent_image_ids |
|
|
|
if isinstance(generator, list) and len(generator) != batch_size: |
|
raise ValueError( |
|
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" |
|
f" size of {batch_size}. Make sure the batch size matches the length of the generators." |
|
) |
|
|
|
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) |
|
latents = self._pack_latents( |
|
latents, batch_size, num_channels_latents, height, width |
|
) |
|
|
|
latent_image_ids = self._prepare_latent_image_ids( |
|
batch_size, height, width, device, dtype |
|
) |
|
|
|
return latents, latent_image_ids |
|
|
|
@property |
|
def guidance_scale(self): |
|
return self._guidance_scale |
|
|
|
@property |
|
def joint_attention_kwargs(self): |
|
return self._joint_attention_kwargs |
|
|
|
@property |
|
def num_timesteps(self): |
|
return self._num_timesteps |
|
|
|
@property |
|
def interrupt(self): |
|
return self._interrupt |
|
|
|
@torch.no_grad() |
|
@replace_example_docstring(EXAMPLE_DOC_STRING) |
|
def __call__( |
|
self, |
|
prompt: Union[str, List[str]] = None, |
|
prompt_mask: Optional[Union[torch.FloatTensor, List[torch.FloatTensor]]] = None, |
|
negative_mask: Optional[ |
|
Union[torch.FloatTensor, List[torch.FloatTensor]] |
|
] = None, |
|
prompt_2: Optional[Union[str, List[str]]] = None, |
|
height: Optional[int] = None, |
|
width: Optional[int] = None, |
|
num_inference_steps: int = 28, |
|
timesteps: List[int] = None, |
|
guidance_scale: float = 3.5, |
|
num_images_per_prompt: Optional[int] = 1, |
|
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, |
|
latents: Optional[torch.FloatTensor] = None, |
|
prompt_embeds: Optional[torch.FloatTensor] = None, |
|
pooled_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
output_type: Optional[str] = "pil", |
|
return_dict: bool = True, |
|
joint_attention_kwargs: Optional[Dict[str, Any]] = None, |
|
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, |
|
callback_on_step_end_tensor_inputs: List[str] = ["latents"], |
|
max_sequence_length: int = 512, |
|
guidance_scale_real: float = 1.0, |
|
negative_prompt: Union[str, List[str]] = "", |
|
negative_prompt_2: Union[str, List[str]] = "", |
|
negative_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
no_cfg_until_timestep: int = 0, |
|
use_prompt_mask: bool = True, |
|
zero_using_prompt_mask: bool = False, |
|
device=torch.device('cuda'), |
|
): |
|
r""" |
|
Function invoked when calling the pipeline for generation. |
|
|
|
Args: |
|
prompt (`str` or `List[str]`, *optional*): |
|
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. |
|
instead. |
|
prompt_mask (`str` or `List[str]`, *optional*): |
|
The prompt or prompts to be used as a mask for the image generation. If not defined, `prompt` is used |
|
instead. |
|
prompt_2 (`str` or `List[str]`, *optional*): |
|
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is |
|
will be used instead |
|
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): |
|
The height in pixels of the generated image. This is set to 1024 by default for the best results. |
|
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): |
|
The width in pixels of the generated image. This is set to 1024 by default for the best results. |
|
num_inference_steps (`int`, *optional*, defaults to 50): |
|
The number of denoising steps. More denoising steps usually lead to a higher quality image at the |
|
expense of slower inference. |
|
timesteps (`List[int]`, *optional*): |
|
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument |
|
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is |
|
passed will be used. Must be in descending order. |
|
guidance_scale (`float`, *optional*, defaults to 7.0): |
|
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). |
|
`guidance_scale` is defined as `w` of equation 2. of [Imagen |
|
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > |
|
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, |
|
usually at the expense of lower image quality. |
|
num_images_per_prompt (`int`, *optional*, defaults to 1): |
|
The number of images to generate per prompt. |
|
generator (`torch.Generator` or `List[torch.Generator]`, *optional*): |
|
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) |
|
to make generation deterministic. |
|
latents (`torch.FloatTensor`, *optional*): |
|
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image |
|
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents |
|
tensor will ge generated by sampling using the supplied random `generator`. |
|
prompt_embeds (`torch.FloatTensor`, *optional*): |
|
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not |
|
provided, text embeddings will be generated from `prompt` input argument. |
|
pooled_prompt_embeds (`torch.FloatTensor`, *optional*): |
|
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. |
|
If not provided, pooled text embeddings will be generated from `prompt` input argument. |
|
output_type (`str`, *optional*, defaults to `"pil"`): |
|
The output format of the generate image. Choose between |
|
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. |
|
return_dict (`bool`, *optional*, defaults to `True`): |
|
Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple. |
|
joint_attention_kwargs (`dict`, *optional*): |
|
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under |
|
`self.processor` in |
|
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). |
|
callback_on_step_end (`Callable`, *optional*): |
|
A function that calls at the end of each denoising steps during the inference. The function is called |
|
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, |
|
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by |
|
`callback_on_step_end_tensor_inputs`. |
|
callback_on_step_end_tensor_inputs (`List`, *optional*): |
|
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list |
|
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the |
|
`._callback_tensor_inputs` attribute of your pipeline class. |
|
max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. |
|
|
|
Examples: |
|
|
|
Returns: |
|
[`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict` |
|
is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated |
|
images. |
|
""" |
|
|
|
height = height or self.default_sample_size * self.vae_scale_factor |
|
width = width or self.default_sample_size * self.vae_scale_factor |
|
|
|
|
|
self.check_inputs( |
|
prompt, |
|
prompt_2, |
|
height, |
|
width, |
|
prompt_embeds=prompt_embeds, |
|
pooled_prompt_embeds=pooled_prompt_embeds, |
|
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, |
|
max_sequence_length=max_sequence_length, |
|
) |
|
|
|
self._guidance_scale = guidance_scale |
|
self._guidance_scale_real = guidance_scale_real |
|
self._joint_attention_kwargs = joint_attention_kwargs |
|
self._interrupt = False |
|
|
|
|
|
if prompt is not None and isinstance(prompt, str): |
|
batch_size = 1 |
|
elif prompt is not None and isinstance(prompt, list): |
|
batch_size = len(prompt) |
|
else: |
|
batch_size = prompt_embeds.shape[0] |
|
|
|
device = device or self._execution_device |
|
|
|
lora_scale = ( |
|
self.joint_attention_kwargs.get("scale", None) |
|
if self.joint_attention_kwargs is not None |
|
else None |
|
) |
|
( |
|
prompt_embeds, |
|
pooled_prompt_embeds, |
|
text_ids, |
|
_prompt_mask, |
|
) = self.encode_prompt( |
|
prompt=prompt, |
|
prompt_2=prompt_2, |
|
prompt_embeds=prompt_embeds, |
|
pooled_prompt_embeds=pooled_prompt_embeds, |
|
device=device, |
|
num_images_per_prompt=num_images_per_prompt, |
|
max_sequence_length=max_sequence_length, |
|
lora_scale=lora_scale, |
|
) |
|
if _prompt_mask is not None: |
|
prompt_mask = _prompt_mask |
|
|
|
if negative_prompt_2 == "" and negative_prompt != "": |
|
negative_prompt_2 = negative_prompt |
|
|
|
negative_text_ids = text_ids |
|
if self._guidance_scale_real > 1.0 and ( |
|
negative_prompt_embeds is None or negative_pooled_prompt_embeds is None |
|
): |
|
( |
|
negative_prompt_embeds, |
|
negative_pooled_prompt_embeds, |
|
negative_text_ids, |
|
_neg_prompt_mask, |
|
) = self.encode_prompt( |
|
prompt=negative_prompt, |
|
prompt_2=negative_prompt_2, |
|
prompt_embeds=None, |
|
pooled_prompt_embeds=None, |
|
device=device, |
|
num_images_per_prompt=num_images_per_prompt, |
|
max_sequence_length=max_sequence_length, |
|
lora_scale=lora_scale, |
|
) |
|
|
|
if _neg_prompt_mask is not None: |
|
negative_mask = _neg_prompt_mask |
|
|
|
|
|
num_channels_latents = self.transformer.config.in_channels // 4 |
|
latents, latent_image_ids = self.prepare_latents( |
|
batch_size * num_images_per_prompt, |
|
num_channels_latents, |
|
height, |
|
width, |
|
prompt_embeds.dtype, |
|
device, |
|
generator, |
|
latents, |
|
) |
|
|
|
|
|
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) |
|
image_seq_len = latents.shape[1] |
|
mu = calculate_shift( |
|
image_seq_len, |
|
self.scheduler.config.base_image_seq_len, |
|
self.scheduler.config.max_image_seq_len, |
|
self.scheduler.config.base_shift, |
|
self.scheduler.config.max_shift, |
|
) |
|
timesteps, num_inference_steps = retrieve_timesteps( |
|
self.scheduler, |
|
num_inference_steps, |
|
device, |
|
timesteps, |
|
sigmas, |
|
mu=mu, |
|
) |
|
num_warmup_steps = max( |
|
len(timesteps) - num_inference_steps * self.scheduler.order, 0 |
|
) |
|
self._num_timesteps = len(timesteps) |
|
|
|
latents = latents |
|
latent_image_ids = latent_image_ids |
|
timesteps = timesteps |
|
text_ids = text_ids.to(device=device) |
|
|
|
|
|
if self.transformer.config.guidance_embeds: |
|
guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32) |
|
guidance = guidance.expand(latents.shape[0]) |
|
else: |
|
guidance = None |
|
|
|
if use_prompt_mask and prompt_mask is not None and not zero_using_prompt_mask: |
|
print('Using masking') |
|
elif use_prompt_mask and prompt_mask is not None and zero_using_prompt_mask: |
|
print('Using zeroed embeds') |
|
else: |
|
print('Not using masking') |
|
|
|
if self._guidance_scale_real > 1.0: |
|
print('Using classifier free guidance', self._guidance_scale_real) |
|
|
|
|
|
with self.progress_bar(total=num_inference_steps) as progress_bar: |
|
for i, t in enumerate(timesteps): |
|
if self.interrupt: |
|
continue |
|
|
|
|
|
timestep = t.expand(latents.shape[0]).to(latents.dtype) |
|
|
|
assert prompt_mask is not None |
|
|
|
extra_transformer_args = {} |
|
if use_prompt_mask and prompt_mask is not None and not zero_using_prompt_mask: |
|
extra_transformer_args["attention_mask"] = prompt_mask |
|
elif use_prompt_mask and prompt_mask is not None and zero_using_prompt_mask: |
|
mask_tens = prompt_mask.unsqueeze(-1).to(device=prompt_embeds.device, dtype=prompt_embeds.dtype) |
|
prompt_embeds = prompt_embeds * mask_tens |
|
|
|
noise_pred = self.transformer( |
|
hidden_states=latents, |
|
|
|
timestep=timestep / 1000, |
|
guidance=guidance, |
|
pooled_projections=pooled_prompt_embeds, |
|
encoder_hidden_states=prompt_embeds, |
|
txt_ids=text_ids, |
|
img_ids=latent_image_ids.to(device=device), |
|
joint_attention_kwargs=self.joint_attention_kwargs, |
|
return_dict=False, |
|
**extra_transformer_args, |
|
)[0] |
|
|
|
|
|
if self._guidance_scale_real > 1.0 and i >= no_cfg_until_timestep: |
|
extra_transformer_args_neg = {} |
|
if negative_mask is not None: |
|
extra_transformer_args_neg["attention_mask"] = negative_mask |
|
extra_transformer_args_neg["attention_mask"] is not None |
|
|
|
noise_pred_uncond = self.transformer( |
|
hidden_states=latents, |
|
|
|
timestep=timestep / 1000, |
|
guidance=guidance, |
|
pooled_projections=negative_pooled_prompt_embeds, |
|
encoder_hidden_states=negative_prompt_embeds, |
|
txt_ids=negative_text_ids, |
|
img_ids=latent_image_ids.to(device=device), |
|
joint_attention_kwargs=self.joint_attention_kwargs, |
|
return_dict=False, |
|
**extra_transformer_args_neg, |
|
)[0] |
|
|
|
noise_pred = noise_pred_uncond + self._guidance_scale_real * ( |
|
noise_pred - noise_pred_uncond |
|
) |
|
progress_bar.set_postfix( |
|
{ |
|
'ts': timestep.detach().item() / 1000, |
|
'cfg': self._guidance_scale_real, |
|
}, |
|
) |
|
else: |
|
progress_bar.set_postfix( |
|
{ |
|
'ts': timestep.detach().item() / 1000, |
|
'cfg': 'N/A', |
|
}, |
|
) |
|
|
|
|
|
latents_dtype = latents.dtype |
|
latents = self.scheduler.step( |
|
noise_pred, t, latents, return_dict=False |
|
)[0] |
|
|
|
if latents.dtype != latents_dtype: |
|
if torch.backends.mps.is_available(): |
|
|
|
latents = latents.to(latents_dtype) |
|
|
|
if callback_on_step_end is not None: |
|
callback_kwargs = {} |
|
for k in callback_on_step_end_tensor_inputs: |
|
callback_kwargs[k] = locals()[k] |
|
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) |
|
|
|
latents = callback_outputs.pop("latents", latents) |
|
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) |
|
|
|
|
|
if i == len(timesteps) - 1 or ( |
|
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 |
|
): |
|
progress_bar.update() |
|
|
|
if XLA_AVAILABLE: |
|
xm.mark_step() |
|
|
|
if output_type == "latent": |
|
image = latents |
|
|
|
else: |
|
latents = self._unpack_latents( |
|
latents, height, width, self.vae_scale_factor |
|
) |
|
latents = ( |
|
latents / self.vae.config.scaling_factor |
|
) + self.vae.config.shift_factor |
|
|
|
image = self.vae.decode( |
|
latents, |
|
return_dict=False, |
|
)[0] |
|
image = self.image_processor.postprocess(image, output_type=output_type) |
|
|
|
|
|
self.maybe_free_model_hooks() |
|
|
|
if not return_dict: |
|
return (image,) |
|
|
|
return FluxPipelineOutput(images=image) |
|
|