235 lines
8.9 KiB
Python
235 lines
8.9 KiB
Python
# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.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
|
|
from torch import nn
|
|
from einops import rearrange
|
|
from torch.nn import functional as F
|
|
from .attention import Attention
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from diffusers.models.attention import FeedForward
|
|
from einops import rearrange
|
|
|
|
|
|
class StableSyncNet(nn.Module):
|
|
def __init__(self, config, gradient_checkpointing=False):
|
|
super().__init__()
|
|
self.audio_encoder = DownEncoder2D(
|
|
in_channels=config["audio_encoder"]["in_channels"],
|
|
block_out_channels=config["audio_encoder"]["block_out_channels"],
|
|
downsample_factors=config["audio_encoder"]["downsample_factors"],
|
|
dropout=config["audio_encoder"]["dropout"],
|
|
attn_blocks=config["audio_encoder"]["attn_blocks"],
|
|
gradient_checkpointing=gradient_checkpointing,
|
|
)
|
|
|
|
self.visual_encoder = DownEncoder2D(
|
|
in_channels=config["visual_encoder"]["in_channels"],
|
|
block_out_channels=config["visual_encoder"]["block_out_channels"],
|
|
downsample_factors=config["visual_encoder"]["downsample_factors"],
|
|
dropout=config["visual_encoder"]["dropout"],
|
|
attn_blocks=config["visual_encoder"]["attn_blocks"],
|
|
gradient_checkpointing=gradient_checkpointing,
|
|
)
|
|
|
|
self.eval()
|
|
|
|
def forward(self, image_sequences, audio_sequences):
|
|
vision_embeds = self.visual_encoder(image_sequences) # (b, c, 1, 1)
|
|
audio_embeds = self.audio_encoder(audio_sequences) # (b, c, 1, 1)
|
|
|
|
vision_embeds = vision_embeds.reshape(vision_embeds.shape[0], -1) # (b, c)
|
|
audio_embeds = audio_embeds.reshape(audio_embeds.shape[0], -1) # (b, c)
|
|
|
|
# Make them unit vectors
|
|
vision_embeds = F.normalize(vision_embeds, p=2, dim=1)
|
|
audio_embeds = F.normalize(audio_embeds, p=2, dim=1)
|
|
|
|
return vision_embeds, audio_embeds
|
|
|
|
|
|
class ResnetBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
dropout: float = 0.0,
|
|
norm_num_groups: int = 32,
|
|
eps: float = 1e-6,
|
|
act_fn: str = "silu",
|
|
downsample_factor=2,
|
|
):
|
|
super().__init__()
|
|
|
|
self.norm1 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True)
|
|
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
|
|
|
self.norm2 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=out_channels, eps=eps, affine=True)
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
|
|
|
if act_fn == "relu":
|
|
self.act_fn = nn.ReLU()
|
|
elif act_fn == "silu":
|
|
self.act_fn = nn.SiLU()
|
|
|
|
if in_channels != out_channels:
|
|
self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
|
else:
|
|
self.conv_shortcut = None
|
|
|
|
if isinstance(downsample_factor, list):
|
|
downsample_factor = tuple(downsample_factor)
|
|
|
|
if downsample_factor == 1:
|
|
self.downsample_conv = None
|
|
else:
|
|
self.downsample_conv = nn.Conv2d(
|
|
out_channels, out_channels, kernel_size=3, stride=downsample_factor, padding=0
|
|
)
|
|
self.pad = (0, 1, 0, 1)
|
|
if isinstance(downsample_factor, tuple):
|
|
if downsample_factor[0] == 1:
|
|
self.pad = (0, 1, 1, 1) # The padding order is from back to front
|
|
elif downsample_factor[1] == 1:
|
|
self.pad = (1, 1, 0, 1)
|
|
|
|
def forward(self, input_tensor):
|
|
hidden_states = input_tensor
|
|
|
|
hidden_states = self.norm1(hidden_states)
|
|
hidden_states = self.act_fn(hidden_states)
|
|
|
|
hidden_states = self.conv1(hidden_states)
|
|
hidden_states = self.norm2(hidden_states)
|
|
hidden_states = self.act_fn(hidden_states)
|
|
|
|
hidden_states = self.dropout(hidden_states)
|
|
hidden_states = self.conv2(hidden_states)
|
|
|
|
if self.conv_shortcut is not None:
|
|
input_tensor = self.conv_shortcut(input_tensor)
|
|
|
|
hidden_states += input_tensor
|
|
|
|
if self.downsample_conv is not None:
|
|
hidden_states = F.pad(hidden_states, self.pad, mode="constant", value=0)
|
|
hidden_states = self.downsample_conv(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class AttentionBlock2D(nn.Module):
|
|
def __init__(self, query_dim, norm_num_groups=32, dropout=0.0):
|
|
super().__init__()
|
|
self.norm1 = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=query_dim, eps=1e-6, affine=True)
|
|
self.norm2 = nn.LayerNorm(query_dim)
|
|
self.norm3 = nn.LayerNorm(query_dim)
|
|
|
|
self.ff = FeedForward(query_dim, dropout=dropout, activation_fn="geglu")
|
|
|
|
self.conv_in = nn.Conv2d(query_dim, query_dim, kernel_size=1, stride=1, padding=0)
|
|
self.conv_out = nn.Conv2d(query_dim, query_dim, kernel_size=1, stride=1, padding=0)
|
|
|
|
self.attn = Attention(query_dim=query_dim, heads=8, dim_head=query_dim // 8, dropout=dropout, bias=True)
|
|
|
|
def forward(self, hidden_states):
|
|
assert hidden_states.dim() == 4, f"Expected hidden_states to have ndim=4, but got ndim={hidden_states.dim()}."
|
|
|
|
batch, channel, height, width = hidden_states.shape
|
|
residual = hidden_states
|
|
|
|
hidden_states = self.norm1(hidden_states)
|
|
hidden_states = self.conv_in(hidden_states)
|
|
hidden_states = rearrange(hidden_states, "b c h w -> b (h w) c")
|
|
|
|
norm_hidden_states = self.norm2(hidden_states)
|
|
|
|
hidden_states = self.attn(norm_hidden_states, attention_mask=None) + hidden_states
|
|
hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
|
|
|
|
hidden_states = rearrange(hidden_states, "b (h w) c -> b c h w", h=height, w=width)
|
|
hidden_states = self.conv_out(hidden_states)
|
|
|
|
hidden_states = hidden_states + residual
|
|
return hidden_states
|
|
|
|
|
|
class DownEncoder2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels=4 * 16,
|
|
block_out_channels=[64, 128, 256, 256],
|
|
downsample_factors=[2, 2, 2, 2],
|
|
layers_per_block=2,
|
|
norm_num_groups=32,
|
|
attn_blocks=[1, 1, 1, 1],
|
|
dropout: float = 0.0,
|
|
act_fn="silu",
|
|
gradient_checkpointing=False,
|
|
):
|
|
super().__init__()
|
|
self.layers_per_block = layers_per_block
|
|
self.gradient_checkpointing = gradient_checkpointing
|
|
|
|
# in
|
|
self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1)
|
|
|
|
# down
|
|
self.down_blocks = nn.ModuleList([])
|
|
|
|
output_channels = block_out_channels[0]
|
|
for i, block_out_channel in enumerate(block_out_channels):
|
|
input_channels = output_channels
|
|
output_channels = block_out_channel
|
|
# is_final_block = i == len(block_out_channels) - 1
|
|
|
|
down_block = ResnetBlock2D(
|
|
in_channels=input_channels,
|
|
out_channels=output_channels,
|
|
downsample_factor=downsample_factors[i],
|
|
norm_num_groups=norm_num_groups,
|
|
dropout=dropout,
|
|
act_fn=act_fn,
|
|
)
|
|
|
|
self.down_blocks.append(down_block)
|
|
|
|
if attn_blocks[i] == 1:
|
|
attention_block = AttentionBlock2D(query_dim=output_channels, dropout=dropout)
|
|
self.down_blocks.append(attention_block)
|
|
|
|
# out
|
|
self.norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
|
|
self.act_fn_out = nn.ReLU()
|
|
|
|
def forward(self, hidden_states):
|
|
hidden_states = self.conv_in(hidden_states)
|
|
|
|
# down
|
|
for down_block in self.down_blocks:
|
|
if self.gradient_checkpointing:
|
|
hidden_states = torch.utils.checkpoint.checkpoint(down_block, hidden_states, use_reentrant=False)
|
|
else:
|
|
hidden_states = down_block(hidden_states)
|
|
|
|
# post-process
|
|
hidden_states = self.norm_out(hidden_states)
|
|
hidden_states = self.act_fn_out(hidden_states)
|
|
|
|
return hidden_states
|