304 lines
11 KiB
Python
304 lines
11 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
# All rights reserved.
|
|
# This file contains code that is adapted from
|
|
# timm: https://github.com/huggingface/pytorch-image-models
|
|
# pixart: https://github.com/PixArt-alpha/PixArt-alpha
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from einops import rearrange
|
|
|
|
from scepter.modules.model.backbone.transformer.attention import drop_path
|
|
|
|
|
|
def modulate(x, shift, scale, unsqueeze=False):
|
|
if unsqueeze:
|
|
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
|
else:
|
|
return x * (1 + scale) + shift
|
|
|
|
|
|
class DropPath(nn.Module):
|
|
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
|
"""
|
|
def __init__(self, drop_prob=None):
|
|
super(DropPath, self).__init__()
|
|
self.drop_prob = drop_prob
|
|
|
|
def forward(self, x):
|
|
return drop_path(x, self.drop_prob, self.training)
|
|
|
|
|
|
class MaskFinalLayer(nn.Module):
|
|
"""
|
|
The final layer of PixArt.
|
|
"""
|
|
def __init__(self, final_hidden_size, c_emb_size, patch_size,
|
|
out_channels):
|
|
super().__init__()
|
|
self.norm_final = nn.LayerNorm(final_hidden_size,
|
|
elementwise_affine=False,
|
|
eps=1e-6)
|
|
self.linear = nn.Linear(final_hidden_size,
|
|
patch_size * patch_size * out_channels,
|
|
bias=True)
|
|
self.adaLN_modulation = nn.Sequential(
|
|
nn.SiLU(), nn.Linear(c_emb_size, 2 * final_hidden_size, bias=True))
|
|
|
|
def forward(self, x, t):
|
|
shift, scale = self.adaLN_modulation(t).chunk(2, dim=1)
|
|
x = modulate(self.norm_final(x), shift, scale, unsqueeze=True)
|
|
x = self.linear(x)
|
|
return x
|
|
|
|
|
|
class DecoderLayer(nn.Module):
|
|
"""
|
|
The final layer of PixArt.
|
|
"""
|
|
def __init__(self, hidden_size, decoder_hidden_size):
|
|
super().__init__()
|
|
self.norm_decoder = nn.LayerNorm(hidden_size,
|
|
elementwise_affine=False,
|
|
eps=1e-6)
|
|
self.linear = nn.Linear(hidden_size, decoder_hidden_size, bias=True)
|
|
self.adaLN_modulation = nn.Sequential(
|
|
nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
|
|
|
|
def forward(self, x, t):
|
|
shift, scale = self.adaLN_modulation(t).chunk(2, dim=1)
|
|
x = modulate(self.norm_decoder(x), shift, scale, unsqueeze=True)
|
|
x = self.linear(x)
|
|
return x
|
|
|
|
|
|
class TimestepEmbedder(nn.Module):
|
|
"""
|
|
Embeds scalar timesteps into vector representations.
|
|
"""
|
|
def __init__(self, hidden_size, frequency_embedding_size=256):
|
|
super().__init__()
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size, hidden_size, bias=True),
|
|
)
|
|
self.frequency_embedding_size = frequency_embedding_size
|
|
|
|
@staticmethod
|
|
def timestep_embedding(t, dim, max_period=10000):
|
|
"""
|
|
Create sinusoidal timestep embeddings.
|
|
:param t: a 1-D Tensor of N indices, one per batch element.
|
|
These may be fractional.
|
|
:param dim: the dimension of the output.
|
|
:param max_period: controls the minimum frequency of the embeddings.
|
|
:return: an (N, D) Tensor of positional embeddings.
|
|
"""
|
|
# https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
|
half = dim // 2
|
|
freqs = torch.exp(
|
|
-math.log(max_period) *
|
|
torch.arange(start=0, end=half, dtype=torch.float32) /
|
|
half).to(device=t.device)
|
|
args = t[:, None].float() * freqs[None]
|
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
if dim % 2:
|
|
embedding = torch.cat(
|
|
[embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
|
return embedding
|
|
|
|
def forward(self, t):
|
|
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
|
t_emb = self.mlp(t_freq)
|
|
return t_emb
|
|
|
|
|
|
class SizeEmbedder(TimestepEmbedder):
|
|
"""
|
|
Embeds scalar timesteps into vector representations.
|
|
"""
|
|
def __init__(self, hidden_size, frequency_embedding_size=256):
|
|
super().__init__(hidden_size=hidden_size,
|
|
frequency_embedding_size=frequency_embedding_size)
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size, hidden_size, bias=True),
|
|
)
|
|
self.frequency_embedding_size = frequency_embedding_size
|
|
self.outdim = hidden_size
|
|
|
|
def forward(self, s, bs):
|
|
if s.ndim == 1:
|
|
s = s[:, None]
|
|
assert s.ndim == 2
|
|
if s.shape[0] != bs:
|
|
s = s.repeat(bs // s.shape[0], 1)
|
|
assert s.shape[0] == bs
|
|
b, dims = s.shape[0], s.shape[1]
|
|
s = rearrange(s, 'b d -> (b d)')
|
|
s_freq = self.timestep_embedding(s, self.frequency_embedding_size).to(
|
|
self.dtype)
|
|
s_emb = self.mlp(s_freq)
|
|
s_emb = rearrange(s_emb,
|
|
'(b d) d2 -> b (d d2)',
|
|
b=b,
|
|
d=dims,
|
|
d2=self.outdim)
|
|
return s_emb
|
|
|
|
@property
|
|
def dtype(self):
|
|
# 返回模型参数的数据类型
|
|
return next(self.parameters()).dtype
|
|
|
|
|
|
class LabelEmbedder(nn.Module):
|
|
"""
|
|
Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
|
|
"""
|
|
def __init__(self, num_classes, hidden_size, dropout_prob):
|
|
super().__init__()
|
|
use_cfg_embedding = dropout_prob > 0
|
|
self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding,
|
|
hidden_size)
|
|
self.num_classes = num_classes
|
|
self.dropout_prob = dropout_prob
|
|
|
|
def token_drop(self, labels, force_drop_ids=None):
|
|
"""
|
|
Drops labels to enable classifier-free guidance.
|
|
"""
|
|
if force_drop_ids is None:
|
|
drop_ids = torch.rand(labels.shape[0]).cuda() < self.dropout_prob
|
|
else:
|
|
drop_ids = force_drop_ids == 1
|
|
labels = torch.where(drop_ids, self.num_classes, labels)
|
|
return labels
|
|
|
|
def forward(self, labels, train, force_drop_ids=None):
|
|
use_dropout = self.dropout_prob > 0
|
|
if (train and use_dropout) or (force_drop_ids is not None):
|
|
labels = self.token_drop(labels, force_drop_ids)
|
|
return self.embedding_table(labels)
|
|
|
|
|
|
class CaptionEmbedder(nn.Module):
|
|
"""
|
|
Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
|
|
"""
|
|
def __init__(self,
|
|
in_channels,
|
|
hidden_size,
|
|
uncond_prob,
|
|
act_layer=nn.GELU(approximate='tanh'),
|
|
token_num=120):
|
|
super().__init__()
|
|
self.y_proj = Mlp(in_features=in_channels,
|
|
hidden_features=hidden_size,
|
|
out_features=hidden_size,
|
|
act_layer=act_layer,
|
|
drop=0)
|
|
self.register_buffer(
|
|
'y_embedding',
|
|
nn.Parameter(
|
|
torch.randn(token_num, in_channels) / in_channels**0.5))
|
|
self.uncond_prob = uncond_prob
|
|
|
|
def token_drop(self, caption, force_drop_ids=None):
|
|
"""
|
|
Drops labels to enable classifier-free guidance.
|
|
"""
|
|
if force_drop_ids is None:
|
|
drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob
|
|
else:
|
|
drop_ids = force_drop_ids == 1
|
|
caption = torch.where(drop_ids[:, None, None], self.y_embedding,
|
|
caption)
|
|
return caption
|
|
|
|
def forward(self, caption, train, force_drop_ids=None):
|
|
if train:
|
|
assert caption.shape[1:] == self.y_embedding.shape
|
|
use_dropout = self.uncond_prob > 0
|
|
if (train and use_dropout) or (force_drop_ids is not None):
|
|
caption = self.token_drop(caption, force_drop_ids)
|
|
caption = self.y_proj(caption)
|
|
return caption
|
|
|
|
|
|
class CaptionEmbedderDoubleBr(nn.Module):
|
|
"""
|
|
Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
|
|
"""
|
|
def __init__(self,
|
|
in_channels,
|
|
hidden_size,
|
|
uncond_prob,
|
|
act_layer=nn.GELU(approximate='tanh'),
|
|
token_num=120):
|
|
super().__init__()
|
|
self.proj = Mlp(in_features=in_channels,
|
|
hidden_features=hidden_size,
|
|
out_features=hidden_size,
|
|
act_layer=act_layer,
|
|
drop=0)
|
|
self.embedding = nn.Parameter(torch.randn(1, in_channels) / 10**0.5)
|
|
self.y_embedding = nn.Parameter(
|
|
torch.randn(token_num, in_channels) / 10**0.5)
|
|
self.uncond_prob = uncond_prob
|
|
|
|
def token_drop(self, global_caption, caption, force_drop_ids=None):
|
|
"""
|
|
Drops labels to enable classifier-free guidance.
|
|
"""
|
|
if force_drop_ids is None:
|
|
drop_ids = torch.rand(
|
|
global_caption.shape[0]).cuda() < self.uncond_prob
|
|
else:
|
|
drop_ids = force_drop_ids == 1
|
|
global_caption = torch.where(drop_ids[:, None], self.embedding,
|
|
global_caption)
|
|
caption = torch.where(drop_ids[:, None, None, None], self.y_embedding,
|
|
caption)
|
|
return global_caption, caption
|
|
|
|
def forward(self, caption, train, force_drop_ids=None):
|
|
assert caption.shape[2:] == self.y_embedding.shape
|
|
global_caption = caption.mean(dim=2).squeeze()
|
|
use_dropout = self.uncond_prob > 0
|
|
if (train and use_dropout) or (force_drop_ids is not None):
|
|
global_caption, caption = self.token_drop(global_caption, caption,
|
|
force_drop_ids)
|
|
y_embed = self.proj(global_caption)
|
|
return y_embed, caption
|
|
|
|
|
|
class Mlp(nn.Module):
|
|
""" MLP as used in Vision Transformer, MLP-Mixer and related networks
|
|
"""
|
|
def __init__(self,
|
|
in_features,
|
|
hidden_features=None,
|
|
out_features=None,
|
|
act_layer=nn.GELU,
|
|
drop=0.):
|
|
super().__init__()
|
|
out_features = out_features or in_features
|
|
hidden_features = hidden_features or in_features
|
|
self.fc1 = nn.Linear(in_features, hidden_features)
|
|
self.act = act_layer()
|
|
self.fc2 = nn.Linear(hidden_features, out_features)
|
|
self.drop = nn.Dropout(drop)
|
|
|
|
def forward(self, x):
|
|
x = self.fc1(x)
|
|
x = self.act(x)
|
|
x = self.drop(x)
|
|
x = self.fc2(x)
|
|
x = self.drop(x)
|
|
return x
|