61 lines
2.7 KiB
Python
61 lines
2.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
from .modules.attention import TransformerModel
|
|
from .modules.audio import build_audio_encoder, check_frame_rate, encoder_projection
|
|
from .modules.projection import LinearHead
|
|
from .modules.visual_encoder import VisualEncoder
|
|
|
|
#: The visual frontend samples faces onto this timeline, so the audio frontend
|
|
#: has to land on it too or the two streams cannot be fused frame for frame.
|
|
MODEL_FRAME_HZ = 25.0
|
|
|
|
|
|
class AudioVisualASD(nn.Module):
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
cross_dim = cfg["cross"]["d_model"]
|
|
temporal_dim = cfg["temporal"]["d_model"]
|
|
if temporal_dim != cross_dim * 2:
|
|
raise ValueError("ASD temporal d_model must be twice cross d_model")
|
|
|
|
self.audio_encoder = build_audio_encoder(cfg)
|
|
check_frame_rate(self.audio_encoder, MODEL_FRAME_HZ, "the ASD visual timeline")
|
|
self.encoder_proj = encoder_projection(self.audio_encoder.out_dim, cross_dim)
|
|
self.visual_encoder = VisualEncoder()
|
|
if self.visual_encoder.out_dim != cross_dim:
|
|
raise ValueError(
|
|
f"ASD cross d_model must be {self.visual_encoder.out_dim} to match the "
|
|
f"visual frontend, got {cross_dim}"
|
|
)
|
|
self.norm_audio = nn.LayerNorm(cross_dim)
|
|
self.norm_visual = nn.LayerNorm(cross_dim)
|
|
self.cross_v2a = TransformerModel(cfg["cross"], cross_attention=True)
|
|
self.cross_a2v = TransformerModel(cfg["cross"], cross_attention=True)
|
|
self.transformer = TransformerModel(cfg["temporal"])
|
|
self.asd_head = LinearHead(temporal_dim, 6)
|
|
self.audio_head = LinearHead(cross_dim, 1)
|
|
self.visual_head = LinearHead(cross_dim, 1)
|
|
|
|
def forward(self, audio, visual, return_embeddings=False):
|
|
audio_embedding = self.norm_audio(
|
|
self.encoder_proj(self.audio_encoder(audio))
|
|
)
|
|
visual_embedding = self.norm_visual(self.visual_encoder(visual))
|
|
if audio_embedding.shape[1] != visual_embedding.shape[1]:
|
|
raise ValueError(
|
|
f"audio and visual lengths differ: {audio_embedding.shape[1]} and "
|
|
f"{visual_embedding.shape[1]}"
|
|
)
|
|
|
|
audio_cross = self.cross_v2a(audio_embedding, src=visual_embedding)
|
|
visual_cross = self.cross_a2v(visual_embedding, src=audio_embedding)
|
|
fused = self.transformer(torch.cat([audio_cross, visual_cross], dim=-1))
|
|
projection_logits = self.asd_head(fused)
|
|
audio_logits = self.audio_head(audio_cross)
|
|
visual_logits = self.visual_head(visual_cross)
|
|
|
|
if return_embeddings:
|
|
return projection_logits, audio_logits, visual_logits, fused, audio_embedding, visual_embedding
|
|
return projection_logits, audio_logits, visual_logits
|