Audio-Visual Speech Separation

import torch
import torch.nn as nn
import torchaudio

class AudioVisualSeparator(nn.Module):
    """
    Simplified Looking to Listen architecture.

    The model takes:
    - Mixed audio spectrogram: (B, 1, F, T)
    - Face embeddings for each speaker: (B, N_speakers, D_face)

    Outputs:
    - Separated spectrograms for each speaker: (B, N_speakers, F, T)
    """

    def __init__(self, n_speakers=2, face_dim=512, audio_channels=257):
        super().__init__()

        # Face encoder (pretrained, e.g., FaceNet)
        self.face_encoder = nn.Sequential(
            nn.Linear(face_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 256)
        )

        # Audio encoder with dilated convolutions
        self.audio_encoder = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=2, dilation=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, padding=4, dilation=4),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )

        # Audio-visual fusion
        self.fusion = nn.Sequential(
            nn.Conv2d(128 + 256, 256, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(256, 128, kernel_size=1),
            nn.ReLU()
        )

        # Mask prediction (one per speaker)
        self.mask_predictor = nn.Sequential(
            nn.Conv2d(128, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, n_speakers, kernel_size=1),
            nn.Sigmoid()
        )

    def forward(self, mixed_audio, face_embeddings):
        """
        Args:
            mixed_audio: (B, 1, F, T) - mixed spectrogram
            face_embeddings: (B, N, D) - face embedding per speaker

        Returns:
            masks: (B, N, F, T) - separation mask per speaker
        """
        B, _, F, T = mixed_audio.shape
        N = face_embeddings.shape[1]

        # Encode audio
        audio_features = self.audio_encoder(mixed_audio)  # (B, 128, F, T)

        # Encode faces and tile across time/frequency
        face_features = self.face_encoder(face_embeddings)  # (B, N, 256)

        # For simplicity, use first speaker's face for conditioning
        # Full model would predict separate masks for each speaker
        face_tiled = face_features[:, 0, :].unsqueeze(-1).unsqueeze(-1)
        face_tiled = face_tiled.expand(-1, -1, F, T)  # (B, 256, F, T)

        # Fuse modalities
        fused = torch.cat([audio_features, face_tiled], dim=1)
        fused = self.fusion(fused)

        # Predict masks
        masks = self.mask_predictor(fused)  # (B, N, F, T)

        return masks

    def separate(self, mixed_audio, face_embeddings):
        """Apply masks to separate audio."""
        masks = self.forward(mixed_audio, face_embeddings)

        # Apply masks to mixed spectrogram
        separated = mixed_audio * masks  # (B, N, F, T)

        return separated, masks