# some modules/classes are copied and modified from https://github.com/mcmonkey4eva/sd3-ref
# the original code is licensed under the MIT License
# and some module/classes are contributed from KohakuBlueleaf. Thanks for the contribution!
from ast import Tuple
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from functools import partial
import math
from types import SimpleNamespace
from typing import Dict, List, Optional, Union
import einops
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from transformers import CLIPTokenizer, T5TokenizerFast
from library import custom_offloading_utils
from library.device_utils import clean_memory_on_device
from .utils import setup_logging
setup_logging()
import logging
logger = logging.getLogger(__name__)
memory_efficient_attention = None
try:
import xformers
except:
pass
try:
from xformers.ops import memory_efficient_attention
except:
memory_efficient_attention = None
# region mmdit
@dataclass
class SD3Params:
patch_size: int
depth: int
num_patches: int
pos_embed_max_size: int
adm_in_channels: int
qk_norm: Optional[str]
x_block_self_attn_layers: list[int]
context_embedder_in_features: int
context_embedder_out_features: int
model_type: str
def get_2d_sincos_pos_embed(
embed_dim,
grid_size,
scaling_factor=None,
offset=None,
):
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) # here w goes first
grid = np.stack(grid, axis=0)
if scaling_factor is not None:
grid = grid / scaling_factor
if offset is not None:
grid = grid - offset
grid = grid.reshape([2, 1, grid_size, grid_size])
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
return pos_embed
def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
assert embed_dim % 2 == 0
# use half of dimensions to encode grid_h
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
return emb
def get_scaled_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, sample_size=64, base_size=16):
"""
This function is contributed by KohakuBlueleaf. Thanks for the contribution!
Creates scaled 2D sinusoidal positional embeddings that maintain consistent relative positions
when the resolution differs from the training resolution.
Args:
embed_dim (int): Dimension of the positional embedding.
grid_size (int or tuple): Size of the position grid (H, W). If int, assumes square grid.
cls_token (bool): Whether to include class token. Defaults to False.
extra_tokens (int): Number of extra tokens (e.g., cls_token). Defaults to 0.
sample_size (int): Reference resolution (typically training resolution). Defaults to 64.
base_size (int): Base grid size used during training. Defaults to 16.
Returns:
numpy.ndarray: Positional embeddings of shape (H*W, embed_dim) or
(H*W + extra_tokens, embed_dim) if cls_token is True.
"""
# Convert grid_size to tuple if it's an integer
if isinstance(grid_size, int):
grid_size = (grid_size, grid_size)
# Create normalized grid coordinates (0 to 1)
grid_h = np.arange(grid_size[0], dtype=np.float32) / grid_size[0]
grid_w = np.arange(grid_size[1], dtype=np.float32) / grid_size[1]
# Calculate scaling factors for height and width
# This ensures that the central region matches the original resolution's embeddings
scale_h = base_size * grid_size[0] / (sample_size)
scale_w = base_size * grid_size[1] / (sample_size)
# Calculate shift values to center the original resolution's embedding region
# This ensures that the central sample_size x sample_size region has similar
# positional embeddings to the original resolution
shift_h = 1 * scale_h * (grid_size[0] - sample_size) / (2 * grid_size[0])
shift_w = 1 * scale_w * (grid_size[1] - sample_size) / (2 * grid_size[1])
# Apply scaling and shifting to create the final grid coordinates
grid_h = grid_h * scale_h - shift_h
grid_w = grid_w * scale_w - shift_w
# Create 2D grid using meshgrid (note: w goes first)
grid = np.meshgrid(grid_w, grid_h)
grid = np.stack(grid, axis=0)
# # Calculate the starting indices for the central region
# # This is used for debugging/visualization of the central region
# st_h = (grid_size[0] - sample_size) // 2
# st_w = (grid_size[1] - sample_size) // 2
# print(grid[:, st_h : st_h + sample_size, st_w : st_w + sample_size])
# Reshape grid for positional embedding calculation
grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
# Generate the sinusoidal positional embeddings
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
# Add zeros for extra tokens (e.g., [CLS] token) if required
if cls_token and extra_tokens > 0:
pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
return pos_embed
# if __name__ == "__main__":
# # This is what you get when you load SD3.5 state dict
# pos_emb = torch.from_numpy(get_scaled_2d_sincos_pos_embed(
# 1536, [384, 384], sample_size=64, base_size=16
# )).float().unsqueeze(0)
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.0
omega = 1.0 / 10000**omega # (D/2,)
pos = pos.reshape(-1) # (M,)
out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product
emb_sin = np.sin(out) # (M, D/2)
emb_cos = np.cos(out) # (M, D/2)
emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
return emb
def get_1d_sincos_pos_embed_from_grid_torch(
embed_dim,
pos,
device=None,
dtype=torch.float32,
):
omega = torch.arange(embed_dim // 2, device=device, dtype=dtype)
omega *= 2.0 / embed_dim
omega = 1.0 / 10000**omega
out = torch.outer(pos.reshape(-1), omega)
emb = torch.cat([out.sin(), out.cos()], dim=1)
return emb
def get_2d_sincos_pos_embed_torch(
embed_dim,
w,
h,
val_center=7.5,
val_magnitude=7.5,
device=None,
dtype=torch.float32,
):
small = min(h, w)
val_h = (h / small) * val_magnitude
val_w = (w / small) * val_magnitude
grid_h, grid_w = torch.meshgrid(
torch.linspace(-val_h + val_center, val_h + val_center, h, device=device, dtype=dtype),
torch.linspace(-val_w + val_center, val_w + val_center, w, device=device, dtype=dtype),
indexing="ij",
)
emb_h = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid_h, device=device, dtype=dtype)
emb_w = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid_w, device=device, dtype=dtype)
emb = torch.cat([emb_w, emb_h], dim=1) # (H*W, D)
return emb
def modulate(x, shift, scale):
if shift is None:
shift = torch.zeros_like(scale)
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
def default(x, default_value):
if x is None:
return default_value
return x
def timestep_embedding(t, dim, max_period=10000):
half = dim // 2
# freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
# device=t.device, dtype=t.dtype
# )
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)
if torch.is_floating_point(t):
embedding = embedding.to(dtype=t.dtype)
return embedding
class PatchEmbed(nn.Module):
def __init__(
self,
img_size=256,
patch_size=4,
in_channels=3,
embed_dim=512,
norm_layer=None,
flatten=True,
bias=True,
strict_img_size=True,
dynamic_img_pad=False,
):
# dynamic_img_pad and norm is omitted in SD3.5
super().__init__()
self.patch_size = patch_size
self.flatten = flatten
self.strict_img_size = strict_img_size
self.dynamic_img_pad = dynamic_img_pad
if img_size is not None:
self.img_size = img_size
self.grid_size = img_size // patch_size
self.num_patches = self.grid_size**2
else:
self.img_size = None
self.grid_size = None
self.num_patches = None
self.proj = nn.Conv2d(in_channels, embed_dim, patch_size, patch_size, bias=bias)
self.norm = nn.Identity() if norm_layer is None else norm_layer(embed_dim)
def forward(self, x):
B, C, H, W = x.shape
if self.dynamic_img_pad:
# Pad input so we won't have partial patch
pad_h = (self.patch_size - H % self.patch_size) % self.patch_size
pad_w = (self.patch_size - W % self.patch_size) % self.patch_size
x = nn.functional.pad(x, (0, pad_w, 0, pad_h), mode="reflect")
x = self.proj(x)
if self.flatten:
x = x.flatten(2).transpose(1, 2)
x = self.norm(x)
return x
# FinalLayer in mmdit.py
class UnPatch(nn.Module):
def __init__(self, hidden_size=512, patch_size=4, out_channels=3):
super().__init__()
self.patch_size = patch_size
self.c = out_channels
# eps is default in mmdit.py
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(hidden_size, patch_size**2 * out_channels)
self.adaLN_modulation = nn.Sequential(
nn.SiLU(),
nn.Linear(hidden_size, 2 * hidden_size),
)
def forward(self, x: torch.Tensor, cmod, H=None, W=None):
b, n, _ = x.shape
p = self.patch_size
c = self.c
if H is None and W is None:
w = h = int(n**0.5)
assert h * w == n
else:
h = H // p if H else n // (W // p)
w = W // p if W else n // h
assert h * w == n
shift, scale = self.adaLN_modulation(cmod).chunk(2, dim=-1)
x = modulate(self.norm_final(x), shift, scale)
x = self.linear(x)
x = x.view(b, h, w, p, p, c)
x = x.permute(0, 5, 1, 3, 2, 4).contiguous()
x = x.view(b, c, h * p, w * p)
return x
class MLP(nn.Module):
def __init__(
self,
in_features,
hidden_features=None,
out_features=None,
act_layer=lambda: nn.GELU(),
norm_layer=None,
bias=True,
use_conv=False,
):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.use_conv = use_conv
layer = partial(nn.Conv1d, kernel_size=1) if use_conv else nn.Linear
self.fc1 = layer(in_features, hidden_features, bias=bias)
self.fc2 = layer(hidden_features, out_features, bias=bias)
self.act = act_layer()
self.norm = norm_layer(hidden_features) if norm_layer else nn.Identity()
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.norm(x)
x = self.fc2(x)
return x
class TimestepEmbedding(nn.Module):
def __init__(self, hidden_size, freq_embed_size=256):
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(freq_embed_size, hidden_size),
nn.SiLU(),
nn.Linear(hidden_size, hidden_size),
)
self.freq_embed_size = freq_embed_size
def forward(self, t, dtype=None, **kwargs):
t_freq = timestep_embedding(t, self.freq_embed_size).to(dtype)
t_emb = self.mlp(t_freq)
return t_emb
class Embedder(nn.Module):
def __init__(self, input_dim, hidden_size):
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(input_dim, hidden_size),
nn.SiLU(),
nn.Linear(hidden_size, hidden_size),
)
def forward(self, x):
return self.mlp(x)
def rmsnorm(x, eps=1e-6):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
class RMSNorm(torch.nn.Module):
def __init__(
self,
dim: int,
elementwise_affine: bool = False,
eps: float = 1e-6,
device=None,
dtype=None,
):
"""
Initialize the RMSNorm normalization layer.
Args:
dim (int): The dimension of the input tensor.
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
Attributes:
eps (float): A small value added to the denominator for numerical stability.
weight (nn.Parameter): Learnable scaling parameter.
"""
super().__init__()
self.eps = eps
self.learnable_scale = elementwise_affine
if self.learnable_scale:
self.weight = nn.Parameter(torch.empty(dim, device=device, dtype=dtype))
else:
self.register_parameter("weight", None)
def forward(self, x):
"""
Forward pass through the RMSNorm layer.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The output tensor after applying RMSNorm.
"""
x = rmsnorm(x, eps=self.eps)
if self.learnable_scale:
return x * self.weight.to(device=x.device, dtype=x.dtype)
else:
return x
class SwiGLUFeedForward(nn.Module):
def __init__(
self,
dim: int,
hidden_dim: int,
multiple_of: int,
ffn_dim_multiplier: float = None,
):
super().__init__()
hidden_dim = int(2 * hidden_dim / 3)
# custom dim factor multiplier
if ffn_dim_multiplier is not None:
hidden_dim = int(ffn_dim_multiplier * hidden_dim)
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
def forward(self, x):
return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))
# Linears for SelfAttention in mmdit.py
class AttentionLinears(nn.Module):
def __init__(
self,
dim: int,
num_heads: int = 8,
qkv_bias: bool = False,
pre_only: bool = False,
qk_norm: Optional[str] = None,
):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
if not pre_only:
self.proj = nn.Linear(dim, dim)
self.pre_only = pre_only
if qk_norm == "rms":
self.ln_q = RMSNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6)
self.ln_k = RMSNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6)
elif qk_norm == "ln":
self.ln_q = nn.LayerNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6)
self.ln_k = nn.LayerNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6)
elif qk_norm is None:
self.ln_q = nn.Identity()
self.ln_k = nn.Identity()
else:
raise ValueError(qk_norm)
def pre_attention(self, x: torch.Tensor) -> torch.Tensor:
"""
output:
q, k, v: [B, L, D]
"""
B, L, C = x.shape
qkv: torch.Tensor = self.qkv(x)
q, k, v = qkv.reshape(B, L, -1, self.head_dim).chunk(3, dim=2)
q = self.ln_q(q).reshape(q.shape[0], q.shape[1], -1)
k = self.ln_k(k).reshape(q.shape[0], q.shape[1], -1)
return (q, k, v)
def post_attention(self, x: torch.Tensor) -> torch.Tensor:
assert not self.pre_only
x = self.proj(x)
return x
MEMORY_LAYOUTS = {
"torch": (
lambda x, head_dim: x.reshape(x.shape[0], x.shape[1], -1, head_dim).transpose(1, 2),
lambda x: x.transpose(1, 2).reshape(x.shape[0], x.shape[2], -1),
lambda x: (1, x, 1, 1),
),
"xformers": (
lambda x, head_dim: x.reshape(x.shape[0], x.shape[1], -1, head_dim),
lambda x: x.reshape(x.shape[0], x.shape[1], -1),
lambda x: (1, 1, x, 1),
),
"math": (
lambda x, head_dim: x.reshape(x.shape[0], x.shape[1], -1, head_dim).transpose(1, 2),
lambda x: x.transpose(1, 2).reshape(x.shape[0], x.shape[2], -1),
lambda x: (1, x, 1, 1),
),
}
# ATTN_FUNCTION = {
# "torch": F.scaled_dot_product_attention,
# "xformers": memory_efficient_attention,
# }
def vanilla_attention(q, k, v, mask, scale=None):
if scale is None:
scale = math.sqrt(q.size(-1))
scores = torch.bmm(q, k.transpose(-1, -2)) / scale
if mask is not None:
mask = einops.rearrange(mask, "b ... -> b (...)")
max_neg_value = -torch.finfo(scores.dtype).max
mask = einops.repeat(mask, "b j -> (b h) j", h=q.size(-3))
scores = scores.masked_fill(~mask, max_neg_value)
p_attn = F.softmax(scores, dim=-1)
return torch.bmm(p_attn, v)
def attention(q, k, v, head_dim, mask=None, scale=None, mode="xformers"):
"""
q, k, v: [B, L, D]
"""
pre_attn_layout = MEMORY_LAYOUTS[mode][0]
post_attn_layout = MEMORY_LAYOUTS[mode][1]
q = pre_attn_layout(q, head_dim)
k = pre_attn_layout(k, head_dim)
v = pre_attn_layout(v, head_dim)
# scores = ATTN_FUNCTION[mode](q, k.to(q), v.to(q), mask, scale=scale)
if mode == "torch":
assert scale is None
scores = F.scaled_dot_product_attention(q, k.to(q), v.to(q), mask) # , scale=scale)
elif mode == "xformers":
scores = memory_efficient_attention(q, k.to(q), v.to(q), mask, scale=scale)
else:
scores = vanilla_attention(q, k.to(q), v.to(q), mask, scale=scale)
scores = post_attn_layout(scores)
return scores
# DismantledBlock in mmdit.py
class SingleDiTBlock(nn.Module):
"""
A DiT block with gated adaptive layer norm (adaLN) conditioning.
"""
def __init__(
self,
hidden_size: int,
num_heads: int,
mlp_ratio: float = 4.0,
attn_mode: str = "xformers",
qkv_bias: bool = False,
pre_only: bool = False,
rmsnorm: bool = False,
scale_mod_only: bool = False,
swiglu: bool = False,
qk_norm: Optional[str] = None,
x_block_self_attn: bool = False,
**block_kwargs,
):
super().__init__()
assert attn_mode in MEMORY_LAYOUTS
self.attn_mode = attn_mode
if not rmsnorm:
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
else:
self.norm1 = RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.attn = AttentionLinears(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, pre_only=pre_only, qk_norm=qk_norm)
self.x_block_self_attn = x_block_self_attn
if self.x_block_self_attn:
assert not pre_only
assert not scale_mod_only
self.attn2 = AttentionLinears(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, pre_only=False, qk_norm=qk_norm)
if not pre_only:
if not rmsnorm:
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
else:
self.norm2 = RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
mlp_hidden_dim = int(hidden_size * mlp_ratio)
if not pre_only:
if not swiglu:
self.mlp = MLP(
in_features=hidden_size,
hidden_features=mlp_hidden_dim,
act_layer=lambda: nn.GELU(approximate="tanh"),
)
else:
self.mlp = SwiGLUFeedForward(
dim=hidden_size,
hidden_dim=mlp_hidden_dim,
multiple_of=256,
)
self.scale_mod_only = scale_mod_only
if self.x_block_self_attn:
n_mods = 9
elif not scale_mod_only:
n_mods = 6 if not pre_only else 2
else:
n_mods = 4 if not pre_only else 1
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, n_mods * hidden_size))
self.pre_only = pre_only
def pre_attention(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
if not self.pre_only:
if not self.scale_mod_only:
(shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = self.adaLN_modulation(c).chunk(6, dim=-1)
else:
shift_msa = None
shift_mlp = None
(scale_msa, gate_msa, scale_mlp, gate_mlp) = self.adaLN_modulation(c).chunk(4, dim=-1)
qkv = self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa))
return qkv, (x, gate_msa, shift_mlp, scale_mlp, gate_mlp)
else:
if not self.scale_mod_only:
(shift_msa, scale_msa) = self.adaLN_modulation(c).chunk(2, dim=-1)
else:
shift_msa = None
scale_msa = self.adaLN_modulation(c)
qkv = self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa))
return qkv, None
def pre_attention_x(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
assert self.x_block_self_attn
(shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2) = self.adaLN_modulation(
c
).chunk(9, dim=1)
x_norm = self.norm1(x)
qkv = self.attn.pre_attention(modulate(x_norm, shift_msa, scale_msa))
qkv2 = self.attn2.pre_attention(modulate(x_norm, shift_msa2, scale_msa2))
return qkv, qkv2, (x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2)
def post_attention(self, attn, x, gate_msa, shift_mlp, scale_mlp, gate_mlp):
assert not self.pre_only
x = x + gate_msa.unsqueeze(1) * self.attn.post_attention(attn)
x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
return x
def post_attention_x(self, attn, attn2, x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2, attn1_dropout: float = 0.0):
assert not self.pre_only
if attn1_dropout > 0.0:
# Use torch.bernoulli to implement dropout, only dropout the batch dimension
attn1_dropout = torch.bernoulli(torch.full((attn.size(0), 1, 1), 1 - attn1_dropout, device=attn.device))
attn_ = gate_msa.unsqueeze(1) * self.attn.post_attention(attn) * attn1_dropout
else:
attn_ = gate_msa.unsqueeze(1) * self.attn.post_attention(attn)
x = x + attn_
attn2_ = gate_msa2.unsqueeze(1) * self.attn2.post_attention(attn2)
x = x + attn2_
mlp_ = gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
x = x + mlp_
return x
# JointBlock + block_mixing in mmdit.py
class MMDiTBlock(nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
pre_only = kwargs.pop("pre_only")
x_block_self_attn = kwargs.pop("x_block_self_attn")
self.context_block = SingleDiTBlock(*args, pre_only=pre_only, **kwargs)
self.x_block = SingleDiTBlock(*args, pre_only=False, x_block_self_attn=x_block_self_attn, **kwargs)
self.head_dim = self.x_block.attn.head_dim
self.mode = self.x_block.attn_mode
self.gradient_checkpointing = False
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
def _forward(self, context, x, c):
ctx_qkv, ctx_intermediate = self.context_block.pre_attention(context, c)
if self.x_block.x_block_self_attn:
x_qkv, x_qkv2, x_intermediates = self.x_block.pre_attention_x(x, c)
else:
x_qkv, x_intermediates = self.x_block.pre_attention(x, c)
ctx_len = ctx_qkv[0].size(1)
q = torch.concat((ctx_qkv[0], x_qkv[0]), dim=1)
k = torch.concat((ctx_qkv[1], x_qkv[1]), dim=1)
v = torch.concat((ctx_qkv[2], x_qkv[2]), dim=1)
attn = attention(q, k, v, head_dim=self.head_dim, mode=self.mode)
ctx_attn_out = attn[:, :ctx_len]
x_attn_out = attn[:, ctx_len:]
if self.x_block.x_block_self_attn:
x_q2, x_k2, x_v2 = x_qkv2
attn2 = attention(x_q2, x_k2, x_v2, self.x_block.attn2.num_heads, mode=self.mode)
x = self.x_block.post_attention_x(x_attn_out, attn2, *x_intermediates)
else:
x = self.x_block.post_attention(x_attn_out, *x_intermediates)
if not self.context_block.pre_only:
context = self.context_block.post_attention(ctx_attn_out, *ctx_intermediate)
else:
context = None
return context, x
def forward(self, *args, **kwargs):
if self.training and self.gradient_checkpointing:
return checkpoint(self._forward, *args, use_reentrant=False, **kwargs)
else:
return self._forward(*args, **kwargs)
class MMDiT(nn.Module):
"""
Diffusion model with a Transformer backbone.
"""
# prepare pos_embed for latent size * 2
POS_EMBED_MAX_RATIO = 1.5
def __init__(
self,
input_size: int = 32,
patch_size: int = 2,
in_channels: int = 4,
depth: int = 28,
# hidden_size: Optional[int] = None,
# num_heads: Optional[int] = None,
mlp_ratio: float = 4.0,
learn_sigma: bool = False,
adm_in_channels: Optional[int] = None,
context_embedder_in_features: Optional[int] = None,
context_embedder_out_features: Optional[int] = None,
use_checkpoint: bool = False,
register_length: int = 0,
attn_mode: str = "torch",
rmsnorm: bool = False,
scale_mod_only: bool = False,
swiglu: bool = False,
out_channels: Optional[int] = None,
pos_embed_scaling_factor: Optional[float] = None,
pos_embed_offset: Optional[float] = None,
pos_embed_max_size: Optional[int] = None,
num_patches=None,
qk_norm: Optional[str] = None,
x_block_self_attn_layers: Optional[list[int]] = [],
qkv_bias: bool = True,
pos_emb_random_crop_rate: float = 0.0,
use_scaled_pos_embed: bool = False,
pos_embed_latent_sizes: Optional[list[int]] = None,
model_type: str = "sd3m",
):
super().__init__()
self._model_type = model_type
self.learn_sigma = learn_sigma
self.in_channels = in_channels
default_out_channels = in_channels * 2 if learn_sigma else in_channels
self.out_channels = default(out_channels, default_out_channels)
self.patch_size = patch_size
self.pos_embed_scaling_factor = pos_embed_scaling_factor
self.pos_embed_offset = pos_embed_offset
self.pos_embed_max_size = pos_embed_max_size
self.x_block_self_attn_layers = x_block_self_attn_layers
self.pos_emb_random_crop_rate = pos_emb_random_crop_rate
self.gradient_checkpointing = use_checkpoint
# hidden_size = default(hidden_size, 64 * depth)
# num_heads = default(num_heads, hidden_size // 64)
# apply magic --> this defines a head_size of 64
self.hidden_size = 64 * depth
num_heads = depth
self.num_heads = num_heads
self.enable_scaled_pos_embed(use_scaled_pos_embed, pos_embed_latent_sizes)
self.x_embedder = PatchEmbed(
input_size,
patch_size,
in_channels,
self.hidden_size,
bias=True,
strict_img_size=self.pos_embed_max_size is None,
)
self.t_embedder = TimestepEmbedding(self.hidden_size)
self.y_embedder = None
if adm_in_channels is not None:
assert isinstance(adm_in_channels, int)
self.y_embedder = Embedder(adm_in_channels, self.hidden_size)
if context_embedder_in_features is not None:
self.context_embedder = nn.Linear(context_embedder_in_features, context_embedder_out_features)
else:
self.context_embedder = nn.Identity()
self.register_length = register_length
if self.register_length > 0:
self.register = nn.Parameter(torch.randn(1, register_length, self.hidden_size))
# num_patches = self.x_embedder.num_patches
# Will use fixed sin-cos embedding:
# just use a buffer already
if num_patches is not None:
self.register_buffer(
"pos_embed",
torch.empty(1, num_patches, self.hidden_size),
)
else:
self.pos_embed = None
self.use_checkpoint = use_checkpoint
self.joint_blocks = nn.ModuleList(
[
MMDiTBlock(
self.hidden_size,
num_heads,
mlp_ratio=mlp_ratio,
attn_mode=attn_mode,
qkv_bias=qkv_bias,
pre_only=i == depth - 1,
rmsnorm=rmsnorm,
scale_mod_only=scale_mod_only,
swiglu=swiglu,
qk_norm=qk_norm,
x_block_self_attn=(i in self.x_block_self_attn_layers),
)
for i in range(depth)
]
)
for block in self.joint_blocks:
block.gradient_checkpointing = use_checkpoint
self.final_layer = UnPatch(self.hidden_size, patch_size, self.out_channels)
# self.initialize_weights()
self.blocks_to_swap = None
self.offloader = None
self.num_blocks = len(self.joint_blocks)
def enable_scaled_pos_embed(self, use_scaled_pos_embed: bool, latent_sizes: Optional[list[int]]):
self.use_scaled_pos_embed = use_scaled_pos_embed
if self.use_scaled_pos_embed:
# remove pos_embed to free up memory up to 0.4 GB
self.pos_embed = None
# remove duplicates and sort latent sizes in ascending order
latent_sizes = list(set(latent_sizes))
latent_sizes = sorted(latent_sizes)
patched_sizes = [latent_size // self.patch_size for latent_size in latent_sizes]
# calculate value range for each latent area: this is used to determine the pos_emb size from the latent shape
max_areas = []
for i in range(1, len(patched_sizes)):
prev_area = patched_sizes[i - 1] ** 2
area = patched_sizes[i] ** 2
max_areas.append((prev_area + area) // 2)
# area of the last latent size, if the latent size exceeds this, error will be raised
max_areas.append(int((patched_sizes[-1] * MMDiT.POS_EMBED_MAX_RATIO) ** 2))
# print("max_areas", max_areas)
self.resolution_area_to_latent_size = [(area, latent_size) for area, latent_size in zip(max_areas, patched_sizes)]
self.resolution_pos_embeds = {}
for patched_size in patched_sizes:
grid_size = int(patched_size * MMDiT.POS_EMBED_MAX_RATIO)
pos_embed = get_scaled_2d_sincos_pos_embed(self.hidden_size, grid_size, sample_size=patched_size)
pos_embed = torch.from_numpy(pos_embed).float().unsqueeze(0)
self.resolution_pos_embeds[patched_size] = pos_embed
# print(f"pos_embed for {patched_size}x{patched_size} latent size: {pos_embed.shape}")
else:
self.resolution_area_to_latent_size = None
self.resolution_pos_embeds = None
@property
def model_type(self):
return self._model_type
@property
def device(self):
return next(self.parameters()).device
@property
def dtype(self):
return next(self.parameters()).dtype
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
for block in self.joint_blocks:
block.enable_gradient_checkpointing()
def disable_gradient_checkpointing(self):
self.gradient_checkpointing = False
for block in self.joint_blocks:
block.disable_gradient_checkpointing()
def initialize_weights(self):
# TODO: Init context_embedder?
# Initialize transformer layers:
def _basic_init(module):
if isinstance(module, nn.Linear):
torch.nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.constant_(module.bias, 0)
self.apply(_basic_init)
# Initialize (and freeze) pos_embed by sin-cos embedding
if self.pos_embed is not None:
pos_embed = get_2d_sincos_pos_embed(
self.pos_embed.shape[-1],
int(self.pos_embed.shape[-2] ** 0.5),
scaling_factor=self.pos_embed_scaling_factor,
)
self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
# Initialize patch_embed like nn.Linear (instead of nn.Conv2d)
w = self.x_embedder.proj.weight.data
nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
nn.init.constant_(self.x_embedder.proj.bias, 0)
if getattr(self, "y_embedder", None) is not None:
nn.init.normal_(self.y_embedder.mlp[0].weight, std=0.02)
nn.init.normal_(self.y_embedder.mlp[2].weight, std=0.02)
# Initialize timestep embedding MLP:
nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
# Zero-out adaLN modulation layers in DiT blocks:
for block in self.joint_blocks:
nn.init.constant_(block.x_block.adaLN_modulation[-1].weight, 0)
nn.init.constant_(block.x_block.adaLN_modulation[-1].bias, 0)
nn.init.constant_(block.context_block.adaLN_modulation[-1].weight, 0)
nn.init.constant_(block.context_block.adaLN_modulation[-1].bias, 0)
# Zero-out output layers:
nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
nn.init.constant_(self.final_layer.linear.weight, 0)
nn.init.constant_(self.final_layer.linear.bias, 0)
def set_pos_emb_random_crop_rate(self, rate: float):
self.pos_emb_random_crop_rate = rate
def cropped_pos_embed(self, h, w, device=None, random_crop: bool = False):
p = self.x_embedder.patch_size
# patched size
h = (h + 1) // p
w = (w + 1) // p
if self.pos_embed is None: # should not happen
return get_2d_sincos_pos_embed_torch(self.hidden_size, w, h, device=device)
assert self.pos_embed_max_size is not None
assert h <= self.pos_embed_max_size, (h, self.pos_embed_max_size)
assert w <= self.pos_embed_max_size, (w, self.pos_embed_max_size)
if not random_crop:
top = (self.pos_embed_max_size - h) // 2
left = (self.pos_embed_max_size - w) // 2
else:
top = torch.randint(0, self.pos_embed_max_size - h + 1, (1,)).item()
left = torch.randint(0, self.pos_embed_max_size - w + 1, (1,)).item()
spatial_pos_embed = self.pos_embed.reshape(
1,
self.pos_embed_max_size,
self.pos_embed_max_size,
self.pos_embed.shape[-1],
)
spatial_pos_embed = spatial_pos_embed[:, top : top + h, left : left + w, :]
spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])
return spatial_pos_embed
def cropped_scaled_pos_embed(self, h, w, device=None, dtype=None, random_crop: bool = False):
p = self.x_embedder.patch_size
# patched size
h = (h + 1) // p
w = (w + 1) // p
# select pos_embed size based on area
area = h * w
patched_size = None
for area_, patched_size_ in self.resolution_area_to_latent_size:
if area <= area_:
patched_size = patched_size_
break
if patched_size is None:
raise ValueError(f"Area {area} is too large for the given latent sizes {self.resolution_area_to_latent_size}.")
pos_embed = self.resolution_pos_embeds[patched_size]
pos_embed_size = round(math.sqrt(pos_embed.shape[1]))
if h > pos_embed_size or w > pos_embed_size:
# # fallback to normal pos_embed
# return self.cropped_pos_embed(h * p, w * p, device=device, random_crop=random_crop)
# extend pos_embed size
logger.warning(
f"Using normal pos_embed for size {h}x{w} as it exceeds the scaled pos_embed size {pos_embed_size}. Image is too tall or wide."
)
pos_embed_size = max(h, w)
pos_embed = get_scaled_2d_sincos_pos_embed(self.hidden_size, pos_embed_size, sample_size=patched_size)
pos_embed = torch.from_numpy(pos_embed).float().unsqueeze(0)
self.resolution_pos_embeds[patched_size] = pos_embed
logger.info(f"Updated pos_embed for size {pos_embed_size}x{pos_embed_size}")
if not random_crop:
top = (pos_embed_size - h) // 2
left = (pos_embed_size - w) // 2
else:
top = torch.randint(0, pos_embed_size - h + 1, (1,)).item()
left = torch.randint(0, pos_embed_size - w + 1, (1,)).item()
if pos_embed.device != device:
pos_embed = pos_embed.to(device)
# which is better to update device, or transfer every time to device? -> 64x64 emb is 96*96*1536*4=56MB. It's okay to update device.
self.resolution_pos_embeds[patched_size] = pos_embed # update device
if pos_embed.dtype != dtype:
pos_embed = pos_embed.to(dtype)
self.resolution_pos_embeds[patched_size] = pos_embed # update dtype
spatial_pos_embed = pos_embed.reshape(1, pos_embed_size, pos_embed_size, pos_embed.shape[-1])
spatial_pos_embed = spatial_pos_embed[:, top : top + h, left : left + w, :]
spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])
# print(
# f"patched size: {h}x{w}, pos_embed size: {pos_embed_size}, pos_embed shape: {pos_embed.shape}, top: {top}, left: {left}"
# )
return spatial_pos_embed
def enable_block_swap(self, num_blocks: int, device: torch.device):
self.blocks_to_swap = num_blocks
assert (
self.blocks_to_swap <= self.num_blocks - 2
), f"Cannot swap more than {self.num_blocks - 2} blocks. Requested: {self.blocks_to_swap} blocks."
self.offloader = custom_offloading_utils.ModelOffloader(
self.joint_blocks, self.num_blocks, self.blocks_to_swap, device # , debug=True
)
print(f"SD3: Block swap enabled. Swapping {num_blocks} blocks, total blocks: {self.num_blocks}, device: {device}.")
def move_to_device_except_swap_blocks(self, device: torch.device):
# assume model is on cpu. do not move blocks to device to reduce temporary memory usage
if self.blocks_to_swap:
save_blocks = self.joint_blocks
self.joint_blocks = None
self.to(device)
if self.blocks_to_swap:
self.joint_blocks = save_blocks
def prepare_block_swap_before_forward(self):
if self.blocks_to_swap is None or self.blocks_to_swap == 0:
return
self.offloader.prepare_block_devices_before_forward(self.joint_blocks)
def forward(
self,
x: torch.Tensor,
t: torch.Tensor,
y: Optional[torch.Tensor] = None,
context: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Forward pass of DiT.
x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
t: (N,) tensor of diffusion timesteps
y: (N, D) tensor of class labels
"""
pos_emb_random_crop = (
False if self.pos_emb_random_crop_rate == 0.0 else torch.rand(1).item() < self.pos_emb_random_crop_rate
)
B, C, H, W = x.shape
# x = self.x_embedder(x) + self.cropped_pos_embed(H, W, device=x.device, random_crop=pos_emb_random_crop).to(dtype=x.dtype)
if not self.use_scaled_pos_embed:
pos_embed = self.cropped_pos_embed(H, W, device=x.device, random_crop=pos_emb_random_crop).to(dtype=x.dtype)
else:
# print(f"Using scaled pos_embed for size {H}x{W}")
pos_embed = self.cropped_scaled_pos_embed(H, W, device=x.device, dtype=x.dtype, random_crop=pos_emb_random_crop)
x = self.x_embedder(x) + pos_embed
del pos_embed
c = self.t_embedder(t, dtype=x.dtype) # (N, D)
if y is not None and self.y_embedder is not None:
y = self.y_embedder(y) # (N, D)
c = c + y # (N, D)
if context is not None:
context = self.context_embedder(context)
if self.register_length > 0:
context = torch.cat(
(einops.repeat(self.register, "1 ... -> b ...", b=x.shape[0]), default(context, torch.Tensor([]).type_as(x))), 1
)
if not self.blocks_to_swap:
for block in self.joint_blocks:
context, x = block(context, x, c)
else:
for block_idx, block in enumerate(self.joint_blocks):
self.offloader.wait_for_block(block_idx)
context, x = block(context, x, c)
self.offloader.submit_move_blocks(self.joint_blocks, block_idx)
x = self.final_layer(x, c, H, W) # Our final layer combined UnPatchify
return x[:, :, :H, :W]
def create_sd3_mmdit(params: SD3Params, attn_mode: str = "torch") -> MMDiT:
mmdit = MMDiT(
input_size=None,
pos_embed_max_size=params.pos_embed_max_size,
patch_size=params.patch_size,
in_channels=16,
adm_in_channels=params.adm_in_channels,
context_embedder_in_features=params.context_embedder_in_features,
context_embedder_out_features=params.context_embedder_out_features,
depth=params.depth,
mlp_ratio=4,
qk_norm=params.qk_norm,
x_block_self_attn_layers=params.x_block_self_attn_layers,
num_patches=params.num_patches,
attn_mode=attn_mode,
model_type=params.model_type,
)
return mmdit
# endregion
# region VAE
VAE_SCALE_FACTOR = 1.5305
VAE_SHIFT_FACTOR = 0.0609
def Normalize(in_channels, num_groups=32, dtype=torch.float32, device=None):
return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
class ResnetBlock(torch.nn.Module):
def __init__(self, *, in_channels, out_channels=None, dtype=torch.float32, device=None):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.norm1 = Normalize(in_channels, dtype=dtype, device=device)
self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
self.norm2 = Normalize(out_channels, dtype=dtype, device=device)
self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
if self.in_channels != self.out_channels:
self.nin_shortcut = torch.nn.Conv2d(
in_channels, out_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device
)
else:
self.nin_shortcut = None
self.swish = torch.nn.SiLU(inplace=True)
def forward(self, x):
hidden = x
hidden = self.norm1(hidden)
hidden = self.swish(hidden)
hidden = self.conv1(hidden)
hidden = self.norm2(hidden)
hidden = self.swish(hidden)
hidden = self.conv2(hidden)
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return x + hidden
class AttnBlock(torch.nn.Module):
def __init__(self, in_channels, dtype=torch.float32, device=None):
super().__init__()
self.norm = Normalize(in_channels, dtype=dtype, device=device)
self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
def forward(self, x):
hidden = self.norm(x)
q = self.q(hidden)
k = self.k(hidden)
v = self.v(hidden)
b, c, h, w = q.shape
q, k, v = map(lambda x: einops.rearrange(x, "b c h w -> b 1 (h w) c").contiguous(), (q, k, v))
hidden = torch.nn.functional.scaled_dot_product_attention(q, k, v) # scale is dim ** -0.5 per default
hidden = einops.rearrange(hidden, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
hidden = self.proj_out(hidden)
return x + hidden
class Downsample(torch.nn.Module):
def __init__(self, in_channels, dtype=torch.float32, device=None):
super().__init__()
self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0, dtype=dtype, device=device)
def forward(self, x):
pad = (0, 1, 0, 1)
x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
x = self.conv(x)
return x
class Upsample(torch.nn.Module):
def __init__(self, in_channels, dtype=torch.float32, device=None):
super().__init__()
self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
def forward(self, x):
org_dtype = x.dtype
if x.dtype == torch.bfloat16:
x = x.to(torch.float32)
x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
if x.dtype != org_dtype:
x = x.to(org_dtype)
x = self.conv(x)
return x
class VAEEncoder(torch.nn.Module):
def __init__(
self, ch=128, ch_mult=(1, 2, 4, 4), num_res_blocks=2, in_channels=3, z_channels=16, dtype=torch.float32, device=None
):
super().__init__()
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
# downsampling
self.conv_in = torch.nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
in_ch_mult = (1,) + tuple(ch_mult)
self.in_ch_mult = in_ch_mult
self.down = torch.nn.ModuleList()
for i_level in range(self.num_resolutions):
block = torch.nn.ModuleList()
attn = torch.nn.ModuleList()
block_in = ch * in_ch_mult[i_level]
block_out = ch * ch_mult[i_level]
for i_block in range(num_res_blocks):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dtype=dtype, device=device))
block_in = block_out
down = torch.nn.Module()
down.block = block
down.attn = attn
if i_level != self.num_resolutions - 1:
down.downsample = Downsample(block_in, dtype=dtype, device=device)
self.down.append(down)
# middle
self.mid = torch.nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dtype=dtype, device=device)
self.mid.attn_1 = AttnBlock(block_in, dtype=dtype, device=device)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dtype=dtype, device=device)
# end
self.norm_out = Normalize(block_in, dtype=dtype, device=device)
self.conv_out = torch.nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
self.swish = torch.nn.SiLU(inplace=True)
def forward(self, x):
# downsampling
hs = [self.conv_in(x)]
for i_level in range(self.num_resolutions):
for i_block in range(self.num_res_blocks):
h = self.down[i_level].block[i_block](hs[-1])
hs.append(h)
if i_level != self.num_resolutions - 1:
hs.append(self.down[i_level].downsample(hs[-1]))
# middle
h = hs[-1]
h = self.mid.block_1(h)
h = self.mid.attn_1(h)
h = self.mid.block_2(h)
# end
h = self.norm_out(h)
h = self.swish(h)
h = self.conv_out(h)
return h
class VAEDecoder(torch.nn.Module):
def __init__(
self,
ch=128,
out_ch=3,
ch_mult=(1, 2, 4, 4),
num_res_blocks=2,
resolution=256,
z_channels=16,
dtype=torch.float32,
device=None,
):
super().__init__()
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
block_in = ch * ch_mult[self.num_resolutions - 1]
curr_res = resolution // 2 ** (self.num_resolutions - 1)
# z to block_in
self.conv_in = torch.nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
# middle
self.mid = torch.nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dtype=dtype, device=device)
self.mid.attn_1 = AttnBlock(block_in, dtype=dtype, device=device)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dtype=dtype, device=device)
# upsampling
self.up = torch.nn.ModuleList()
for i_level in reversed(range(self.num_resolutions)):
block = torch.nn.ModuleList()
block_out = ch * ch_mult[i_level]
for i_block in range(self.num_res_blocks + 1):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dtype=dtype, device=device))
block_in = block_out
up = torch.nn.Module()
up.block = block
if i_level != 0:
up.upsample = Upsample(block_in, dtype=dtype, device=device)
curr_res = curr_res * 2
self.up.insert(0, up) # prepend to get consistent order
# end
self.norm_out = Normalize(block_in, dtype=dtype, device=device)
self.conv_out = torch.nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
self.swish = torch.nn.SiLU(inplace=True)
def forward(self, z):
# z to block_in
hidden = self.conv_in(z)
# middle
hidden = self.mid.block_1(hidden)
hidden = self.mid.attn_1(hidden)
hidden = self.mid.block_2(hidden)
# upsampling
for i_level in reversed(range(self.num_resolutions)):
for i_block in range(self.num_res_blocks + 1):
hidden = self.up[i_level].block[i_block](hidden)
if i_level != 0:
hidden = self.up[i_level].upsample(hidden)
# end
hidden = self.norm_out(hidden)
hidden = self.swish(hidden)
hidden = self.conv_out(hidden)
return hidden
class SDVAE(torch.nn.Module):
def __init__(self, dtype=torch.float32, device=None):
super().__init__()
self.encoder = VAEEncoder(dtype=dtype, device=device)
self.decoder = VAEDecoder(dtype=dtype, device=device)
@property
def device(self):
return next(self.parameters()).device
@property
def dtype(self):
return next(self.parameters()).dtype
# @torch.autocast("cuda", dtype=torch.float16)
def decode(self, latent):
return self.decoder(latent)
# @torch.autocast("cuda", dtype=torch.float16)
def encode(self, image):
hidden = self.encoder(image)
mean, logvar = torch.chunk(hidden, 2, dim=1)
logvar = torch.clamp(logvar, -30.0, 20.0)
std = torch.exp(0.5 * logvar)
return mean + std * torch.randn_like(mean)
@staticmethod
def process_in(latent):
return (latent - VAE_SHIFT_FACTOR) * VAE_SCALE_FACTOR
@staticmethod
def process_out(latent):
return (latent / VAE_SCALE_FACTOR) + VAE_SHIFT_FACTOR
# endregion