Spaces:
Sleeping
Sleeping
| """ | |
| MultiSense-DF — Lip-Sync Consistency Module | |
| SyncNet-derived dual-stream network: mouth crops + mel-spectrogram | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class MouthVisualNet(nn.Module): | |
| """5-layer CNN for 96×96 mouth-region crops.""" | |
| def __init__(self, out_dim=256): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Conv2d(3, 32, 3, 2, 1), nn.BatchNorm2d(32), nn.ReLU(), # 48 | |
| nn.Conv2d(32, 64, 3, 2, 1), nn.BatchNorm2d(64), nn.ReLU(), # 24 | |
| nn.Conv2d(64, 128, 3, 2, 1), nn.BatchNorm2d(128), nn.ReLU(),# 12 | |
| nn.Conv2d(128, 256, 3, 2, 1), nn.BatchNorm2d(256), nn.ReLU(),# 6 | |
| nn.Conv2d(256, 512, 3, 2, 1), nn.BatchNorm2d(512), nn.ReLU(),# 3 | |
| nn.AdaptiveAvgPool2d(1) | |
| ) | |
| self.fc = nn.Linear(512, out_dim) | |
| def forward(self, x): | |
| # x: (B*T, 3, 96, 96) | |
| return F.normalize(self.fc(self.net(x).squeeze(-1).squeeze(-1)), dim=-1) | |
| class AudioSpecNet(nn.Module): | |
| """5-layer CNN for mel-spectrogram windows aligned to mouth frames.""" | |
| def __init__(self, out_dim=256): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Conv2d(1, 32, 3, 2, 1), nn.BatchNorm2d(32), nn.ReLU(), | |
| nn.Conv2d(32, 64, 3, 2, 1), nn.BatchNorm2d(64), nn.ReLU(), | |
| nn.Conv2d(64, 128, 3, 2, 1), nn.BatchNorm2d(128), nn.ReLU(), | |
| nn.Conv2d(128, 256, 3, 2, 1), nn.BatchNorm2d(256), nn.ReLU(), | |
| nn.Conv2d(256, 512, 3, (2, 1), 1), nn.BatchNorm2d(512), nn.ReLU(), | |
| nn.AdaptiveAvgPool2d(1) | |
| ) | |
| self.fc = nn.Linear(512, out_dim) | |
| def forward(self, x): | |
| # x: (B*T, 1, 80, W) mel-spectrogram window | |
| return F.normalize(self.fc(self.net(x).squeeze(-1).squeeze(-1)), dim=-1) | |
| class LipSyncBranch(nn.Module): | |
| """ | |
| SyncNet-derived lip-sync consistency module. | |
| Computes per-frame cosine similarity between audio and visual embeddings, | |
| then summarises temporal synchrony profile with a 2-layer LSTM. | |
| Input : mouth_crops (B, T, 3, 96, 96) | |
| mel_specs (B, T, 1, 80, W) | |
| Output: (B, 512) lip-sync embedding + (B, 1) logit | |
| """ | |
| def __init__(self, sync_dim=256, embed_dim=512, dropout=0.1): | |
| super().__init__() | |
| self.visual_net = MouthVisualNet(out_dim=sync_dim) | |
| self.audio_net = AudioSpecNet(out_dim=sync_dim) | |
| # Summarise temporal synchrony profile | |
| self.lstm = nn.LSTM( | |
| input_size=sync_dim * 2 + 1, # [v_feat, a_feat, cos_sim] | |
| hidden_size=embed_dim // 2, | |
| num_layers=2, | |
| batch_first=True, | |
| bidirectional=True, | |
| dropout=dropout | |
| ) | |
| self.norm = nn.LayerNorm(embed_dim) | |
| self.classifier = nn.Sequential( | |
| nn.Linear(embed_dim, 256), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(256, 1) | |
| ) | |
| def forward(self, mouth_crops, mel_specs): | |
| B, T, C, H, W = mouth_crops.shape | |
| _, _, _, F_mel, W_mel = mel_specs.shape | |
| # Per-frame embeddings | |
| v = self.visual_net(mouth_crops.reshape(B * T, C, H, W)) # (B*T, 256) | |
| a = self.audio_net(mel_specs.reshape(B * T, 1, F_mel, W_mel)) # (B*T, 256) | |
| v = v.reshape(B, T, -1) # (B, T, 256) | |
| a = a.reshape(B, T, -1) # (B, T, 256) | |
| # Per-frame cosine similarity (scalar) | |
| cos_sim = F.cosine_similarity(v, a, dim=-1, eps=1e-8).unsqueeze(-1) # (B, T, 1) | |
| # Concatenate features + similarity for LSTM input | |
| seq = torch.cat([v, a, cos_sim], dim=-1) # (B, T, 513) | |
| out, _ = self.lstm(seq) # (B, T, 512) | |
| embed = self.norm(out[:, -1]) # last timestep → (B, 512) | |
| logit = self.classifier(embed) # (B, 1) | |
| return embed, logit, cos_sim.squeeze(-1) # also return sync profile | |
| def contrastive_loss(self, mouth_crops, mel_specs, is_synced): | |
| """ | |
| Contrastive loss for pre-training: in-sync pairs pull together, | |
| offset pairs push apart. | |
| is_synced: (B,) bool tensor | |
| """ | |
| B, T, C, H, W = mouth_crops.shape | |
| _, _, _, F_mel, W_mel = mel_specs.shape | |
| v = self.visual_net(mouth_crops.reshape(B * T, C, H, W)).reshape(B, T, -1).mean(1) | |
| a = self.audio_net(mel_specs.reshape(B * T, 1, F_mel, W_mel)).reshape(B, T, -1).mean(1) | |
| sim = F.cosine_similarity(v, a, dim=-1) # (B,) | |
| margin = 0.4 | |
| loss_pos = (1 - sim) * is_synced.float() | |
| loss_neg = torch.clamp(sim - margin, min=0) * (~is_synced).float() | |
| return (loss_pos + loss_neg).mean() | |
| if __name__ == '__main__': | |
| model = LipSyncBranch() | |
| mouths = torch.randn(2, 125, 3, 96, 96) | |
| mels = torch.randn(2, 125, 1, 80, 16) | |
| emb, logit, sync = model(mouths, mels) | |
| print(f'LipSync embed: {emb.shape}, logit: {logit.shape}, sync profile: {sync.shape}') | |