|
import tensorflow as tf |
|
from tensorflow.keras.layers import Conv2d,Dense,Dropout,LayerNormalization,Activation |
|
from tensorflow.keras.initializers import RandomNormal |
|
from tensorflow.keras import Model |
|
import collections.abc |
|
from itertools import repeat |
|
from typing import Optional |
|
import numpy as np |
|
import math |
|
|
|
|
|
def modulate(x, shift, scale): |
|
return x * (1 + tf.expand_dims(scale, 1)) + tf.expand_dims(shift, 1) |
|
|
|
|
|
|
|
|
|
|
|
|
|
class TimestepEmbedder: |
|
""" |
|
Embeds scalar timesteps into vector representations. |
|
""" |
|
def __init__(self, hidden_size, frequency_embedding_size=256): |
|
self.mlp = tf.keras.Sequential() |
|
self.mlp.add(Dense(hidden_size, kernel_initializer=RandomNormal(stddev=0.02), use_bias=True)) |
|
self.mlp.add(Activation('silu')) |
|
self.mlp.add(Dense(hidden_size, kernel_initializer=RandomNormal(stddev=0.02), use_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. |
|
""" |
|
|
|
half = dim // 2 |
|
freqs = tf.math.exp( |
|
-math.log(max_period) * tf.range(start=0, limit=half, dtype=tf.float32) / half |
|
) |
|
args = tf.cast(t[:, None], 'float32') * freqs[None] |
|
embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], axis=-1) |
|
if dim % 2: |
|
embedding = tf.concat([embedding, tf.zeros_like(embedding[:, :1])], axis=-1) |
|
return embedding |
|
|
|
def __call__(self, t): |
|
t_freq = self.timestep_embedding(t, self.frequency_embedding_size) |
|
t_emb = self.mlp(t_freq) |
|
return t_emb |
|
|
|
|
|
class LabelEmbedder(tf.keras.layers.Layer): |
|
""" |
|
Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance. |
|
""" |
|
def __init__(self, num_classes, hidden_size, dropout_prob): |
|
use_cfg_embedding = dropout_prob > 0 |
|
self.embedding_table = self.add_weight( |
|
name='embedding_table', |
|
shape=(num_classes + use_cfg_embedding, hidden_size), |
|
initializer=tf.keras.initializers.RandomNormal(stddev=0.02), |
|
trainable=True |
|
) |
|
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 = tf.random.uniform([labels.shape[0]]) < self.dropout_prob |
|
else: |
|
drop_ids = force_drop_ids == 1 |
|
labels = tf.where(drop_ids, self.num_classes, labels) |
|
return labels |
|
|
|
def __call__(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) |
|
embeddings = tf.gather(self.embedding_table, labels) |
|
return embeddings |
|
|
|
|
|
|
|
|
|
|
|
|
|
class DiTBlock: |
|
""" |
|
A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning. |
|
""" |
|
def __init__(self, hidden_size, num_heads, mlp_ratio=4.0): |
|
self.norm1 = LayerNormalization(epsilon=1e-6) |
|
self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True) |
|
self.norm2 = LayerNormalization(epsilon=1e-6) |
|
mlp_hidden_dim = int(hidden_size * mlp_ratio) |
|
self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, drop=0) |
|
self.adaLN_modulation = tf.keras.Sequential() |
|
self.adaLN_modulation.add(Activation('silu')) |
|
self.adaLN_modulation.add(Dense(6 * hidden_size, kernel_initializer='zeros', use_bias=True)) |
|
|
|
def __call__(self, x, c): |
|
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = tf.split(self.adaLN_modulation(c), num_or_size_splits=6, axis=1) |
|
x = x + tf.expand_dims(gate_msa, 1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa)) |
|
x = x + tf.expand_dims(gate_mlp, 1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp)) |
|
return x |
|
|
|
|
|
class FinalLayer: |
|
""" |
|
The final layer of DiT. |
|
""" |
|
def __init__(self, hidden_size, patch_size, out_channels): |
|
self.norm_final = LayerNormalization(epsilon=1e-6) |
|
self.linear = Dense(patch_size * patch_size * out_channels, kernel_initializer='zeros', use_bias=True) |
|
self.adaLN_modulation = tf.keras.Sequential() |
|
self.adaLN_modulation.add(Activation('silu')) |
|
self.adaLN_modulation.add(Dense(2 * hidden_size, kernel_initializer='zeros', use_bias=True)) |
|
|
|
def __call__(self, x, c): |
|
shift, scale = tf.split(self.adaLN_modulation(c), num_or_size_splits=2, axis=1) |
|
x = modulate(self.norm_final(x), shift, scale) |
|
x = self.linear(x) |
|
return x |
|
|
|
|
|
class DiT(Model): |
|
""" |
|
Diffusion model with a Transformer backbone. |
|
""" |
|
def __init__( |
|
self, |
|
input_size=32, |
|
patch_size=2, |
|
in_channels=4, |
|
hidden_size=1152, |
|
depth=28, |
|
num_heads=16, |
|
mlp_ratio=4.0, |
|
class_dropout_prob=0.1, |
|
num_classes=1000, |
|
learn_sigma=True, |
|
): |
|
super(DiT, self).__init__() |
|
self.learn_sigma = learn_sigma |
|
self.in_channels = in_channels |
|
self.out_channels = in_channels * 2 if learn_sigma else in_channels |
|
self.patch_size = patch_size |
|
self.num_heads = num_heads |
|
|
|
self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True) |
|
self.t_embedder = TimestepEmbedder(hidden_size) |
|
self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob) |
|
num_patches = self.x_embedder.num_patches |
|
|
|
self.pos_embed = self.add_weight( |
|
name='pos_embed', |
|
shape=(1, num_patches, hidden_size), |
|
initializer=tf.keras.initializers.Zeros(), |
|
trainable=False |
|
) |
|
|
|
self.blocks = [ |
|
DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth) |
|
] |
|
self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels) |
|
self.initialize_weights() |
|
|
|
def initialize_weights(self): |
|
|
|
pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5)) |
|
self.pos_embed.assign(tf.convert_to_tensor(pos_embed, dtype=tf.float32)[tf.newaxis, :]) |
|
|
|
def unpatchify(self, x): |
|
""" |
|
x: (N, T, patch_size**2 * C) |
|
imgs: (N, H, W, C) |
|
""" |
|
c = self.out_channels |
|
p = self.x_embedder.patch_size[0] |
|
h = w = int(x.shape[1] ** 0.5) |
|
assert h * w == x.shape[1] |
|
|
|
x = tf.reshape(x, (x.shape[0], h, w, p, p, c)) |
|
x = tf.einsum('nhwpqc->nchpwq', x) |
|
imgs = tf.reshape(x, (x.shape[0], h * p, h * p, c)) |
|
return imgs |
|
|
|
def __call__(self, x, t, y): |
|
""" |
|
Forward pass of DiT. |
|
x: (N, H, W, C) tensor of spatial inputs (images or latent representations of images) |
|
t: (N,) tensor of diffusion timesteps |
|
y: (N,) tensor of class labels |
|
""" |
|
x = self.x_embedder(x) + self.pos_embed |
|
t = self.t_embedder(t) |
|
y = self.y_embedder(y, self.training) |
|
c = t + y |
|
for block in self.blocks: |
|
x = block(x, c) |
|
x = self.final_layer(x, c) |
|
x = self.unpatchify(x) |
|
return x |
|
|
|
def forward_with_cfg(self, x, t, y, cfg_scale): |
|
""" |
|
Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance. |
|
""" |
|
|
|
half = x[: len(x) // 2] |
|
combined = tf.concat([half, half], axis=0) |
|
model_out = self.forward(combined, t, y) |
|
|
|
|
|
|
|
|
|
eps, rest = model_out[:, :3], model_out[:, 3:] |
|
cond_eps, uncond_eps = tf.split(eps, len(eps) // 2, dim=0) |
|
half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) |
|
eps = tf.concat([half_eps, half_eps], axis=0) |
|
return tf.concat([eps, rest], axis=1) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0): |
|
""" |
|
grid_size: int of the grid height and width |
|
return: |
|
pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) |
|
""" |
|
grid_h = np.arange(grid_size, dtype=np.float32) |
|
grid_w = np.arange(grid_size, dtype=np.float32) |
|
grid = np.meshgrid(grid_w, grid_h) |
|
grid = np.stack(grid, axis=0) |
|
|
|
grid = grid.reshape([2, 1, grid_size, grid_size]) |
|
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) |
|
if cls_token and extra_tokens > 0: |
|
pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) |
|
return pos_embed |
|
|
|
|
|
def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): |
|
assert embed_dim % 2 == 0 |
|
|
|
|
|
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) |
|
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) |
|
|
|
emb = np.concatenate([emb_h, emb_w], axis=1) |
|
return emb |
|
|
|
|
|
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): |
|
""" |
|
embed_dim: output dimension for each position |
|
pos: a list of positions to be encoded: size (M,) |
|
out: (M, D) |
|
""" |
|
assert embed_dim % 2 == 0 |
|
omega = np.arange(embed_dim // 2, dtype=np.float64) |
|
omega /= embed_dim / 2. |
|
omega = 1. / 10000**omega |
|
|
|
pos = pos.reshape(-1) |
|
out = np.einsum('m,d->md', pos, omega) |
|
|
|
emb_sin = np.sin(out) |
|
emb_cos = np.cos(out) |
|
|
|
emb = np.concatenate([emb_sin, emb_cos], axis=1) |
|
return emb |
|
|
|
|
|
|
|
|
|
|
|
|
|
def DiT_XL_2(): |
|
return DiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16) |
|
|
|
def DiT_XL_4(): |
|
return DiT(depth=28, hidden_size=1152, patch_size=4, num_heads=16) |
|
|
|
def DiT_XL_8(): |
|
return DiT(depth=28, hidden_size=1152, patch_size=8, num_heads=16) |
|
|
|
def DiT_L_2(): |
|
return DiT(depth=24, hidden_size=1024, patch_size=2, num_heads=16) |
|
|
|
def DiT_L_4(): |
|
return DiT(depth=24, hidden_size=1024, patch_size=4, num_heads=16) |
|
|
|
def DiT_L_8(): |
|
return DiT(depth=24, hidden_size=1024, patch_size=8, num_heads=16) |
|
|
|
def DiT_B_2(): |
|
return DiT(depth=12, hidden_size=768, patch_size=2, num_heads=12) |
|
|
|
def DiT_B_4(): |
|
return DiT(depth=12, hidden_size=768, patch_size=4, num_heads=12) |
|
|
|
def DiT_B_8(): |
|
return DiT(depth=12, hidden_size=768, patch_size=8, num_heads=12) |
|
|
|
def DiT_S_2(): |
|
return DiT(depth=12, hidden_size=384, patch_size=2, num_heads=6) |
|
|
|
def DiT_S_4(): |
|
return DiT(depth=12, hidden_size=384, patch_size=4, num_heads=6) |
|
|
|
def DiT_S_8(): |
|
return DiT(depth=12, hidden_size=384, patch_size=8, num_heads=6) |
|
|
|
|
|
DiT_models = { |
|
'DiT-XL/2': DiT_XL_2, 'DiT-XL/4': DiT_XL_4, 'DiT-XL/8': DiT_XL_8, |
|
'DiT-L/2': DiT_L_2, 'DiT-L/4': DiT_L_4, 'DiT-L/8': DiT_L_8, |
|
'DiT-B/2': DiT_B_2, 'DiT-B/4': DiT_B_4, 'DiT-B/8': DiT_B_8, |
|
'DiT-S/2': DiT_S_2, 'DiT-S/4': DiT_S_4, 'DiT-S/8': DiT_S_8, |
|
} |
|
|
|
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 PatchEmbed: |
|
""" 2D Image to Patch Embedding |
|
""" |
|
def __init__( |
|
self, |
|
img_size: Optional[int] = 224, |
|
patch_size: int = 16, |
|
in_chans: int = 3, |
|
embed_dim: int = 768, |
|
flatten: bool = True, |
|
bias: bool = True, |
|
): |
|
self.patch_size = to_2tuple(patch_size) |
|
if img_size is not None: |
|
self.img_size = to_2tuple(img_size) |
|
self.grid_size = tuple([s // p for s, p in zip(self.img_size, self.patch_size)]) |
|
self.num_patches = self.grid_size[0] * self.grid_size[1] |
|
else: |
|
self.img_size = None |
|
self.grid_size = None |
|
self.num_patches = None |
|
|
|
|
|
self.flatten = flatten |
|
|
|
self.proj = Conv2d(embed_dim, kernel_size=patch_size, strides=patch_size, use_bias=bias) |
|
|
|
def __call__(self, x): |
|
x = self.proj(x) |
|
B, H, W, C = x.shape |
|
if self.flatten: |
|
x = tf.reshape(x, [B, H*W, C]) |
|
return x |
|
|
|
|
|
class Mlp: |
|
""" MLP as used in Vision Transformer, MLP-Mixer and related networks |
|
""" |
|
def __init__( |
|
self, |
|
in_features, |
|
hidden_features=None, |
|
out_features=None, |
|
act_layer=tf.nn.gelu, |
|
norm_layer=None, |
|
bias=True, |
|
drop=0., |
|
use_conv=False, |
|
): |
|
out_features = out_features or in_features |
|
hidden_features = hidden_features or in_features |
|
bias = to_2tuple(bias) |
|
drop_probs = to_2tuple(drop) |
|
|
|
self.fc1 = Dense(hidden_features, use_bias=bias[0]) |
|
self.act = act_layer |
|
self.drop1 = Dropout(drop_probs[0]) |
|
self.fc2 = Dense(out_features, use_bias=bias[1]) |
|
self.drop2 = Dropout(drop_probs[1]) |
|
|
|
def __call__(self, x): |
|
x = self.fc1(x) |
|
x = self.act(x, approximate="tanh") |
|
x = self.drop1(x) |
|
x = self.fc2(x) |
|
x = self.drop2(x) |
|
return x |
|
|
|
|
|
class Attention: |
|
def __init__( |
|
self, |
|
dim: int, |
|
num_heads: int = 8, |
|
qkv_bias: bool = False, |
|
attn_drop: float = 0., |
|
proj_drop: float = 0., |
|
): |
|
assert dim % num_heads == 0, 'dim should be divisible by num_heads' |
|
self.num_heads = num_heads |
|
self.head_dim = dim // num_heads |
|
self.scale = self.head_dim ** -0.5 |
|
|
|
self.qkv = Dense(dim * 3, use_bias=qkv_bias) |
|
self.attn_drop = Dropout(attn_drop) |
|
self.proj = Dense(dim) |
|
self.proj_drop = Dropout(proj_drop) |
|
|
|
def __call__(self, x): |
|
B, N, C = x.shape |
|
qkv = tf.transpose(tf.reshape(self.qkv(x), (B, N, 3, self.num_heads, self.head_dim)), (2, 0, 3, 1, 4)) |
|
q, k, v = tf.unstack(qkv) |
|
|
|
q = q * self.scale |
|
attn = tf.matmul(q, tf.transpose(k, (0, 1, 3, 2))) |
|
attn = tf.nn.softmax(attn) |
|
attn = self.attn_drop(attn) |
|
x = tf.matmul(attn, v) |
|
|
|
x = tf.reshape(tf.transpose(x, (0, 2, 1, 3)), (B, N, C)) |
|
x = self.proj(x) |
|
x = self.proj_drop(x) |
|
return x |