Upload STDiT

3485a6a verified 2 months ago

No virus

16.4 kB

	import math
	import collections.abc

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	from einops import rearrange
	from itertools import repeat
	from functools import partial
	from .utils import get_layernorm, t2i_modulate, approx_gelu


	try:
	import xformers
	HAS_XFORMERS = True
	except:
	HAS_XFORMERS = False

	# ============================
	# Attention
	# ============================
	class Attention(nn.Module):
	def __init__(
	self,
	dim: int,
	num_heads: int = 8,
	qkv_bias: bool = False,
	qk_norm: bool = False,
	attn_drop: float = 0.0,
	proj_drop: float = 0.0,
	norm_layer: nn.Module = nn.LayerNorm,
	enable_flash_attn: bool = False,
	) -> None:
	super().__init__()
	assert dim % num_heads == 0, "dim should be divisible by num_heads"
	self.dim = dim
	self.num_heads = num_heads
	self.head_dim = dim // num_heads
	self.scale = self.head_dim**-0.5
	self.enable_flash_attn = enable_flash_attn

	self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
	self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
	self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
	self.attn_drop = nn.Dropout(attn_drop)
	self.proj = nn.Linear(dim, dim)
	self.proj_drop = nn.Dropout(proj_drop)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	B, N, C = x.shape
	qkv = self.qkv(x)
	qkv_shape = (B, N, 3, self.num_heads, self.head_dim)
	if self.enable_flash_attn:
	qkv_permute_shape = (2, 0, 1, 3, 4)
	else:
	qkv_permute_shape = (2, 0, 3, 1, 4)
	qkv = qkv.view(qkv_shape).permute(qkv_permute_shape)
	q, k, v = qkv.unbind(0)
	q, k = self.q_norm(q), self.k_norm(k)
	if self.enable_flash_attn:
	from flash_attn import flash_attn_func

	x = flash_attn_func(
	q,
	k,
	v,
	dropout_p=self.attn_drop.p if self.training else 0.0,
	softmax_scale=self.scale,
	)
	else:
	dtype = q.dtype
	q = q * self.scale
	attn = q @ k.transpose(-2, -1) # translate attn to float32
	attn = attn.to(torch.float32)
	attn = attn.softmax(dim=-1)
	attn = attn.to(dtype) # cast back attn to original dtype
	attn = self.attn_drop(attn)
	x = attn @ v

	x_output_shape = (B, N, C)
	if not self.enable_flash_attn:
	x = x.transpose(1, 2)
	x = x.reshape(x_output_shape)
	x = self.proj(x)
	x = self.proj_drop(x)
	return x


	# ============================
	# Caption Embedding
	# ============================
	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, None], self.y_embedding, caption)
	return caption

	def forward(self, caption, train, force_drop_ids=None):
	if train:
	assert caption.shape[2:] == 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


	# ============================
	# Drop Path
	# ============================
	def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
	"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

	This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
	the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
	See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
	changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
	'survival rate' as the argument.

	"""
	if drop_prob == 0. or not training:
	return x
	keep_prob = 1 - drop_prob
	shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
	random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
	if keep_prob > 0.0 and scale_by_keep:
	random_tensor.div_(keep_prob)
	return x * random_tensor

	class DropPath(nn.Module):
	"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
	"""
	def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
	super(DropPath, self).__init__()
	self.drop_prob = drop_prob
	self.scale_by_keep = scale_by_keep

	def forward(self, x):
	return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)

	def extra_repr(self):
	return f'drop_prob={round(self.drop_prob,3):0.3f}'


	# ============================
	# MHA
	# ============================
	class MultiHeadCrossAttention(nn.Module):
	def __init__(self, d_model, num_heads, attn_drop=0.0, proj_drop=0.0):
	super(MultiHeadCrossAttention, self).__init__()
	assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

	self.d_model = d_model
	self.num_heads = num_heads
	self.head_dim = d_model // num_heads

	self.q_linear = nn.Linear(d_model, d_model)
	self.kv_linear = nn.Linear(d_model, d_model * 2)
	self.attn_drop = nn.Dropout(attn_drop)
	self.proj = nn.Linear(d_model, d_model)
	self.proj_drop = nn.Dropout(proj_drop)

	def forward(self, x, cond, mask=None):
	# query/value: img tokens; key: condition; mask: if padding tokens
	B, N, C = x.shape

	q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim)
	kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim)
	k, v = kv.unbind(2)

	attn_bias = None
	assert HAS_XFORMERS, "Please install xformers to use this module."
	if mask is not None:
	attn_bias = xformers.ops.fmha.BlockDiagonalMask.from_seqlens([N] * B, mask)
	x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias)

	x = x.view(B, -1, C)
	x = self.proj(x)
	x = self.proj_drop(x)
	return x


	# ============================
	# MLP
	# ============================
	def _ntuple(n):
	def parse(x):
	if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
	return tuple(x)
	return tuple(repeat(x, n))
	return parse

	to_2tuple = _ntuple(2)

	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,
	norm_layer=None,
	bias=True,
	drop=0.,
	use_conv=False,
	):
	super().__init__()
	out_features = out_features or in_features
	hidden_features = hidden_features or in_features
	bias = to_2tuple(bias)
	drop_probs = to_2tuple(drop)
	linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear

	self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0])
	self.act = act_layer()
	self.drop1 = nn.Dropout(drop_probs[0])
	self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity()
	self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1])
	self.drop2 = nn.Dropout(drop_probs[1])

	def forward(self, x):
	x = self.fc1(x)
	x = self.act(x)
	x = self.drop1(x)
	x = self.norm(x)
	x = self.fc2(x)
	x = self.drop2(x)
	return x


	# ============================
	# Patch Embedding
	# ============================
	class PatchEmbed3D(nn.Module):
	"""Video to Patch Embedding.

	Args:
	patch_size (int): Patch token size. Default: (2,4,4).
	in_chans (int): Number of input video channels. Default: 3.
	embed_dim (int): Number of linear projection output channels. Default: 96.
	norm_layer (nn.Module, optional): Normalization layer. Default: None
	"""

	def __init__(
	self,
	patch_size=(2, 4, 4),
	in_chans=3,
	embed_dim=96,
	norm_layer=None,
	flatten=True,
	):
	super().__init__()
	self.patch_size = patch_size
	self.flatten = flatten

	self.in_chans = in_chans
	self.embed_dim = embed_dim

	self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
	if norm_layer is not None:
	self.norm = norm_layer(embed_dim)
	else:
	self.norm = None

	def forward(self, x):
	"""Forward function."""
	# padding
	_, _, D, H, W = x.size()
	if W % self.patch_size[2] != 0:
	x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
	if H % self.patch_size[1] != 0:
	x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
	if D % self.patch_size[0] != 0:
	x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))

	x = self.proj(x) # (B C T H W)
	if self.norm is not None:
	D, Wh, Ww = x.size(2), x.size(3), x.size(4)
	x = x.flatten(2).transpose(1, 2)
	x = self.norm(x)
	x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)
	if self.flatten:
	x = x.flatten(2).transpose(1, 2) # BCTHW -> BNC
	return x


	# ============================
	# T2I
	# ============================
	def t2i_modulate(x, shift, scale):
	return x * (1 + scale) + shift

	class T2IFinalLayer(nn.Module):
	"""
	The final layer of PixArt.
	"""

	def __init__(self, hidden_size, num_patch, out_channels):
	super().__init__()
	self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
	self.linear = nn.Linear(hidden_size, num_patch * out_channels, bias=True)
	self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
	self.out_channels = out_channels

	def forward(self, x, t):
	shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1)
	x = t2i_modulate(self.norm_final(x), shift, scale)
	x = self.linear(x)
	return x


	# ============================
	# Time Step Embedding
	# ============================
	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)
	freqs = freqs.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, dtype):
	t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
	if t_freq.dtype != dtype:
	t_freq = t_freq.to(dtype)
	t_emb = self.mlp(t_freq)
	return t_emb


	# ============================
	# STDiT Block
	# ============================
	class STDiTBlock(nn.Module):
	def __init__(
	self,
	hidden_size,
	num_heads,
	d_s=None,
	d_t=None,
	mlp_ratio=4.0,
	drop_path=0.0,
	enable_flash_attn=False,
	enable_layernorm_kernel=False,
	enable_sequence_parallelism=False,
	):
	super().__init__()
	self.hidden_size = hidden_size
	self.enable_flash_attn = enable_flash_attn
	self._enable_sequence_parallelism = enable_sequence_parallelism

	if enable_sequence_parallelism:
	self.attn_cls = SeqParallelAttention
	self.mha_cls = SeqParallelMultiHeadCrossAttention
	else:
	self.attn_cls = Attention
	self.mha_cls = MultiHeadCrossAttention

	self.norm1 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
	self.attn = self.attn_cls(
	hidden_size,
	num_heads=num_heads,
	qkv_bias=True,
	enable_flash_attn=enable_flash_attn,
	)
	self.cross_attn = self.mha_cls(hidden_size, num_heads)
	self.norm2 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
	self.mlp = Mlp(
	in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0
	)
	self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
	self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size**0.5)

	# temporal attention
	self.d_s = d_s
	self.d_t = d_t

	if self._enable_sequence_parallelism:
	sp_size = dist.get_world_size(get_sequence_parallel_group())
	# make sure d_t is divisible by sp_size
	assert d_t % sp_size == 0
	self.d_t = d_t // sp_size

	self.attn_temp = self.attn_cls(
	hidden_size,
	num_heads=num_heads,
	qkv_bias=True,
	enable_flash_attn=self.enable_flash_attn,
	)

	def forward(self, x, y, t, mask=None, tpe=None):
	B, N, C = x.shape

	shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
	self.scale_shift_table[None] + t.reshape(B, 6, -1)
	).chunk(6, dim=1)
	x_m = t2i_modulate(self.norm1(x), shift_msa, scale_msa)

	# spatial branch
	x_s = rearrange(x_m, "B (T S) C -> (B T) S C", T=self.d_t, S=self.d_s)
	x_s = self.attn(x_s)
	x_s = rearrange(x_s, "(B T) S C -> B (T S) C", T=self.d_t, S=self.d_s)
	x = x + self.drop_path(gate_msa * x_s)

	# temporal branch
	x_t = rearrange(x, "B (T S) C -> (B S) T C", T=self.d_t, S=self.d_s)
	if tpe is not None:
	x_t = x_t + tpe
	x_t = self.attn_temp(x_t)
	x_t = rearrange(x_t, "(B S) T C -> B (T S) C", T=self.d_t, S=self.d_s)
	x = x + self.drop_path(gate_msa * x_t)

	# cross attn
	x = x + self.cross_attn(x, y, mask)

	# mlp
	x = x + self.drop_path(gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp)))

	return x