147 lines
6.0 KiB
Python
147 lines
6.0 KiB
Python
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
|
|
|
# Licensed under the GNU General Public License, Version 3.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
|
|
# https://www.gnu.org/licenses/gpl-3.0
|
|
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers.models.attention import FeedForward
|
|
from diffusers.models.attention_processor import Attention
|
|
from typing import Tuple
|
|
|
|
|
|
class PerceiverAttentionBlock(nn.Module):
|
|
def __init__(self, dim, dim_head=64, heads=8, ff_mult=4, n_conds=1):
|
|
super().__init__()
|
|
|
|
self.norm_attn1 = nn.LayerNorm(dim)
|
|
self.norm_attn2_l = nn.LayerNorm(dim)
|
|
self.norm_attn2_x = nn.LayerNorm(dim)
|
|
if n_conds == 2:
|
|
self.norm_attn3_l = nn.LayerNorm(dim)
|
|
self.norm_attn3_y = nn.LayerNorm(dim)
|
|
self.norm_ff = nn.LayerNorm(dim)
|
|
|
|
self.attn1 = Attention(query_dim=dim, cross_attention_dim=dim, heads=heads, dim_head=dim_head, out_bias=False)
|
|
self.attn2 = Attention(query_dim=dim, cross_attention_dim=dim, heads=heads, dim_head=dim_head, out_bias=False)
|
|
if n_conds == 2:
|
|
self.attn3 = Attention(query_dim=dim, cross_attention_dim=dim, heads=heads, dim_head=dim_head, out_bias=False)
|
|
self.ff = FeedForward(dim=dim, mult=ff_mult, activation_fn='gelu')
|
|
|
|
def forward(self, x, latents, y=None):
|
|
# attention 1
|
|
norm_latents = self.norm_attn1(latents)
|
|
latents = self.attn1(norm_latents, encoder_hidden_states=norm_latents) + latents
|
|
|
|
# attention 2
|
|
norm_latents = self.norm_attn2_l(latents)
|
|
norm_x = self.norm_attn2_x(x)
|
|
latents = self.attn2(norm_latents, encoder_hidden_states=norm_x) + latents
|
|
|
|
# attention 3
|
|
if not y is None:
|
|
norm_latents = self.norm_attn3_l(latents)
|
|
norm_y = self.norm_attn3_y(y)
|
|
latents = self.attn3(norm_latents, encoder_hidden_states=norm_y) + latents
|
|
|
|
# feed forward
|
|
norm_latents = self.norm_ff(latents)
|
|
latents = self.ff(norm_latents) + latents
|
|
|
|
return latents
|
|
|
|
|
|
class Resampler(nn.Module):
|
|
def __init__(self, dim=1024, depth=8, dim_head=64, heads=8, num_queries=16, embedding_dim=768, output_dim=1024, ff_mult=4, n_conds=1):
|
|
super().__init__()
|
|
|
|
self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim ** 0.5)
|
|
self.proj_in = nn.Linear(embedding_dim, dim)
|
|
self.proj_out = nn.Linear(dim, output_dim)
|
|
|
|
self.layers = nn.ModuleList([
|
|
PerceiverAttentionBlock(dim=dim, dim_head=dim_head, heads=heads, ff_mult=ff_mult, n_conds=n_conds)
|
|
for _ in range(depth)
|
|
])
|
|
|
|
if n_conds == 2:
|
|
self.proj_in_2 = nn.Linear(embedding_dim, dim)
|
|
|
|
def norm_out(self, x):
|
|
return x * (x.pow(2.0).sum(dim=-1, keepdim=True) + 1e-8).rsqrt()
|
|
|
|
def forward(self, x, y=None):
|
|
latents = self.latents.repeat(x.shape[0], 1, 1)
|
|
|
|
x = self.proj_in(x)
|
|
if not y is None:
|
|
y = self.proj_in_2(y)
|
|
for layer in self.layers:
|
|
latents = layer(x, latents, y)
|
|
|
|
latents = self.proj_out(latents)
|
|
latents = torch.mean(latents, dim=0, keepdim=True) # get mean latents across batch
|
|
latents = self.norm_out(latents)
|
|
|
|
return latents
|
|
|
|
|
|
class HyperLoRALayer(nn.Module):
|
|
def __init__(self, dim: int, in_features: int, out_features: int, lora_rank: int = 4, has_base_lora: bool = False):
|
|
super().__init__()
|
|
|
|
self.id_down_basis = nn.Parameter(torch.zeros(dim, in_features, lora_rank))
|
|
self.id_up_basis = nn.Parameter(torch.zeros(dim, lora_rank, out_features))
|
|
|
|
if has_base_lora:
|
|
self.base_down_basis = nn.Parameter(torch.zeros(dim, in_features, lora_rank // 2))
|
|
self.base_up_basis = nn.Parameter(torch.zeros(dim, lora_rank // 2, out_features))
|
|
|
|
def linear_mix(self, coeffs: torch.Tensor, basis: torch.Tensor):
|
|
return torch.einsum('bi, ijk -> bjk', coeffs, basis)
|
|
|
|
def forward(self, tokens: torch.Tensor, mode='id'):
|
|
down_coeffs, up_coeffs = tokens[:,0,:], tokens[:,1,:]
|
|
if mode == 'id':
|
|
down, up = self.linear_mix(down_coeffs, self.id_down_basis), self.linear_mix(up_coeffs, self.id_up_basis)
|
|
elif mode == 'base':
|
|
down, up = self.linear_mix(down_coeffs, self.base_down_basis), self.linear_mix(up_coeffs, self.base_up_basis)
|
|
return down[0].t().contiguous(), up[0].t().contiguous()
|
|
|
|
|
|
class HyperLoRAModule(nn.Module):
|
|
def __init__(self, dim: int, hidden_size: int, cross_attention_dim: int, lora_rank: int = 4, has_base_lora: bool = False):
|
|
super().__init__()
|
|
|
|
self.to_q = HyperLoRALayer(dim, hidden_size, hidden_size, lora_rank=lora_rank, has_base_lora=has_base_lora)
|
|
self.to_k = HyperLoRALayer(dim, cross_attention_dim or hidden_size, hidden_size, lora_rank=lora_rank, has_base_lora=has_base_lora)
|
|
self.to_v = HyperLoRALayer(dim, cross_attention_dim or hidden_size, hidden_size, lora_rank=lora_rank, has_base_lora=has_base_lora)
|
|
self.to_out = HyperLoRALayer(dim, hidden_size, hidden_size, lora_rank=lora_rank, has_base_lora=has_base_lora)
|
|
|
|
def forward(self, tokens: torch.Tensor, mode='id'):
|
|
q, k, v, out = tokens.chunk(4, dim=1)
|
|
return {
|
|
'to_q': self.to_q(q, mode=mode),
|
|
'to_k': self.to_k(k, mode=mode),
|
|
'to_v': self.to_v(v, mode=mode),
|
|
'to_out_0': self.to_out(out, mode=mode)
|
|
}
|
|
|
|
|
|
class Reshape(nn.Module):
|
|
def __init__(self, shape: Tuple[int]):
|
|
super().__init__()
|
|
self.shape = shape
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
return x.reshape(*self.shape)
|