File size: 9,035 Bytes
4a9ad28 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 |
import torch
import torch.nn as nn
import torch.nn.functional as F
# This architecture was my attempt at the following Simple Diffusion paper with some modifications:
# https://arxiv.org/pdf/2410.19324v1
# Very similar to GeGLU or SwiGLU, there's a learned gate FN, uses arctan as the activation fn.
class xATGLU(nn.Module):
def __init__(self, input_dim, output_dim, bias=True):
super().__init__()
# GATE path | VALUE path
self.proj = nn.Linear(input_dim, output_dim * 2, bias=bias)
nn.init.kaiming_normal_(self.proj.weight, nonlinearity='linear')
self.alpha = nn.Parameter(torch.zeros(1))
self.half_pi = torch.pi / 2
self.inv_pi = 1 / torch.pi
def forward(self, x):
projected = self.proj(x)
gate_path, value_path = projected.chunk(2, dim=-1)
# Apply arctan gating with expanded range via learned alpha -- https://arxiv.org/pdf/2405.20768
gate = (torch.arctan(gate_path) + self.half_pi) * self.inv_pi
expanded_gate = gate * (1 + 2 * self.alpha) - self.alpha
return expanded_gate * value_path # g(x) × y
class ResBlock(nn.Module):
def __init__(self, channels):
super().__init__()
self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
self.norm1 = nn.GroupNorm(32, channels)
self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
self.norm2 = nn.GroupNorm(32, channels)
def forward(self, x):
h = self.conv1(F.silu(self.norm1(x)))
h = self.conv2(F.silu(self.norm2(h)))
return x + h
class TransformerBlock(nn.Module):
def __init__(self, channels, num_heads=8):
super().__init__()
self.norm1 = nn.LayerNorm(channels)
self.attn = nn.MultiheadAttention(channels, num_heads)
self.norm2 = nn.LayerNorm(channels)
self.mlp = nn.Sequential(
xATGLU(channels, 4 * channels),
nn.Linear(4 * channels, channels)
)
def forward(self, x):
# Reshape for attention [B, C, H, W] -> [H*W, B, C]
b, c, h, w = x.shape
spatial_size = h * w
x = x.flatten(2).permute(2, 0, 1)
# Self attention
h_attn = self.norm1(x)
h_attn, _ = self.attn(h_attn, h_attn, h_attn)
x = x + h_attn
# MLP
h_mlp = self.norm2(x)
h_mlp = self.mlp(h_mlp)
x = x + h_mlp
# Reshape back [H*W, B, C] -> [B, C, H, W]
return x.permute(1, 2, 0).reshape(b, c, h, w)
class LevelBlock(nn.Module):
def __init__(self, channels, num_blocks, block_type='res'):
super().__init__()
self.blocks = nn.ModuleList()
for _ in range(num_blocks):
if block_type == 'transformer':
self.blocks.append(TransformerBlock(channels))
else:
self.blocks.append(ResBlock(channels))
def forward(self, x):
for block in self.blocks:
x = block(x)
return x
class AsymmetricResidualUDiT(nn.Module):
def __init__(self,
in_channels=3, # Input color channels
base_channels=128, # Initial feature size, dramatically increases parameter size of network.
patch_size=2, # Smaller patches dramatically increases flops and compute expenses. Recommend >=4 unless you have real compute.
num_levels=3, # Feature downsample, essentially the unet depth -- so we down/upsample three times. Dramatically increases parameters as you increase.
encoder_blocks=3, # Can be different number of blocks VS decoder_blocks
decoder_blocks=7, # Can be different number of blocks VS encoder_blocks
encoder_transformer_thresh=2, #When to start using transformer blocks instead of res blocks in the encoder. (>=)
decoder_transformer_thresh=4, #When to stop using transformer blocks instead of res blocks in the decoder. (<=)
mid_blocks=16 # Number of middle transformer blocks. Relatively cheap as this is at the bottom of the unet feature bottleneck.
):
super().__init__()
# Initial projection from image space
self.patch_embed = nn.Conv2d(in_channels, base_channels,
kernel_size=patch_size, stride=patch_size)
# Create encoder levels
self.encoders = nn.ModuleList()
curr_channels = base_channels
for level in range(num_levels):
# Create the main processing blocks for this level
use_transformer = level >= encoder_transformer_thresh # Use transformers for latter levels
# Encoder blocks -- encoder_blocks
self.encoders.append(
LevelBlock(curr_channels, encoder_blocks, use_transformer)
)
# Add channel scaling for next level
# Doubles the size of the feature space for each step, except for the last level.
if level < num_levels - 1:
self.encoders.append(
nn.Conv2d(curr_channels, curr_channels * 2, 1)
)
curr_channels *= 2
# Middle transformer blocks -- mid_blocks
self.middle = nn.ModuleList([
TransformerBlock(curr_channels) for _ in range(mid_blocks)
])
# Create decoder levels
self.decoders = nn.ModuleList()
for level in range(num_levels):
# Create the main processing blocks for this level
use_transformer = level <= decoder_transformer_thresh # Use transformers for early levels (inverse of encoder)
# Decoder blocks -- decoder_blocks
self.decoders.append(
LevelBlock(curr_channels, decoder_blocks, use_transformer)
)
# Add channel scaling for next level
# Halves the size of the feature space for each step, except for the last level.
if level < num_levels - 1:
self.decoders.append(
nn.Conv2d(curr_channels, curr_channels // 2, 1)
)
curr_channels //= 2
# Final projection back to image space
self.final_proj = nn.ConvTranspose2d(base_channels, in_channels,
kernel_size=patch_size, stride=patch_size)
def downsample(self, x):
return F.avg_pool2d(x, kernel_size=2)
def upsample(self, x):
return F.interpolate(x, scale_factor=2, mode='nearest')
def forward(self, x, t=None):
# Start by patch embedding the inputs.
x = self.patch_embed(x)
# Track residual path and features at each spatial level
# The paper was very specific about the residual flow path, I tried my best to copy how they described it.
# *Per resolution e.g. per num_level resolution block more or less
# f(x) = fu( U(fm(D(h)) - D(h)) + h ) where h = fd(x)
#
# Where
# 1. h = fd(x) : Encoder path processes input
# 2. D(h) : Downsample the encoded features
# 3. fm(D(h)) : Middle transformer blocks process downsampled features
# 4. fm(D(h))-D(h): Subtract original downsampled features (residual connection)
# 5. U(...) : Upsample the processed features
# 6. ... + h : Add back original encoder features (skip connection)
# 7. fu(...) : Decoder path processes the combined features
residuals = []
curr_res = x
# Encoder path (computing h = fd(x))
h = x
for i, blocks in enumerate(self.encoders):
if isinstance(blocks, LevelBlock):
h = blocks(h)
else:
# Save residual before downsampling
residuals.append(curr_res)
# Downsample and update current residual
h = self.downsample(blocks(h))
curr_res = h
# Middle blocks (fm)
x = h
for block in self.middle:
x = block(x)
# Subtract the residual at this level (D(h))
x = x - curr_res
# Decoder path (fu)
for i, blocks in enumerate(self.decoders):
if isinstance(blocks, LevelBlock):
x = blocks(x)
else:
# Channel reduction
x = blocks(x)
# Upsample
x = self.upsample(x)
# Add residual from encoder at this level, LIFO, last residual added is the first we want, since it's this u-shape.
curr_res = residuals.pop()
x = x + curr_res
# Final projection
x = self.final_proj(x)
return x |