modules/syncnet/models.py

import torch
from torch import nn
from torch.nn import functional as F
import numpy as np
import math


class Conv1d(nn.Module):
    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conv_block = nn.Sequential(
                            nn.Conv1d(cin, cout, kernel_size, stride, padding),
                            nn.BatchNorm1d(cout)
                            )
        self.act = nn.ReLU()
        self.residual = residual

    def forward(self, x):
        out = self.conv_block(x)
        if self.residual:
            out += x
        return self.act(out)

class LossScale(nn.Module):
    def __init__(self, init_w=10.0, init_b=-5.0):
        super(LossScale, self).__init__()
        
        self.wC = nn.Parameter(torch.tensor(init_w))
        self.bC = nn.Parameter(torch.tensor(init_b))

class CLIPLoss(nn.Module):
    def __init__(self,):
        super().__init__()

    def forward(self, audio_features, motion_features, logit_scale, clip_mask=None):
        logits_per_audio = logit_scale * audio_features @ motion_features.T # [b,c]
        logits_per_motion = logit_scale * motion_features @ audio_features.T # [b,c]
        if clip_mask is not None:
            logits_per_audio += clip_mask        
            logits_per_motion += clip_mask
        labels = torch.arange(logits_per_motion.shape[0]).to(logits_per_motion.device)
        motion_loss = F.cross_entropy(logits_per_motion, labels)
        audio_loss = F.cross_entropy(logits_per_audio, labels)
        clip_loss = (motion_loss + audio_loss) / 2
        ret = {
            "audio_loss": audio_loss,
            "motion_loss": motion_loss,
            "clip_loss": clip_loss
        }
        return ret
        
    def compute_sync_conf(self, audio_features, motion_features, return_matrix=False):
        logits_per_audio = audio_features @ motion_features.T # [b,c]
        if return_matrix:
            return logits_per_audio
        return logits_per_audio[range(len(audio_features)), range(len(audio_features))]

class LandmarkHubertSyncNet(nn.Module):
    def __init__(self, lm_dim=60, audio_dim=1024, num_layers_per_block=3, base_hid_size=128, out_dim=512):
        super(LandmarkHubertSyncNet, self).__init__()

        self.clip_loss_fn = CLIPLoss()
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) * 0
        self.logit_scale_2 = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) * 0
        self.logit_scale_max = math.log(1. / 0.01)

        # hubert = torch.rand(B, 1024, t=10)
        hubert_layers = [
            Conv1d(audio_dim, base_hid_size, kernel_size=3, stride=1, padding=1)
        ]

        hubert_layers.append(
            Conv1d(base_hid_size, base_hid_size, kernel_size=3, stride=1, padding=1),
        )
        hubert_layers += [
            Conv1d(base_hid_size, base_hid_size, kernel_size=3, stride=1, padding=1, residual=True) for _ in range(num_layers_per_block-1)
        ]

        hubert_layers.append(
            Conv1d(base_hid_size, 2*base_hid_size, kernel_size=3, stride=2, padding=1),
        )
        hubert_layers += [
            Conv1d(2*base_hid_size, 2*base_hid_size, kernel_size=3, stride=1, padding=1, residual=True) for _ in range(num_layers_per_block-1)
        ]

        hubert_layers.append(
            Conv1d(2*base_hid_size, 4*base_hid_size, kernel_size=3, stride=2, padding=1),
        )
        hubert_layers += [
            Conv1d(4*base_hid_size, 4*base_hid_size, kernel_size=3, stride=1, padding=1, residual=True) for _ in range(num_layers_per_block-1)
        ]

        hubert_layers += [
            Conv1d(4*base_hid_size, 4*base_hid_size, kernel_size=3, stride=1, padding=1),
            Conv1d(4*base_hid_size, 4*base_hid_size, kernel_size=3, stride=1, padding=0),
            Conv1d(4*base_hid_size, 4*base_hid_size, kernel_size=1, stride=1, padding=0),
            Conv1d(4*base_hid_size, out_dim, kernel_size=1, stride=1, padding=0),
        ]
        self.hubert_encoder = nn.Sequential(*hubert_layers)

        # mouth = torch.rand(B, 20*3, t=5)
        mouth_layers = [
            Conv1d(lm_dim, 96, kernel_size=3, stride=1, padding=1)
        ]

        mouth_layers.append(
            Conv1d(96, base_hid_size, kernel_size=3, stride=1, padding=1),
        )
        mouth_layers += [
            Conv1d(base_hid_size, base_hid_size, kernel_size=3, stride=1, padding=1, residual=True) for _ in range(num_layers_per_block-1)
        ]

        mouth_layers.append(
            Conv1d(base_hid_size, 2*base_hid_size, kernel_size=3, stride=2, padding=1),
        )
        mouth_layers += [
            Conv1d(2*base_hid_size, 2*base_hid_size, kernel_size=3, stride=1, padding=1, residual=True) for _ in range(num_layers_per_block-1)
        ]

        mouth_layers.append(
            Conv1d(2*base_hid_size, 4*base_hid_size, kernel_size=3, stride=1, padding=1),
        )
        mouth_layers += [
            Conv1d(4*base_hid_size, 4*base_hid_size, kernel_size=3, stride=1, padding=1, residual=True) for _ in range(num_layers_per_block-1)
        ]

        mouth_layers += [
            Conv1d(4*base_hid_size, 4*base_hid_size, kernel_size=3, stride=1, padding=1),
            Conv1d(4*base_hid_size, 4*base_hid_size, kernel_size=3, stride=1, padding=0),
            Conv1d(4*base_hid_size, 4*base_hid_size, kernel_size=1, stride=1, padding=0),
            Conv1d(4*base_hid_size, out_dim, kernel_size=1, stride=1, padding=0),
        ]
        self.mouth_encoder = nn.Sequential(*mouth_layers)

        self.lm_dim = lm_dim
        self.audio_dim = audio_dim
        self.logloss = nn.BCELoss()

    def forward(self, hubert, mouth_lm): 
        # hubert := (B, T=10, C=1024)
        # mouth_lm3d := (B, T=5, C=60)
        hubert = hubert.transpose(1,2)
        mouth_lm = mouth_lm.transpose(1,2)
        mouth_embedding = self.mouth_encoder(mouth_lm)
        audio_embedding = self.hubert_encoder(hubert)
        audio_embedding = audio_embedding.view(audio_embedding.size(0), -1)
        mouth_embedding = mouth_embedding.view(mouth_embedding.size(0), -1)
        audio_embedding = F.normalize(audio_embedding, p=2, dim=1)
        mouth_embedding = F.normalize(mouth_embedding, p=2, dim=1)
        return audio_embedding, mouth_embedding

    def cal_sync_loss(self, audio_embedding, mouth_embedding, label, reduction='none'):
        if isinstance(label, torch.Tensor): # finegrained label
            gt_d = label.float().view(-1).to(audio_embedding.device)
        else: # int to represent global label, 1 denotes positive, and 0 denotes negative, used when calculate sync loss for other models
            gt_d = (torch.ones([audio_embedding.shape[0]]) * label).float().to(audio_embedding.device) # int
        d = F.cosine_similarity(audio_embedding, mouth_embedding) # [B]
        loss = F.binary_cross_entropy(d.reshape([audio_embedding.shape[0],]), gt_d, reduction=reduction)
        return loss, d

    def cal_clip_loss(self, audio_embedding, mouth_embedding, clip_mask=None):
        # logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
        logit_scale = 1
        clip_ret = self.clip_loss_fn(audio_embedding, mouth_embedding, logit_scale, clip_mask=clip_mask)
        loss = clip_ret['clip_loss']
        return loss

    def cal_clip_loss_local(self, audio_embedding, mouth_embedding, clip_mask=None):
        # logit_scale = torch.clamp(self.logit_scale_2, max=self.logit_scale_max).exp()
        logit_scale = 1
        clip_ret = self.clip_loss_fn(audio_embedding, mouth_embedding, logit_scale, clip_mask=clip_mask)
        loss = clip_ret['clip_loss']
        return loss

    def compute_sync_conf(self, audio_embedding, mouth_embedding, return_matrix=False):
        # logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
        logit_scale = 1
        clip_ret = self.clip_loss_fn.compute_sync_conf(audio_embedding, mouth_embedding, return_matrix)
        return clip_ret

if __name__ == '__main__':
    syncnet = LandmarkHubertSyncNet(lm_dim=204)
    hubert = torch.rand(2, 10, 1024)
    lm = torch.rand(2, 5, 204)
    mel_embedding, exp_embedding = syncnet(hubert, lm)
    label = torch.tensor([1., 0.])
    loss = syncnet.cal_sync_loss(mel_embedding, exp_embedding, label)
    print(" ")