residual_diffusion.py

import math
import os
import time

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from transformer_denoiser import lengths_to_mask, TransformerDenoiser
from react_dataset import ReactDataset


#DIFFUSION

def cosine_beta_schedule(num_timesteps, max_beta=0.99):
    #Schedule from Improved DDPM; Nichol, Dhariwal 2021

    betas = []
    for i in range(num_timesteps):
        t1 = i / num_timesteps
        t2 = (i + 1) / num_timesteps
        alpha_bar_t1 = math.cos(((t1 + 0.008) / 1.008) * (math.pi / 2))
        alpha_bar_t2 = math.cos(((t2 + 0.008) / 1.008) * (math.pi / 2))
        
        if alpha_bar_t1 < 1e-5:
            alpha_bar_t1 = 1e-5
        
        beta = 1 - (alpha_bar_t2 / alpha_bar_t1)
        if beta <= max_beta:
            betas.append(beta)
        else:
            betas.append(max_beta)
    return betas

class ResidualDiffusion:
    def __init__(self, num_train_timesteps=1000, num_inference_timesteps=50):
        self.num_train_timesteps = num_train_timesteps
        self.num_inference = num_inference_timesteps

        betas = cosine_beta_schedule(num_train_timesteps)
        self.betas = torch.tensor(betas, dtype=torch.float32) #(1000, )
        alphas = 1.0 - self.betas 
        self.alpha_bars = torch.cumprod(alphas, dim=0).float() #(1000, )

        step_ratio = num_train_timesteps / num_inference_timesteps
        self.ddim_timesteps = torch.arange(num_inference_timesteps - 1, -1, -1) * step_ratio

    def extract(self, schedule, t):
        out = schedule.to(t.device).gather(0, t)
        return out.view(-1, 1, 1)

    def q_sample(
            self, 
            x_0, #(B, T, 25)
            t, #(B, )
            noise=None
    ):
        if noise is None:
            noise = torch.randn_like(x_0) #(B, T, 25)
        alpha_bar = self.extract(self.alpha_bars, t) #(B, 1, 1)
        return torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1.0 - alpha_bar) * noise #(B, T, 25)
    
    def training_losses(
            self,
            model, #TransformerDenoiser 
            x_0, #(B, T, 25) clean listener reaction
            speaker_audio, #(B, T, 768)
            speaker_3dmm, #(B, T, 58)
            speaker_emotion, #(B, T, 25)
            motion_length, #(B, )
            mean_r, #(25, )
            std_r #(25, )
    ):
        b = x_0.shape[0]
        device = x_0.device

        mean_r = mean_r.to(device)
        std_r = std_r.to(device)

        #sample timesteps from uniform random distribution
        t = torch.randint(0, self.num_train_timesteps, (b, ), device=device) #(B, )

        #Calculate residual
        mean_r = mean_r.unsqueeze(0).unsqueeze(0) #(1, 1, 25)
        std_r = std_r.unsqueeze(0).unsqueeze(0) #(1, 1, 25)
        residual = (x_0 - speaker_emotion) 
        r_0 = (residual - mean_r) / std_r
        
        r_t = self.q_sample(r_0, t) #(B, T, 25)
        pred_r0 = model(r_t, t, speaker_audio, speaker_3dmm, speaker_emotion, motion_length) #(B, T, 25)
        
        return {"pred":pred_r0, "target": r_0}
    
    @torch.no_grad()
    def ddim_sample(self, model, shape, speaker_audio, speaker_3dmm, speaker_emotion,
                    mean_r, std_r, motion_length, eta=0.0, guidance_weight=1.0, stopping_frequency=None):
        device = speaker_audio.device
        x = torch.randn(shape, device=device)
        mean_r = mean_r.to(device)
        std_r = std_r.to(device)

        for i in range(x.ndim - 1):
            mean_r = mean_r.unsqueeze(0)
            std_r = std_r.unsqueeze(0)

        if stopping_frequency is not None:
            path = [x * std_r + mean_r + speaker_emotion]

        for idx, t_val in enumerate(self.ddim_timesteps):
            t = torch.full((shape[0],), int(t_val), device=device, dtype=torch.long)
            t_prev_val = int(self.ddim_timesteps[idx + 1]) if idx < len(self.ddim_timesteps) - 1 else 0
            t_prev = torch.full((shape[0],), t_prev_val, device=device, dtype=torch.long)

            ab = self.extract(self.alpha_bars, t)
            ab_prev = self.extract(self.alpha_bars, t_prev)

            pred_cond = model(x, t, speaker_audio, speaker_3dmm,
                            speaker_emotion, motion_length)
            if guidance_weight != 1.0:
                pred_uncond = model(x, t, speaker_audio, speaker_3dmm,
                                    speaker_emotion, motion_length, force_uncond=True)
                pred_x0 = pred_uncond + guidance_weight * (pred_cond - pred_uncond)
            else:
                pred_x0 = pred_cond

            eps = (x - torch.sqrt(ab) * pred_x0) / torch.sqrt(1.0 - ab)
            sigma = eta * torch.sqrt((1 - ab_prev) / (1 - ab)) * torch.sqrt(1 - ab / ab_prev)
            mean_pred = pred_x0 * torch.sqrt(ab_prev) + torch.sqrt(1 - ab_prev - sigma ** 2) * eps
            x = mean_pred + sigma * torch.randn_like(x) if t_val > 0 else mean_pred

            if stopping_frequency is not None and (idx + 1) % stopping_frequency == 0:
                path.append(x * std_r + mean_r + speaker_emotion)

        
        # x -> DDIM predicted residual

        if stopping_frequency is not None:
            return path
        
        return x * std_r + mean_r + speaker_emotion


# Weighted MSE Loss
def compute_loss(pred, target, motion_length=None, w_au=1.0, w_va=5.0, w_em=5.0):
    if motion_length is not None:
        # (B, T, 1): 1.0 on real frames, 0.0 on padded frames
        m = lengths_to_mask(motion_length, pred.device, max_len=pred.shape[1]).unsqueeze(-1).float()
        def mse(a, b):
            return (((a - b) ** 2) * m).sum() / m.sum().clamp(min=1.0) / a.shape[-1]
    else:
        mse = F.mse_loss

    loss_au = mse(pred[..., :15],  target[..., :15])
    loss_va = mse(pred[..., 15:17], target[..., 15:17])
    loss_em = mse(pred[..., 17:],  target[..., 17:])
    loss = w_au * loss_au + w_va * loss_va + w_em * loss_em
    return {"loss": loss, "loss_au": loss_au, "loss_va": loss_va, "loss_em": loss_em}


def compute_eval_loss(pred, target, motion_length=None):
    assert pred.size() == target.size(), f"Prediction target mismatch: {pred.size()} {target.size()}"

    B, T, C = pred.shape

    if motion_length is not None:
        valid = lengths_to_mask(motion_length, pred.device, max_len=T)  # (B, T)
        m = valid.unsqueeze(-1).float()                                 # (B, T, 1)
    else:
        m = torch.ones(B, T, 1, device=pred.device)

    denom = m.sum(dim=1).clamp(min=1.0)                # (B, 1) valid frames per sample

    # per-channel means over valid frames
    pred_mean = (pred * m).sum(dim=1) / denom          # (B, C)
    target_mean = (target * m).sum(dim=1) / denom      # (B, C)

    # zero out padded frames so they contribute nothing
    pred_centered = (pred - pred_mean.unsqueeze(1)) * m       # (B, T, C)
    target_centered = (target - target_mean.unsqueeze(1)) * m # (B, T, C)

    pred_var = pred_centered.square().sum(dim=1) / denom      # (B, C)
    target_var = target_centered.square().sum(dim=1) / denom  # (B, C)
    covariance = (pred_centered * target_centered).sum(dim=1) / denom  # (B, C)

    diff_squared = (pred_mean - target_mean).square()         # (B, C)

    ccc_per_channel = 2 * covariance / (pred_var + target_var + diff_squared + 1e-8)  # (B, C)

    return 1 - ccc_per_channel.mean()  # mean over channels and batch


# TRAINING

@torch.no_grad()
def evaluate(model, loader, diffusion, device, mean_r, std_r, seed=888):
    model.eval()
    rng_state = torch.get_rng_state()

    if seed is not None:
        torch.manual_seed(seed)
    
    loss_ccc = 0.0
    n = 0

    for batch in loader:
        x0 = batch["listener_emotion"].to(device)
        sp_audio = batch["speaker_audio"].to(device)
        sp_dmm = batch["speaker_3dmm"].to(device)
        sp_emo = batch["speaker_emotion"].to(device)
        motion_len = batch["motion_length"].to(device)

        out = diffusion.training_losses(model, x0, sp_audio, sp_dmm, sp_emo, motion_len, mean_r, std_r)
        L_ccc = compute_eval_loss(out["pred"] * std_r + mean_r + sp_emo, out["target"] * std_r + mean_r + sp_emo, motion_length=motion_len)
            
        b = x0.shape[0]
        loss_ccc += L_ccc.item() * b
        n += b
    
    torch.set_rng_state(rng_state)

    loss_ccc = loss_ccc / max(n, 1)
    return loss_ccc

def main(
        data_dir="./data",
        clip_length=750,
        batch_size=2,
        epochs=500,
        distribution_path=None,
        save_ckpt_path=os.path.join("checkpoints", "residual_diffusion_best.pt"),
        num_workers=8,
        use_amp=True,
        seed=888,
        num_layers=9
):
    torch.manual_seed(seed)

    if torch.cuda.is_available():
        device = "cuda"
        print("Connected to GPU")
    else:
        device = "cpu"
        print("Unable to connected to GPU; connected to CPU")

    amp_enabled = use_amp and device == "cuda"
    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)

    os.makedirs(os.path.dirname(save_ckpt_path), exist_ok=True)
    
    train_set = ReactDataset(data_dir, split="train", clip_length=clip_length)
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, persistent_workers=num_workers > 0)

    val_workers = max(num_workers // 2, 0)
    val_set = ReactDataset(data_dir, split="val", clip_length=clip_length, random_crop=False)
    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=val_workers, pin_memory=True, persistent_workers=val_workers > 0)

    model = TransformerDenoiser(num_layers=num_layers).to(device)
    diffusion = ResidualDiffusion()

    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)

    best_val = float("inf")

    if distribution_path is not None:
        distribution = torch.load(distribution_path)
        mean_r = distribution["mean"].to(device)
        std_r = distribution["std"].to(device)
        direct_residual = False
    else:
        mean_r = torch.zeros((25, ), device=device)
        std_r = torch.ones((25, ), device=device)
        direct_residual = True

    print("Set up complete")

    for epoch in range(epochs):
        start_time = time.time()

        model.train()
        running = 0.0

        for step, batch in enumerate(train_loader):
            x0 = batch["listener_emotion"].to(device)
            sp_audio = batch["speaker_audio"].to(device)
            sp_dmm = batch["speaker_3dmm"].to(device)
            sp_emo = batch["speaker_emotion"].to(device)
            motion_len = batch["motion_length"].to(device)

            opt.zero_grad()
            with torch.amp.autocast("cuda", enabled=amp_enabled):
                out = diffusion.training_losses(model, x0, sp_audio, sp_dmm, sp_emo, motion_len, mean_r, std_r)

                if direct_residual:
                    losses = compute_loss(out["pred"], out["target"], motion_length=motion_len)
                else:
                    losses = compute_loss(out["pred"], out["target"], motion_length=motion_len, w_au=1.0, w_va=1.0, w_em=1.0)
                loss = losses["loss"]
            
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()
            running += loss.item()
        
        sched.step()

        print(f"Time train epoch {epoch + 1}: {time.time() - start_time}")

        if (epoch + 1) % 10 == 0:
            start_time = time.time()
            val_loss = evaluate(model, val_loader, diffusion, device, mean_r, std_r)
            
            print(f"Validation: epoch {epoch + 1}; loss_ccc: {val_loss}")

            if val_loss < best_val:
                best_val = val_loss
                save = {
                    "epoch": epoch,
                    "model": model.state_dict(),
                    "opt": opt.state_dict(), 
                    "val_loss": val_loss
                }

                torch.save(save, save_ckpt_path)
            
            print(f"Time test epoch {epoch + 1}: {time.time() - start_time}")

if __name__=="__main__":
    main(epochs=10)