| import torch |
| from torch import nn |
| from functools import partial |
| from utils.interpolation_modules.attention import SelfAttention, CrossAttention, TemporalAttention |
|
|
|
|
| def modulate(x, shift, scale): |
| return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
|
|
|
|
| class Mlp(nn.Module): |
| def __init__(self, dim=768, mlp_ratio=4., act_layer=nn.GELU, proj_drop_rate=0.): |
| super().__init__() |
| hidden_dim = int(dim * mlp_ratio) |
| self.fc1 = nn.Linear(dim, hidden_dim) |
| self.act = act_layer() |
| self.fc2 = nn.Linear(hidden_dim, dim) |
| self.drop = nn.Dropout(proj_drop_rate) |
|
|
| def forward(self, x): |
| x = self.drop(self.act(self.fc1(x))) |
| x = self.drop(self.fc2(x)) |
| return x |
|
|
|
|
| class VAEBlock(nn.Module): |
| def __init__(self, dim=768, num_heads=12, mlp_ratio=4., qkv_bias=False, attn_drop_rate=0., proj_drop_rate=0., act_layer=nn.GELU, |
| norm_layer=partial(nn.LayerNorm, eps=1e-5), use_xformers=True, is_encoder=True, add_attn_encoder=False, |
| add_attn_decoder=False, add_attn_type="temporal_attn"): |
| super().__init__() |
| self.is_encoder = is_encoder |
| self.add_attn_encoder = add_attn_encoder |
| self.add_attn_decoder = add_attn_decoder |
| self.add_attn_type = add_attn_type |
| self.norm_self_attn = norm_layer(dim) |
| self.self_attn = SelfAttention(dim, num_heads, qkv_bias, attn_drop_rate, proj_drop_rate, use_xformers) |
| if (add_attn_encoder and is_encoder) or (add_attn_decoder and not is_encoder): |
| if add_attn_type == "cross_attn": |
| self.norm_cross_attn = norm_layer(dim) |
| self.cross_attn = CrossAttention(dim, num_heads, qkv_bias, attn_drop_rate, proj_drop_rate, use_xformers) |
| elif add_attn_type == "temporal_attn": |
| self.norm_temp_attn = norm_layer(dim) |
| self.temp_attn = TemporalAttention(dim, num_heads, qkv_bias, attn_drop_rate, proj_drop_rate, use_xformers) |
| else: |
| raise f"No implementation of {add_attn_type}" |
| self.norm_mlp = norm_layer(dim) |
| self.mlp = Mlp(dim, mlp_ratio, act_layer, proj_drop_rate) |
|
|
| def forward(self, query_tokens, mid_tokens, cond_tokens, ph=None, pw=None): |
| b, n, d = mid_tokens.shape |
| x = torch.cat((mid_tokens, query_tokens), dim=1) |
| x = x + self.self_attn(self.norm_self_attn(x)) |
| if self.is_encoder: |
| y = x[:, n:, :] |
| else: |
| y = x[:, :n, :] |
| if (self.add_attn_encoder and self.is_encoder) or (self.add_attn_decoder and not self.is_encoder): |
| if self.add_attn_type == "cross_attn": |
| y = y + self.cross_attn(self.norm_cross_attn(y), y=cond_tokens) |
| elif self.add_attn_type == "temporal_attn": |
| tokens0, tokens1 = cond_tokens.chunk(2, dim=1) |
| y = y + self.temp_attn(self.norm_temp_attn(y), x0=tokens0, x1=tokens1, ph=ph, pw=pw) |
| y = y + self.mlp(self.norm_mlp(y)) |
| return y |
|
|
|
|
| class DiTBlock(nn.Module): |
| def __init__(self, dim=768, num_heads=12, mlp_ratio=4., qkv_bias=False, attn_drop_rate=0., proj_drop_rate=0., act_layer=nn.GELU, |
| norm_layer=partial(nn.LayerNorm, eps=1e-5), use_xformers=True): |
| super().__init__() |
| self.norm_self_attn = norm_layer(dim) |
| self.self_attn = SelfAttention(dim, num_heads, qkv_bias, attn_drop_rate, proj_drop_rate, use_xformers) |
| self.norm_temp_attn = norm_layer(dim) |
| self.temp_attn = TemporalAttention(dim, num_heads, qkv_bias, attn_drop_rate, proj_drop_rate, use_xformers) |
| self.norm_mlp = norm_layer(dim) |
| self.mlp = Mlp(dim, mlp_ratio, act_layer, proj_drop_rate) |
|
|
| def forward(self, query_tokens, tokens0, tokens1, ph, pw, modulations): |
| shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulations.chunk(6, dim=1) |
| query_tokens = query_tokens + gate_msa.unsqueeze(1) * self.self_attn(modulate(self.norm_self_attn(query_tokens), |
| shift_msa, scale_msa)) |
| query_tokens = query_tokens + self.temp_attn(self.norm_temp_attn(query_tokens), x0=tokens0, x1=tokens1, |
| ph=ph, pw=pw) |
| query_tokens = query_tokens + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm_mlp(query_tokens), |
| shift_mlp, scale_mlp)) |
| return query_tokens |
|
|
