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)