222 lines
8.1 KiB
Python
222 lines
8.1 KiB
Python
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.utils import deprecate
|
|
from diffusers.models.activations import FP32SiLU, get_activation
|
|
from diffusers.models.attention_processor import Attention
|
|
|
|
|
|
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.
|
|
|
|
:param timesteps: a 1-D Tensor of N indices, one per batch element.
|
|
These may be fractional.
|
|
:param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the
|
|
embeddings. :return: 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):
|
|
super().__init__()
|
|
self.num_channels = num_channels
|
|
self.flip_sin_to_cos = flip_sin_to_cos
|
|
self.downscale_freq_shift = downscale_freq_shift
|
|
|
|
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,
|
|
)
|
|
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_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
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class HunyuanDiTAttentionPool(nn.Module):
|
|
def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
|
|
super().__init__()
|
|
self.positional_embedding = nn.Parameter(torch.randn(spacial_dim + 1, embed_dim) / embed_dim**0.5)
|
|
self.k_proj = nn.Linear(embed_dim, embed_dim)
|
|
self.q_proj = nn.Linear(embed_dim, embed_dim)
|
|
self.v_proj = nn.Linear(embed_dim, embed_dim)
|
|
self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
|
|
self.num_heads = num_heads
|
|
|
|
def forward(self, x):
|
|
x = torch.cat([x.mean(dim=1, keepdim=True), x], dim=1)
|
|
x = x + self.positional_embedding[None, :, :].to(x.dtype)
|
|
|
|
query = self.q_proj(x[:, :1])
|
|
key = self.k_proj(x)
|
|
value = self.v_proj(x)
|
|
batch_size, _, _ = query.size()
|
|
|
|
query = query.reshape(batch_size, -1, self.num_heads, query.size(-1) // self.num_heads).transpose(1, 2) # (1, H, N, E/H)
|
|
key = key.reshape(batch_size, -1, self.num_heads, key.size(-1) // self.num_heads).transpose(1, 2) # (L+1, H, N, E/H)
|
|
value = value.reshape(batch_size, -1, self.num_heads, value.size(-1) // self.num_heads).transpose(1, 2) # (L+1, H, N, E/H)
|
|
|
|
x = F.scaled_dot_product_attention(query=query, key=key, value=value, attn_mask=None, dropout_p=0.0, is_causal=False)
|
|
x = x.transpose(1, 2).reshape(batch_size, 1, -1)
|
|
x = x.to(query.dtype)
|
|
x = self.c_proj(x)
|
|
|
|
return x.squeeze(1)
|
|
|
|
|
|
class HunyuanCombinedTimestepTextSizeStyleEmbedding(nn.Module):
|
|
def __init__(self, embedding_dim, pooled_projection_dim=1024, seq_len=256, cross_attention_dim=2048):
|
|
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.pooler = HunyuanDiTAttentionPool(
|
|
seq_len, cross_attention_dim, num_heads=8, output_dim=pooled_projection_dim
|
|
)
|
|
# Here we use a default learned embedder layer for future extension.
|
|
self.style_embedder = nn.Embedding(1, embedding_dim)
|
|
extra_in_dim = 256 * 6 + embedding_dim + pooled_projection_dim
|
|
self.extra_embedder = PixArtAlphaTextProjection(
|
|
in_features=extra_in_dim,
|
|
hidden_size=embedding_dim * 4,
|
|
out_features=embedding_dim,
|
|
act_fn="silu_fp32",
|
|
)
|
|
|
|
def forward(self, timestep, encoder_hidden_states, image_meta_size, style, hidden_dtype=None):
|
|
timesteps_proj = self.time_proj(timestep)
|
|
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, 256)
|
|
|
|
# extra condition1: text
|
|
pooled_projections = self.pooler(encoder_hidden_states) # (N, 1024)
|
|
|
|
# extra condition2: image meta size embdding
|
|
image_meta_size = get_timestep_embedding(image_meta_size.view(-1), 256, True, 0)
|
|
image_meta_size = image_meta_size.to(dtype=hidden_dtype)
|
|
image_meta_size = image_meta_size.view(-1, 6 * 256) # (N, 1536)
|
|
|
|
# extra condition3: style embedding
|
|
style_embedding = self.style_embedder(style) # (N, embedding_dim)
|
|
|
|
# Concatenate all extra vectors
|
|
extra_cond = torch.cat([pooled_projections, image_meta_size, style_embedding], dim=1)
|
|
conditioning = timesteps_emb + self.extra_embedder(extra_cond) # [B, D]
|
|
|
|
return conditioning |