File size: 9,756 Bytes
eb339cb |
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 |
from typing import Optional, Union
import torch
import torch.nn as nn
from mld.models.operator.embeddings import TimestepEmbedding, Timesteps
from mld.models.operator.attention import (SkipTransformerEncoder,
SkipTransformerDecoder,
TransformerDecoder,
TransformerDecoderLayer,
TransformerEncoder,
TransformerEncoderLayer)
from mld.models.operator.moe import MoeTransformerEncoderLayer, MoeTransformerDecoderLayer
from mld.models.operator.utils import get_clones, get_activation_fn, zero_module
from mld.models.operator.position_encoding import build_position_encoding
def load_balancing_loss_func(router_logits: tuple, num_experts: int = 4, topk: int = 2):
router_logits = torch.cat(router_logits, dim=0)
routing_weights = torch.nn.functional.softmax(router_logits, dim=-1)
_, selected_experts = torch.topk(routing_weights, topk, dim=-1)
expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
router_prob_per_expert = torch.mean(routing_weights, dim=0)
overall_loss = num_experts * torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
return overall_loss
class MldDenoiser(nn.Module):
def __init__(self,
latent_dim: list = [1, 256],
hidden_dim: Optional[int] = None,
text_dim: int = 768,
time_dim: int = 768,
ff_size: int = 1024,
num_layers: int = 9,
num_heads: int = 4,
dropout: float = 0.1,
normalize_before: bool = False,
norm_eps: float = 1e-5,
activation: str = "gelu",
norm_post: bool = True,
activation_post: Optional[str] = None,
flip_sin_to_cos: bool = True,
freq_shift: float = 0,
time_act_fn: str = 'silu',
time_post_act_fn: Optional[str] = None,
position_embedding: str = "learned",
arch: str = "trans_enc",
add_mem_pos: bool = True,
force_pre_post_proj: bool = False,
text_act_fn: str = 'relu',
time_cond_proj_dim: Optional[int] = None,
zero_init_cond: bool = True,
is_controlnet: bool = False,
controlnet_embed_dim: Optional[int] = None,
controlnet_act_fn: str = 'silu',
moe: bool = False,
moe_num_experts: int = 4,
moe_topk: int = 2,
moe_loss_weight: float = 1e-2,
moe_jitter_noise: Optional[float] = None
) -> None:
super(MldDenoiser, self).__init__()
self.latent_dim = latent_dim[-1] if hidden_dim is None else hidden_dim
add_pre_post_proj = force_pre_post_proj or (hidden_dim is not None and hidden_dim != latent_dim[-1])
self.latent_pre = nn.Linear(latent_dim[-1], self.latent_dim) if add_pre_post_proj else nn.Identity()
self.latent_post = nn.Linear(self.latent_dim, latent_dim[-1]) if add_pre_post_proj else nn.Identity()
self.arch = arch
self.time_cond_proj_dim = time_cond_proj_dim
self.moe_num_experts = moe_num_experts
self.moe_topk = moe_topk
self.moe_loss_weight = moe_loss_weight
self.time_proj = Timesteps(time_dim, flip_sin_to_cos, freq_shift)
self.time_embedding = TimestepEmbedding(time_dim, self.latent_dim, time_act_fn, post_act_fn=time_post_act_fn,
cond_proj_dim=time_cond_proj_dim, zero_init_cond=zero_init_cond)
self.emb_proj = nn.Sequential(get_activation_fn(text_act_fn), nn.Linear(text_dim, self.latent_dim))
self.query_pos = build_position_encoding(self.latent_dim, position_embedding=position_embedding)
if self.arch == "trans_enc":
if moe:
encoder_layer = MoeTransformerEncoderLayer(
self.latent_dim, num_heads, moe_num_experts, moe_topk, ff_size,
dropout, activation, normalize_before, norm_eps, moe_jitter_noise)
else:
encoder_layer = TransformerEncoderLayer(
self.latent_dim, num_heads, ff_size, dropout,
activation, normalize_before, norm_eps)
encoder_norm = nn.LayerNorm(self.latent_dim, eps=norm_eps) if norm_post and not is_controlnet else None
self.encoder = SkipTransformerEncoder(encoder_layer, num_layers, encoder_norm, activation_post,
is_controlnet=is_controlnet, is_moe=moe)
elif self.arch == 'trans_dec':
if add_mem_pos:
self.mem_pos = build_position_encoding(self.latent_dim, position_embedding=position_embedding)
else:
self.mem_pos = None
if moe:
decoder_layer = MoeTransformerDecoderLayer(
self.latent_dim, num_heads, moe_num_experts, moe_topk, ff_size,
dropout, activation, normalize_before, norm_eps, moe_jitter_noise)
else:
decoder_layer = TransformerDecoderLayer(
self.latent_dim, num_heads, ff_size, dropout,
activation, normalize_before, norm_eps)
decoder_norm = nn.LayerNorm(self.latent_dim, eps=norm_eps) if norm_post and not is_controlnet else None
self.decoder = SkipTransformerDecoder(decoder_layer, num_layers, decoder_norm, activation_post,
is_controlnet=is_controlnet, is_moe=moe)
else:
raise ValueError(f"Not supported architecture: {self.arch}!")
self.is_controlnet = is_controlnet
if self.is_controlnet:
embed_dim = controlnet_embed_dim if controlnet_embed_dim is not None else self.latent_dim
modules = [
nn.Linear(latent_dim[-1], embed_dim),
get_activation_fn(controlnet_act_fn) if controlnet_act_fn else None,
nn.Linear(embed_dim, embed_dim),
get_activation_fn(controlnet_act_fn) if controlnet_act_fn else None,
zero_module(nn.Linear(embed_dim, latent_dim[-1]))
]
self.controlnet_cond_embedding = nn.Sequential(*[m for m in modules if m is not None])
self.controlnet_down_mid_blocks = nn.ModuleList([
zero_module(nn.Linear(self.latent_dim, self.latent_dim)) for _ in range(num_layers)])
def forward(self,
sample: torch.Tensor,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep_cond: Optional[torch.Tensor] = None,
controlnet_cond: Optional[torch.Tensor] = None,
controlnet_residuals: Optional[list[torch.Tensor]] = None
) -> tuple:
# 0. check if controlnet
if self.is_controlnet:
sample = sample + self.controlnet_cond_embedding(controlnet_cond)
# 1. dimension matching (pre)
sample = sample.permute(1, 0, 2)
sample = self.latent_pre(sample)
# 2. time_embedding
timesteps = timestep.expand(sample.shape[1]).clone()
time_emb = self.time_proj(timesteps)
time_emb = time_emb.to(dtype=sample.dtype)
# [1, bs, latent_dim] <= [bs, latent_dim]
time_emb = self.time_embedding(time_emb, timestep_cond).unsqueeze(0)
# 3. condition + time embedding
# text_emb [seq_len, batch_size, text_dim] <= [batch_size, seq_len, text_dim]
encoder_hidden_states = encoder_hidden_states.permute(1, 0, 2)
# text embedding projection
text_emb_latent = self.emb_proj(encoder_hidden_states)
emb_latent = torch.cat((time_emb, text_emb_latent), 0)
# 4. transformer
if self.arch == "trans_enc":
xseq = torch.cat((sample, emb_latent), axis=0)
xseq = self.query_pos(xseq)
tokens, intermediates, router_logits = self.encoder(xseq, controlnet_residuals=controlnet_residuals)
elif self.arch == 'trans_dec':
sample = self.query_pos(sample)
if self.mem_pos:
emb_latent = self.mem_pos(emb_latent)
tokens, intermediates, router_logits = self.decoder(sample, emb_latent,
controlnet_residuals=controlnet_residuals)
else:
raise TypeError(f"{self.arch} is not supported")
router_loss = None
if router_logits is not None:
router_loss = load_balancing_loss_func(router_logits, self.moe_num_experts, self.moe_topk)
router_loss = self.moe_loss_weight * router_loss
if self.is_controlnet:
control_res_samples = []
for res, block in zip(intermediates, self.controlnet_down_mid_blocks):
r = block(res)
control_res_samples.append(r)
return control_res_samples, router_loss
elif self.arch == "trans_enc":
sample = tokens[:sample.shape[0]]
elif self.arch == 'trans_dec':
sample = tokens
else:
raise TypeError(f"{self.arch} is not supported")
# 5. dimension matching (post)
sample = self.latent_post(sample)
sample = sample.permute(1, 0, 2)
return sample, router_loss
|