support miniflux
This commit is contained in:
kijai
@@ -0,0 +1,3 @@
|
||||
from .modeling_pyramid_flux import PyramidFluxTransformer
|
||||
from .modeling_text_encoder import FluxTextEncoderWithMask
|
||||
from .modeling_flux_block import FluxSingleTransformerBlock, FluxTransformerBlock
|
||||
@@ -0,0 +1,201 @@
|
||||
import math
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from diffusers.models.activations import get_activation, FP32SiLU
|
||||
|
||||
def get_timestep_embedding(
|
||||
timesteps: torch.Tensor,
|
||||
embedding_dim: int,
|
||||
flip_sin_to_cos: bool = False,
|
||||
downscale_freq_shift: float = 1,
|
||||
scale: float = 1,
|
||||
max_period: int = 10000,
|
||||
):
|
||||
"""
|
||||
This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
|
||||
|
||||
Args
|
||||
timesteps (torch.Tensor):
|
||||
a 1-D Tensor of N indices, one per batch element. These may be fractional.
|
||||
embedding_dim (int):
|
||||
the dimension of the output.
|
||||
flip_sin_to_cos (bool):
|
||||
Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
|
||||
downscale_freq_shift (float):
|
||||
Controls the delta between frequencies between dimensions
|
||||
scale (float):
|
||||
Scaling factor applied to the embeddings.
|
||||
max_period (int):
|
||||
Controls the maximum frequency of the embeddings
|
||||
Returns
|
||||
torch.Tensor: an [N x dim] Tensor of positional embeddings.
|
||||
"""
|
||||
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
|
||||
|
||||
half_dim = embedding_dim // 2
|
||||
exponent = -math.log(max_period) * torch.arange(
|
||||
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
|
||||
)
|
||||
exponent = exponent / (half_dim - downscale_freq_shift)
|
||||
|
||||
emb = torch.exp(exponent)
|
||||
emb = timesteps[:, None].float() * emb[None, :]
|
||||
|
||||
# scale embeddings
|
||||
emb = scale * emb
|
||||
|
||||
# concat sine and cosine embeddings
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
|
||||
|
||||
# flip sine and cosine embeddings
|
||||
if flip_sin_to_cos:
|
||||
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
|
||||
|
||||
# zero pad
|
||||
if embedding_dim % 2 == 1:
|
||||
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
|
||||
return emb
|
||||
|
||||
|
||||
class Timesteps(nn.Module):
|
||||
def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
|
||||
super().__init__()
|
||||
self.num_channels = num_channels
|
||||
self.flip_sin_to_cos = flip_sin_to_cos
|
||||
self.downscale_freq_shift = downscale_freq_shift
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, timesteps):
|
||||
t_emb = get_timestep_embedding(
|
||||
timesteps,
|
||||
self.num_channels,
|
||||
flip_sin_to_cos=self.flip_sin_to_cos,
|
||||
downscale_freq_shift=self.downscale_freq_shift,
|
||||
scale=self.scale,
|
||||
)
|
||||
return t_emb
|
||||
|
||||
|
||||
class TimestepEmbedding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
time_embed_dim: int,
|
||||
act_fn: str = "silu",
|
||||
out_dim: int = None,
|
||||
post_act_fn: Optional[str] = None,
|
||||
cond_proj_dim=None,
|
||||
sample_proj_bias=True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
|
||||
|
||||
if cond_proj_dim is not None:
|
||||
self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
|
||||
else:
|
||||
self.cond_proj = None
|
||||
|
||||
self.act = get_activation(act_fn)
|
||||
|
||||
if out_dim is not None:
|
||||
time_embed_dim_out = out_dim
|
||||
else:
|
||||
time_embed_dim_out = time_embed_dim
|
||||
self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
|
||||
|
||||
if post_act_fn is None:
|
||||
self.post_act = None
|
||||
else:
|
||||
self.post_act = get_activation(post_act_fn)
|
||||
|
||||
def forward(self, sample, condition=None):
|
||||
if condition is not None:
|
||||
sample = sample + self.cond_proj(condition)
|
||||
sample = self.linear_1(sample)
|
||||
|
||||
if self.act is not None:
|
||||
sample = self.act(sample)
|
||||
|
||||
sample = self.linear_2(sample)
|
||||
|
||||
if self.post_act is not None:
|
||||
sample = self.post_act(sample)
|
||||
return sample
|
||||
|
||||
|
||||
class PixArtAlphaTextProjection(nn.Module):
|
||||
"""
|
||||
Projects caption embeddings. Also handles dropout for classifier-free guidance.
|
||||
|
||||
Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, hidden_size, out_features=None, act_fn="gelu_tanh"):
|
||||
super().__init__()
|
||||
if out_features is None:
|
||||
out_features = hidden_size
|
||||
self.linear_1 = nn.Linear(in_features=in_features, out_features=hidden_size, bias=True)
|
||||
if act_fn == "gelu_tanh":
|
||||
self.act_1 = nn.GELU(approximate="tanh")
|
||||
elif act_fn == "silu":
|
||||
self.act_1 = nn.SiLU()
|
||||
elif act_fn == "silu_fp32":
|
||||
self.act_1 = FP32SiLU()
|
||||
else:
|
||||
raise ValueError(f"Unknown activation function: {act_fn}")
|
||||
self.linear_2 = nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)
|
||||
|
||||
def forward(self, caption):
|
||||
hidden_states = self.linear_1(caption)
|
||||
hidden_states = self.act_1(hidden_states)
|
||||
hidden_states = self.linear_2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
|
||||
def __init__(self, embedding_dim, pooled_projection_dim):
|
||||
super().__init__()
|
||||
|
||||
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
|
||||
self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
|
||||
self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
|
||||
|
||||
def forward(self, timestep, guidance, pooled_projection):
|
||||
timesteps_proj = self.time_proj(timestep)
|
||||
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) # (N, D)
|
||||
|
||||
guidance_proj = self.time_proj(guidance)
|
||||
guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype)) # (N, D)
|
||||
|
||||
time_guidance_emb = timesteps_emb + guidance_emb
|
||||
|
||||
pooled_projections = self.text_embedder(pooled_projection)
|
||||
conditioning = time_guidance_emb + pooled_projections
|
||||
|
||||
return conditioning
|
||||
|
||||
|
||||
class CombinedTimestepTextProjEmbeddings(nn.Module):
|
||||
def __init__(self, embedding_dim, pooled_projection_dim):
|
||||
super().__init__()
|
||||
|
||||
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
|
||||
self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
|
||||
|
||||
def forward(self, timestep, pooled_projection):
|
||||
timesteps_proj = self.time_proj(timestep)
|
||||
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) # (N, D)
|
||||
|
||||
pooled_projections = self.text_embedder(pooled_projection)
|
||||
|
||||
conditioning = timesteps_emb + pooled_projections
|
||||
|
||||
return conditioning
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,249 @@
|
||||
import numbers
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
from diffusers.utils import is_torch_version
|
||||
|
||||
|
||||
if is_torch_version(">=", "2.1.0"):
|
||||
LayerNorm = nn.LayerNorm
|
||||
else:
|
||||
# Has optional bias parameter compared to torch layer norm
|
||||
# TODO: replace with torch layernorm once min required torch version >= 2.1
|
||||
class LayerNorm(nn.Module):
|
||||
def __init__(self, dim, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True):
|
||||
super().__init__()
|
||||
|
||||
self.eps = eps
|
||||
|
||||
if isinstance(dim, numbers.Integral):
|
||||
dim = (dim,)
|
||||
|
||||
self.dim = torch.Size(dim)
|
||||
|
||||
if elementwise_affine:
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.bias = nn.Parameter(torch.zeros(dim)) if bias else None
|
||||
else:
|
||||
self.weight = None
|
||||
self.bias = None
|
||||
|
||||
def forward(self, input):
|
||||
return F.layer_norm(input, self.dim, self.weight, self.bias, self.eps)
|
||||
|
||||
|
||||
class FP32LayerNorm(nn.LayerNorm):
|
||||
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
||||
origin_dtype = inputs.dtype
|
||||
return F.layer_norm(
|
||||
inputs.float(),
|
||||
self.normalized_shape,
|
||||
self.weight.float() if self.weight is not None else None,
|
||||
self.bias.float() if self.bias is not None else None,
|
||||
self.eps,
|
||||
).to(origin_dtype)
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim, eps: float, elementwise_affine: bool = True):
|
||||
super().__init__()
|
||||
|
||||
self.eps = eps
|
||||
|
||||
if isinstance(dim, numbers.Integral):
|
||||
dim = (dim,)
|
||||
|
||||
self.dim = torch.Size(dim)
|
||||
|
||||
if elementwise_affine:
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
else:
|
||||
self.weight = None
|
||||
|
||||
def forward(self, hidden_states):
|
||||
input_dtype = hidden_states.dtype
|
||||
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
|
||||
|
||||
if self.weight is not None:
|
||||
# convert into half-precision if necessary
|
||||
if self.weight.dtype in [torch.float16, torch.bfloat16]:
|
||||
hidden_states = hidden_states.to(self.weight.dtype)
|
||||
hidden_states = hidden_states * self.weight
|
||||
else:
|
||||
hidden_states = hidden_states.to(input_dtype)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class AdaLayerNormContinuous(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
conditioning_embedding_dim: int,
|
||||
# NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
|
||||
# because the output is immediately scaled and shifted by the projected conditioning embeddings.
|
||||
# Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
|
||||
# However, this is how it was implemented in the original code, and it's rather likely you should
|
||||
# set `elementwise_affine` to False.
|
||||
elementwise_affine=True,
|
||||
eps=1e-5,
|
||||
bias=True,
|
||||
norm_type="layer_norm",
|
||||
):
|
||||
super().__init__()
|
||||
self.silu = nn.SiLU()
|
||||
self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
|
||||
if norm_type == "layer_norm":
|
||||
self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias)
|
||||
elif norm_type == "rms_norm":
|
||||
self.norm = RMSNorm(embedding_dim, eps, elementwise_affine)
|
||||
else:
|
||||
raise ValueError(f"unknown norm_type {norm_type}")
|
||||
|
||||
def forward_with_pad(self, x: torch.Tensor, conditioning_embedding: torch.Tensor, hidden_length=None) -> torch.Tensor:
|
||||
assert hidden_length is not None
|
||||
|
||||
emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
|
||||
batch_emb = torch.zeros_like(x).repeat(1, 1, 2)
|
||||
|
||||
i_sum = 0
|
||||
num_stages = len(hidden_length)
|
||||
for i_p, length in enumerate(hidden_length):
|
||||
batch_emb[:, i_sum:i_sum+length] = emb[i_p::num_stages][:,None]
|
||||
i_sum += length
|
||||
|
||||
batch_scale, batch_shift = torch.chunk(batch_emb, 2, dim=2)
|
||||
x = self.norm(x) * (1 + batch_scale) + batch_shift
|
||||
return x
|
||||
|
||||
def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor, hidden_length=None) -> torch.Tensor:
|
||||
# convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
|
||||
if hidden_length is not None:
|
||||
return self.forward_with_pad(x, conditioning_embedding, hidden_length)
|
||||
emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
|
||||
scale, shift = torch.chunk(emb, 2, dim=1)
|
||||
x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
|
||||
return x
|
||||
|
||||
|
||||
class AdaLayerNormZero(nn.Module):
|
||||
r"""
|
||||
Norm layer adaptive layer norm zero (adaLN-Zero).
|
||||
|
||||
Parameters:
|
||||
embedding_dim (`int`): The size of each embedding vector.
|
||||
num_embeddings (`int`): The size of the embeddings dictionary.
|
||||
"""
|
||||
|
||||
def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None):
|
||||
super().__init__()
|
||||
self.emb = None
|
||||
|
||||
self.silu = nn.SiLU()
|
||||
self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)
|
||||
self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
|
||||
|
||||
def forward_with_pad(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
hidden_dtype: Optional[torch.dtype] = None,
|
||||
emb: Optional[torch.Tensor] = None,
|
||||
hidden_length: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
# hidden_length: [[20, 30], [30, 40], [50, 60]]
|
||||
# x: [bs, seq_len, dim]
|
||||
if self.emb is not None:
|
||||
emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
|
||||
|
||||
emb = self.linear(self.silu(emb))
|
||||
batch_emb = torch.zeros_like(x).repeat(1, 1, 6)
|
||||
|
||||
i_sum = 0
|
||||
num_stages = len(hidden_length)
|
||||
for i_p, length in enumerate(hidden_length):
|
||||
batch_emb[:, i_sum:i_sum+length] = emb[i_p::num_stages][:,None]
|
||||
i_sum += length
|
||||
|
||||
batch_shift_msa, batch_scale_msa, batch_gate_msa, batch_shift_mlp, batch_scale_mlp, batch_gate_mlp = batch_emb.chunk(6, dim=2)
|
||||
x = self.norm(x) * (1 + batch_scale_msa) + batch_shift_msa
|
||||
return x, batch_gate_msa, batch_shift_mlp, batch_scale_mlp, batch_gate_mlp
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
hidden_dtype: Optional[torch.dtype] = None,
|
||||
emb: Optional[torch.Tensor] = None,
|
||||
hidden_length: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if hidden_length is not None:
|
||||
return self.forward_with_pad(x, timestep, class_labels, hidden_dtype, emb, hidden_length)
|
||||
if self.emb is not None:
|
||||
emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
|
||||
emb = self.linear(self.silu(emb))
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)
|
||||
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
|
||||
return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
|
||||
|
||||
|
||||
class AdaLayerNormZeroSingle(nn.Module):
|
||||
r"""
|
||||
Norm layer adaptive layer norm zero (adaLN-Zero).
|
||||
|
||||
Parameters:
|
||||
embedding_dim (`int`): The size of each embedding vector.
|
||||
num_embeddings (`int`): The size of the embeddings dictionary.
|
||||
"""
|
||||
|
||||
def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
|
||||
super().__init__()
|
||||
|
||||
self.silu = nn.SiLU()
|
||||
self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias)
|
||||
if norm_type == "layer_norm":
|
||||
self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'."
|
||||
)
|
||||
|
||||
def forward_with_pad(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
emb: Optional[torch.Tensor] = None,
|
||||
hidden_length: Optional[torch.Tensor] = None,
|
||||
):
|
||||
emb = self.linear(self.silu(emb))
|
||||
batch_emb = torch.zeros_like(x).repeat(1, 1, 3)
|
||||
|
||||
i_sum = 0
|
||||
num_stages = len(hidden_length)
|
||||
for i_p, length in enumerate(hidden_length):
|
||||
batch_emb[:, i_sum:i_sum+length] = emb[i_p::num_stages][:,None]
|
||||
i_sum += length
|
||||
|
||||
batch_shift_msa, batch_scale_msa, batch_gate_msa = batch_emb.chunk(3, dim=2)
|
||||
|
||||
x = self.norm(x) * (1 + batch_scale_msa) + batch_shift_msa
|
||||
|
||||
return x, batch_gate_msa
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
emb: Optional[torch.Tensor] = None,
|
||||
hidden_length: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if hidden_length is not None:
|
||||
return self.forward_with_pad(x, emb, hidden_length)
|
||||
emb = self.linear(self.silu(emb))
|
||||
shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
|
||||
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
|
||||
return x, gate_msa
|
||||
@@ -0,0 +1,538 @@
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import torch
|
||||
import os
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
from tqdm import tqdm
|
||||
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.utils import is_torch_version
|
||||
|
||||
from .modeling_normalization import AdaLayerNormContinuous
|
||||
from .modeling_embedding import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings
|
||||
from .modeling_flux_block import FluxTransformerBlock, FluxSingleTransformerBlock
|
||||
|
||||
from ...trainer_misc import (
|
||||
is_sequence_parallel_initialized,
|
||||
get_sequence_parallel_group,
|
||||
get_sequence_parallel_world_size,
|
||||
get_sequence_parallel_rank,
|
||||
all_to_all,
|
||||
)
|
||||
|
||||
|
||||
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(2)
|
||||
|
||||
|
||||
class PyramidFluxTransformer(ModelMixin, ConfigMixin):
|
||||
"""
|
||||
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`.
|
||||
"""
|
||||
|
||||
_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 = 64,
|
||||
num_attention_heads: int = 24,
|
||||
joint_attention_dim: int = 4096,
|
||||
pooled_projection_dim: int = 768,
|
||||
axes_dims_rope: List[int] = [16, 24, 24],
|
||||
use_flash_attn: bool = False,
|
||||
use_temporal_causal: bool = True,
|
||||
interp_condition_pos: bool = True,
|
||||
use_gradient_checkpointing: bool = False,
|
||||
gradient_checkpointing_ratio: float = 0.6,
|
||||
):
|
||||
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)
|
||||
self.time_text_embed = CombinedTimestepTextProjEmbeddings(
|
||||
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,
|
||||
use_flash_attn=use_flash_attn,
|
||||
)
|
||||
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,
|
||||
use_flash_attn=use_flash_attn,
|
||||
)
|
||||
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 = use_gradient_checkpointing
|
||||
self.gradient_checkpointing_ratio = gradient_checkpointing_ratio
|
||||
|
||||
self.use_temporal_causal = use_temporal_causal
|
||||
if self.use_temporal_causal:
|
||||
print("Using temporal causal attention")
|
||||
|
||||
self.use_flash_attn = use_flash_attn
|
||||
if self.use_flash_attn:
|
||||
print("Using Flash attention")
|
||||
|
||||
self.patch_size = 2 # hard-code for now
|
||||
|
||||
# init weights
|
||||
self.initialize_weights()
|
||||
|
||||
def initialize_weights(self):
|
||||
# Initialize transformer layers:
|
||||
def _basic_init(module):
|
||||
if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv3d)):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
self.apply(_basic_init)
|
||||
|
||||
# Initialize all the conditioning to normal init
|
||||
nn.init.normal_(self.time_text_embed.timestep_embedder.linear_1.weight, std=0.02)
|
||||
nn.init.normal_(self.time_text_embed.timestep_embedder.linear_2.weight, std=0.02)
|
||||
nn.init.normal_(self.time_text_embed.text_embedder.linear_1.weight, std=0.02)
|
||||
nn.init.normal_(self.time_text_embed.text_embedder.linear_2.weight, std=0.02)
|
||||
nn.init.normal_(self.context_embedder.weight, std=0.02)
|
||||
|
||||
# Zero-out adaLN modulation layers in DiT blocks:
|
||||
for block in self.transformer_blocks:
|
||||
nn.init.constant_(block.norm1.linear.weight, 0)
|
||||
nn.init.constant_(block.norm1.linear.bias, 0)
|
||||
nn.init.constant_(block.norm1_context.linear.weight, 0)
|
||||
nn.init.constant_(block.norm1_context.linear.bias, 0)
|
||||
|
||||
for block in self.single_transformer_blocks:
|
||||
nn.init.constant_(block.norm.linear.weight, 0)
|
||||
nn.init.constant_(block.norm.linear.bias, 0)
|
||||
|
||||
# Zero-out output layers:
|
||||
nn.init.constant_(self.norm_out.linear.weight, 0)
|
||||
nn.init.constant_(self.norm_out.linear.bias, 0)
|
||||
nn.init.constant_(self.proj_out.weight, 0)
|
||||
nn.init.constant_(self.proj_out.bias, 0)
|
||||
|
||||
@torch.no_grad()
|
||||
def _prepare_image_ids(self, batch_size, temp, height, width, train_height, train_width, device, start_time_stamp=0):
|
||||
latent_image_ids = torch.zeros(temp, height, width, 3)
|
||||
|
||||
# Temporal Rope
|
||||
latent_image_ids[..., 0] = latent_image_ids[..., 0] + torch.arange(start_time_stamp, start_time_stamp + temp)[:, None, None]
|
||||
|
||||
# height Rope
|
||||
if height != train_height:
|
||||
height_pos = F.interpolate(torch.arange(train_height)[None, None, :].float(), height, mode='linear').squeeze(0, 1)
|
||||
else:
|
||||
height_pos = torch.arange(train_height).float()
|
||||
|
||||
latent_image_ids[..., 1] = latent_image_ids[..., 1] + height_pos[None, :, None]
|
||||
|
||||
# width rope
|
||||
if width != train_width:
|
||||
width_pos = F.interpolate(torch.arange(train_width)[None, None, :].float(), width, mode='linear').squeeze(0, 1)
|
||||
else:
|
||||
width_pos = torch.arange(train_width).float()
|
||||
|
||||
latent_image_ids[..., 2] = latent_image_ids[..., 2] + width_pos[None, None, :]
|
||||
|
||||
latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1, 1)
|
||||
latent_image_ids = rearrange(latent_image_ids, 'b t h w c -> b (t h w) c')
|
||||
|
||||
return latent_image_ids.to(device=device)
|
||||
|
||||
@torch.no_grad()
|
||||
def _prepare_pyramid_image_ids(self, sample, batch_size, device):
|
||||
image_ids_list = []
|
||||
|
||||
for i_b, sample_ in enumerate(sample):
|
||||
if not isinstance(sample_, list):
|
||||
sample_ = [sample_]
|
||||
|
||||
cur_image_ids = []
|
||||
start_time_stamp = 0
|
||||
|
||||
train_height = sample_[-1].shape[-2] // self.patch_size
|
||||
train_width = sample_[-1].shape[-1] // self.patch_size
|
||||
|
||||
for clip_ in sample_:
|
||||
_, _, temp, height, width = clip_.shape
|
||||
height = height // self.patch_size
|
||||
width = width // self.patch_size
|
||||
cur_image_ids.append(self._prepare_image_ids(batch_size, temp, height, width, train_height, train_width, device, start_time_stamp=start_time_stamp))
|
||||
start_time_stamp += temp
|
||||
|
||||
cur_image_ids = torch.cat(cur_image_ids, dim=1)
|
||||
image_ids_list.append(cur_image_ids)
|
||||
|
||||
return image_ids_list
|
||||
|
||||
def merge_input(self, sample, encoder_hidden_length, encoder_attention_mask):
|
||||
"""
|
||||
Merge the input video with different resolutions into one sequence
|
||||
Sample: From low resolution to high resolution
|
||||
"""
|
||||
if isinstance(sample[0], list):
|
||||
device = sample[0][-1].device
|
||||
pad_batch_size = sample[0][-1].shape[0]
|
||||
else:
|
||||
device = sample[0].device
|
||||
pad_batch_size = sample[0].shape[0]
|
||||
|
||||
num_stages = len(sample)
|
||||
height_list = [];width_list = [];temp_list = []
|
||||
trainable_token_list = []
|
||||
|
||||
for i_b, sample_ in enumerate(sample):
|
||||
if isinstance(sample_, list):
|
||||
sample_ = sample_[-1]
|
||||
_, _, temp, height, width = sample_.shape
|
||||
height = height // self.patch_size
|
||||
width = width // self.patch_size
|
||||
temp_list.append(temp)
|
||||
height_list.append(height)
|
||||
width_list.append(width)
|
||||
trainable_token_list.append(height * width * temp)
|
||||
|
||||
# prepare the RoPE IDs,
|
||||
image_ids_list = self._prepare_pyramid_image_ids(sample, pad_batch_size, device)
|
||||
text_ids = torch.zeros(pad_batch_size, encoder_attention_mask.shape[1], 3).to(device=device)
|
||||
input_ids_list = [torch.cat([text_ids, image_ids], dim=1) for image_ids in image_ids_list]
|
||||
image_rotary_emb = [self.pos_embed(input_ids) for input_ids in input_ids_list] # [bs, seq_len, 1, head_dim // 2, 2, 2]
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
concat_output = True if self.training else False
|
||||
image_rotary_emb = [all_to_all(x_.repeat(1, 1, sp_group_size, 1, 1, 1), sp_group, sp_group_size, scatter_dim=2, gather_dim=0, concat_output=concat_output) for x_ in image_rotary_emb]
|
||||
input_ids_list = [all_to_all(input_ids.repeat(1, 1, sp_group_size), sp_group, sp_group_size, scatter_dim=2, gather_dim=0, concat_output=concat_output) for input_ids in input_ids_list]
|
||||
|
||||
hidden_states, hidden_length = [], []
|
||||
|
||||
for sample_ in sample:
|
||||
video_tokens = []
|
||||
|
||||
for each_latent in sample_:
|
||||
each_latent = rearrange(each_latent, 'b c t h w -> b t h w c')
|
||||
each_latent = rearrange(each_latent, 'b t (h p1) (w p2) c -> b (t h w) (p1 p2 c)', p1=self.patch_size, p2=self.patch_size)
|
||||
video_tokens.append(each_latent)
|
||||
|
||||
video_tokens = torch.cat(video_tokens, dim=1)
|
||||
video_tokens = self.x_embedder(video_tokens)
|
||||
hidden_states.append(video_tokens)
|
||||
hidden_length.append(video_tokens.shape[1])
|
||||
|
||||
# prepare the attention mask
|
||||
if self.use_flash_attn:
|
||||
attention_mask = None
|
||||
indices_list = []
|
||||
for i_p, length in enumerate(hidden_length):
|
||||
pad_attention_mask = torch.ones((pad_batch_size, length), dtype=encoder_attention_mask.dtype).to(device)
|
||||
pad_attention_mask = torch.cat([encoder_attention_mask[i_p::num_stages], pad_attention_mask], dim=1)
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
pad_attention_mask = all_to_all(pad_attention_mask.unsqueeze(2).repeat(1, 1, sp_group_size), sp_group, sp_group_size, scatter_dim=2, gather_dim=0)
|
||||
pad_attention_mask = pad_attention_mask.squeeze(2)
|
||||
|
||||
seqlens_in_batch = pad_attention_mask.sum(dim=-1, dtype=torch.int32)
|
||||
indices = torch.nonzero(pad_attention_mask.flatten(), as_tuple=False).flatten()
|
||||
|
||||
indices_list.append(
|
||||
{
|
||||
'indices': indices,
|
||||
'seqlens_in_batch': seqlens_in_batch,
|
||||
}
|
||||
)
|
||||
encoder_attention_mask = indices_list
|
||||
else:
|
||||
assert encoder_attention_mask.shape[1] == encoder_hidden_length
|
||||
real_batch_size = encoder_attention_mask.shape[0]
|
||||
|
||||
# prepare text ids
|
||||
text_ids = torch.arange(1, real_batch_size + 1, dtype=encoder_attention_mask.dtype).unsqueeze(1).repeat(1, encoder_hidden_length)
|
||||
text_ids = text_ids.to(device)
|
||||
text_ids[encoder_attention_mask == 0] = 0
|
||||
|
||||
# prepare image ids
|
||||
image_ids = torch.arange(1, real_batch_size + 1, dtype=encoder_attention_mask.dtype).unsqueeze(1).repeat(1, max(hidden_length))
|
||||
image_ids = image_ids.to(device)
|
||||
image_ids_list = []
|
||||
for i_p, length in enumerate(hidden_length):
|
||||
image_ids_list.append(image_ids[i_p::num_stages][:, :length])
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
concat_output = True if self.training else False
|
||||
text_ids = all_to_all(text_ids.unsqueeze(2).repeat(1, 1, sp_group_size), sp_group, sp_group_size, scatter_dim=2, gather_dim=0, concat_output=concat_output).squeeze(2)
|
||||
image_ids_list = [all_to_all(image_ids_.unsqueeze(2).repeat(1, 1, sp_group_size), sp_group, sp_group_size, scatter_dim=2, gather_dim=0, concat_output=concat_output).squeeze(2) for image_ids_ in image_ids_list]
|
||||
|
||||
attention_mask = []
|
||||
for i_p in range(len(hidden_length)):
|
||||
image_ids = image_ids_list[i_p]
|
||||
token_ids = torch.cat([text_ids[i_p::num_stages], image_ids], dim=1)
|
||||
stage_attention_mask = rearrange(token_ids, 'b i -> b 1 i 1') == rearrange(token_ids, 'b j -> b 1 1 j') # [bs, 1, q_len, k_len]
|
||||
if self.use_temporal_causal:
|
||||
input_order_ids = input_ids_list[i_p][:,:,0]
|
||||
temporal_causal_mask = rearrange(input_order_ids, 'b i -> b 1 i 1') >= rearrange(input_order_ids, 'b j -> b 1 1 j')
|
||||
stage_attention_mask = stage_attention_mask & temporal_causal_mask
|
||||
attention_mask.append(stage_attention_mask)
|
||||
|
||||
return hidden_states, hidden_length, temp_list, height_list, width_list, trainable_token_list, encoder_attention_mask, attention_mask, image_rotary_emb
|
||||
|
||||
def split_output(self, batch_hidden_states, hidden_length, temps, heights, widths, trainable_token_list):
|
||||
# To split the hidden states
|
||||
batch_size = batch_hidden_states.shape[0]
|
||||
output_hidden_list = []
|
||||
batch_hidden_states = torch.split(batch_hidden_states, hidden_length, dim=1)
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
batch_size = batch_size // sp_group_size
|
||||
|
||||
for i_p, length in enumerate(hidden_length):
|
||||
width, height, temp = widths[i_p], heights[i_p], temps[i_p]
|
||||
trainable_token_num = trainable_token_list[i_p]
|
||||
hidden_states = batch_hidden_states[i_p]
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
hidden_states = all_to_all(hidden_states, sp_group, sp_group_size, scatter_dim=0, gather_dim=1)
|
||||
|
||||
# only the trainable token are taking part in loss computation
|
||||
hidden_states = hidden_states[:, -trainable_token_num:]
|
||||
|
||||
# unpatchify
|
||||
hidden_states = hidden_states.reshape(
|
||||
shape=(batch_size, temp, height, width, self.patch_size, self.patch_size, self.out_channels // 4)
|
||||
)
|
||||
hidden_states = rearrange(hidden_states, "b t h w p1 p2 c -> b t (h p1) (w p2) c")
|
||||
hidden_states = rearrange(hidden_states, "b t h w c -> b c t h w")
|
||||
output_hidden_list.append(hidden_states)
|
||||
|
||||
return output_hidden_list
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.FloatTensor, # [num_stages]
|
||||
encoder_hidden_states: torch.Tensor = None,
|
||||
encoder_attention_mask: torch.FloatTensor = None,
|
||||
pooled_projections: torch.Tensor = None,
|
||||
timestep_ratio: torch.LongTensor = None,
|
||||
):
|
||||
temb = self.time_text_embed(timestep_ratio, pooled_projections)
|
||||
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
|
||||
encoder_hidden_length = encoder_hidden_states.shape[1]
|
||||
|
||||
# Get the input sequence
|
||||
hidden_states, hidden_length, temps, heights, widths, trainable_token_list, encoder_attention_mask, attention_mask, \
|
||||
image_rotary_emb = self.merge_input(sample, encoder_hidden_length, encoder_attention_mask)
|
||||
|
||||
# split the long latents if necessary
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
concat_output = True if self.training else False
|
||||
|
||||
# sync the input hidden states
|
||||
batch_hidden_states = []
|
||||
for i_p, hidden_states_ in enumerate(hidden_states):
|
||||
assert hidden_states_.shape[1] % sp_group_size == 0, "The sequence length should be divided by sequence parallel size"
|
||||
hidden_states_ = all_to_all(hidden_states_, sp_group, sp_group_size, scatter_dim=1, gather_dim=0, concat_output=concat_output)
|
||||
hidden_length[i_p] = hidden_length[i_p] // sp_group_size
|
||||
batch_hidden_states.append(hidden_states_)
|
||||
|
||||
# sync the encoder hidden states
|
||||
hidden_states = torch.cat(batch_hidden_states, dim=1)
|
||||
encoder_hidden_states = all_to_all(encoder_hidden_states, sp_group, sp_group_size, scatter_dim=1, gather_dim=0, concat_output=concat_output)
|
||||
temb = all_to_all(temb.unsqueeze(1).repeat(1, sp_group_size, 1), sp_group, sp_group_size, scatter_dim=1, gather_dim=0, concat_output=concat_output)
|
||||
temb = temb.squeeze(1)
|
||||
else:
|
||||
hidden_states = torch.cat(hidden_states, dim=1)
|
||||
|
||||
for index_block, block in enumerate(self.transformer_blocks):
|
||||
if self.training and self.gradient_checkpointing and (index_block <= int(len(self.transformer_blocks) * self.gradient_checkpointing_ratio)):
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs):
|
||||
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,
|
||||
encoder_attention_mask,
|
||||
temb,
|
||||
attention_mask,
|
||||
hidden_length,
|
||||
image_rotary_emb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
encoder_hidden_states, hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
temb=temb,
|
||||
attention_mask=attention_mask,
|
||||
hidden_length=hidden_length,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
)
|
||||
|
||||
# remerge for single attention block
|
||||
num_stages = len(hidden_length)
|
||||
batch_hidden_states = list(torch.split(hidden_states, hidden_length, dim=1))
|
||||
concat_hidden_length = []
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
encoder_hidden_states = all_to_all(encoder_hidden_states, sp_group, sp_group_size, scatter_dim=0, gather_dim=1)
|
||||
|
||||
for i_p in range(len(hidden_length)):
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
batch_hidden_states[i_p] = all_to_all(batch_hidden_states[i_p], sp_group, sp_group_size, scatter_dim=0, gather_dim=1)
|
||||
|
||||
batch_hidden_states[i_p] = torch.cat([encoder_hidden_states[i_p::num_stages], batch_hidden_states[i_p]], dim=1)
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
batch_hidden_states[i_p] = all_to_all(batch_hidden_states[i_p], sp_group, sp_group_size, scatter_dim=1, gather_dim=0)
|
||||
|
||||
concat_hidden_length.append(batch_hidden_states[i_p].shape[1])
|
||||
|
||||
hidden_states = torch.cat(batch_hidden_states, dim=1)
|
||||
|
||||
for index_block, block in enumerate(self.single_transformer_blocks):
|
||||
if self.training and self.gradient_checkpointing and (index_block <= int(len(self.single_transformer_blocks) * self.gradient_checkpointing_ratio)):
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs):
|
||||
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,
|
||||
encoder_attention_mask,
|
||||
attention_mask,
|
||||
concat_hidden_length,
|
||||
image_rotary_emb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
temb=temb,
|
||||
encoder_attention_mask=encoder_attention_mask, # used for
|
||||
attention_mask=attention_mask,
|
||||
hidden_length=concat_hidden_length,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
)
|
||||
|
||||
batch_hidden_states = list(torch.split(hidden_states, concat_hidden_length, dim=1))
|
||||
|
||||
for i_p in range(len(concat_hidden_length)):
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
batch_hidden_states[i_p] = all_to_all(batch_hidden_states[i_p], sp_group, sp_group_size, scatter_dim=0, gather_dim=1)
|
||||
|
||||
batch_hidden_states[i_p] = batch_hidden_states[i_p][:, encoder_hidden_length :, ...]
|
||||
|
||||
if is_sequence_parallel_initialized():
|
||||
sp_group = get_sequence_parallel_group()
|
||||
sp_group_size = get_sequence_parallel_world_size()
|
||||
batch_hidden_states[i_p] = all_to_all(batch_hidden_states[i_p], sp_group, sp_group_size, scatter_dim=1, gather_dim=0)
|
||||
|
||||
hidden_states = torch.cat(batch_hidden_states, dim=1)
|
||||
hidden_states = self.norm_out(hidden_states, temb, hidden_length=hidden_length)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
|
||||
output = self.split_output(hidden_states, hidden_length, temps, heights, widths, trainable_token_list)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,134 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import os
|
||||
|
||||
from transformers import (
|
||||
CLIPTextModel,
|
||||
CLIPTokenizer,
|
||||
T5EncoderModel,
|
||||
T5TokenizerFast,
|
||||
)
|
||||
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
|
||||
class FluxTextEncoderWithMask(nn.Module):
|
||||
def __init__(self, model_path, torch_dtype):
|
||||
super().__init__()
|
||||
# CLIP-G
|
||||
self.tokenizer = CLIPTokenizer.from_pretrained(os.path.join(model_path, 'tokenizer'), torch_dtype=torch_dtype)
|
||||
self.tokenizer_max_length = (
|
||||
self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
|
||||
)
|
||||
self.text_encoder = CLIPTextModel.from_pretrained(os.path.join(model_path, 'text_encoder'), torch_dtype=torch_dtype)
|
||||
|
||||
# T5
|
||||
self.tokenizer_2 = T5TokenizerFast.from_pretrained(os.path.join(model_path, 'tokenizer_2'))
|
||||
self.text_encoder_2 = T5EncoderModel.from_pretrained(os.path.join(model_path, 'text_encoder_2'), torch_dtype=torch_dtype)
|
||||
|
||||
self._freeze()
|
||||
|
||||
def _freeze(self):
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def _get_t5_prompt_embeds(
|
||||
self,
|
||||
prompt: Union[str, List[str]] = None,
|
||||
num_images_per_prompt: int = 1,
|
||||
max_sequence_length: int = 128,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
|
||||
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",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
prompt_attention_mask = text_inputs.attention_mask
|
||||
prompt_attention_mask = prompt_attention_mask.to(device)
|
||||
|
||||
prompt_embeds = self.text_encoder_2(text_input_ids.to(device), attention_mask=prompt_attention_mask, 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
|
||||
|
||||
# duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
|
||||
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)
|
||||
prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
|
||||
prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 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,
|
||||
):
|
||||
|
||||
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
|
||||
prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
|
||||
|
||||
# Use pooled output of CLIPTextModel
|
||||
prompt_embeds = prompt_embeds.pooler_output
|
||||
prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
|
||||
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
|
||||
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
|
||||
|
||||
return prompt_embeds
|
||||
|
||||
def encode_prompt(self,
|
||||
prompt,
|
||||
num_images_per_prompt=1,
|
||||
device=None,
|
||||
):
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
|
||||
batch_size = len(prompt)
|
||||
|
||||
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,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
device=device,
|
||||
)
|
||||
|
||||
return prompt_embeds, prompt_attention_mask, pooled_prompt_embeds
|
||||
|
||||
def forward(self, input_prompts, device):
|
||||
with torch.no_grad():
|
||||
prompt_embeds, prompt_attention_mask, pooled_prompt_embeds = self.encode_prompt(input_prompts, 1, device=device)
|
||||
|
||||
return prompt_embeds, prompt_attention_mask, pooled_prompt_embeds
|
||||
Reference in New Issue
Block a user