# Copyright (c) Meta Platforms, Inc. # All rights reserved. import os import json import argparse from pathlib import Path from tqdm import tqdm import torch import torch.distributed as dist_torch import torch.nn.functional as F import numpy as np from PIL import Image import lpips from dreamsim import dreamsim from torchvision import transforms from torcheval.metrics import FrechetInceptionDistance import soundfile as sf import resampy import distributed as dist import librosa from skimage.metrics import structural_similarity as sk_ssim from mel_scale import MelScale # ----------------------------- # Safe, lazy import for FAD (avoid argparse conflicts from dependencies) # ----------------------------- def safe_import_fad(): """ Import frechet_audio_distance.FrechetAudioDistance without letting downstream libraries parse our CLI args during import time. """ import importlib, sys argv_backup = sys.argv[:] try: sys.argv = [argv_backup[0]] # hide our CLI flags from misbehaving imports fad_mod = importlib.import_module("frechet_audio_distance") return getattr(fad_mod, "FrechetAudioDistance") finally: sys.argv = argv_backup # ----------------------------- # Distributed init # ----------------------------- def setup_distributed(): if "RANK" in os.environ and "WORLD_SIZE" in os.environ and "LOCAL_RANK" in os.environ: rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) local_rank = int(os.environ["LOCAL_RANK"]) else: return 0, 1, 0 os.environ.setdefault("MASTER_ADDR", "127.0.0.1") os.environ.setdefault("MASTER_PORT", "29500") assert torch.cuda.is_available(), "CUDA Unavailable" assert torch.cuda.device_count() > local_rank, "local_rank out of the number of GPUs" torch.cuda.set_device(local_rank) dist_torch.init_process_group( backend="nccl", init_method="env://", rank=rank, world_size=world_size, ) dist_torch.barrier() if rank == 0: print(f"[init] world_size={world_size} | rank->gpu OK") return rank, world_size, local_rank # ----------------------------- # Vision metrics factory # ----------------------------- def get_loss_fn(loss_fn_type, secs, device): if loss_fn_type == 'lpips': general_lpips_loss_fn = lpips.LPIPS(net='alex').to(device).eval() def loss_fn(img0_paths, img1_paths): img0_list, img1_list = [], [] for p0, p1 in zip(img0_paths, img1_paths): img0 = lpips.im2tensor(lpips.load_image(p0)).to(device) # [-1,1] img1 = lpips.im2tensor(lpips.load_image(p1)).to(device) img0_list.append(img0) img1_list.append(img1) all_img0 = torch.cat(img0_list, dim=0) all_img1 = torch.cat(img1_list, dim=0) with torch.no_grad(): dist_val = general_lpips_loss_fn.forward(all_img0, all_img1) return dist_val.mean() elif loss_fn_type == 'dreamsim': dreamsim_loss_fn, preprocess = dreamsim(pretrained=True, device=device) dreamsim_loss_fn.eval() def loss_fn(img0_paths, img1_paths): img0_list, img1_list = [], [] for p0, p1 in zip(img0_paths, img1_paths): img0 = preprocess(Image.open(p0)).to(device) img1 = preprocess(Image.open(p1)).to(device) img0_list.append(img0) img1_list.append(img1) all_img0 = torch.cat(img0_list, dim=0) all_img1 = torch.cat(img1_list, dim=0) with torch.no_grad(): dist_val = dreamsim_loss_fn(all_img0, all_img1) return dist_val.mean() elif loss_fn_type == 'fid': fid_metrics = {} for sec in secs: fid_metrics[sec] = FrechetInceptionDistance(feature_dim=2048).to(device) return fid_metrics else: raise NotImplementedError return loss_fn # ===== Helpers for LSD/SSIM (reproducing AudioMetrics behavior) ===== _EPS = 1e-12 def _ensure_stereo_np(y: np.ndarray): if y.ndim == 1: y = np.stack([y, y], axis=0) elif y.ndim == 2: if y.shape[0] == 1: y = np.concatenate([y, y], axis=0) elif y.shape[0] > 2: y = y[:2, :] else: raise ValueError("Unsupported audio array shape") return y def _wav_to_spectrogram(wav: np.ndarray, rate: int): if rate == 44100: hop_length = 441 n_fft = 2048 elif rate == 16000: hop_length = 160 n_fft = 743 else: raise ValueError("Bad Samplerate (expected 16000 or 44100)") f = np.abs(librosa.stft(wav, hop_length=hop_length, n_fft=n_fft)) # [F, T] f = np.transpose(f, (1, 0)) # [T, F] f_torch = torch.tensor(f[None, None, ...], dtype=torch.float32) # [1,1,T,F] return f_torch def _lsd_from_specs(est: torch.Tensor, target: torch.Tensor) -> torch.Tensor: ratio = (target ** 2) / ((est + _EPS) ** 2) + _EPS lsd = torch.log10(ratio) ** 2 lsd = torch.mean(torch.mean(lsd, dim=3) ** 0.5, dim=2) return lsd.mean() def _mel_lsd_ssim_single( e_wav: np.ndarray, g_wav: np.ndarray, mel_tf: MelScale, n_fft: int = 743, hop_length: int = 160, ) -> tuple[float, float]: est_mag = np.abs(librosa.stft(e_wav, n_fft=n_fft, hop_length=hop_length)) ref_mag = np.abs(librosa.stft(g_wav, n_fft=n_fft, hop_length=hop_length)) est_mag_t = torch.from_numpy(est_mag).float() ref_mag_t = torch.from_numpy(ref_mag).float() est_mel = mel_tf(est_mag_t) ref_mel = mel_tf(ref_mag_t) ex_m = est_mel.transpose(0, 1).unsqueeze(0).unsqueeze(0) gt_m = ref_mel.transpose(0, 1).unsqueeze(0).unsqueeze(0) mel_lsd = float(_lsd_from_specs(ex_m, gt_m)) mel_ssim = float(_ssim_from_specs(ex_m, gt_m)) return mel_lsd, mel_ssim def _to_log_specs(x: torch.Tensor) -> torch.Tensor: return torch.log10(x + _EPS) def _pow_p_norm(x: torch.Tensor) -> torch.Tensor: return torch.mean(x.pow(2), dim=(2, 3)) def _energy_unify(est: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: p_est = _pow_p_norm(est) p_tgt = _pow_p_norm(target) scale = torch.sqrt((p_tgt + _EPS) / (p_est + _EPS)) scale = scale[..., None, None] est_scaled = est * scale return est_scaled, target def _sispec_from_specs(est: torch.Tensor, target: torch.Tensor, log_domain: bool) -> torch.Tensor: if log_domain: est = _to_log_specs(est) target = _to_log_specs(target) est_u, tgt_u = _energy_unify(est, target) noise = est_u - tgt_u snr = ( _pow_p_norm(tgt_u) / (_pow_p_norm(noise) + _EPS) ) + _EPS sp_loss = 10.0 * torch.log10(snr) return sp_loss.mean() # ===== Image PSNR (RGB on [0,1]) ===== def _psnr_from_tensors(gt: torch.Tensor, pred: torch.Tensor, data_range: float = 1.0, eps: float = 1e-10) -> torch.Tensor: mse = torch.mean((gt - pred) ** 2, dim=(1, 2, 3)) dr = torch.as_tensor(data_range, device=gt.device, dtype=gt.dtype) psnr = 10.0 * torch.log10((dr * dr) / (mse + eps)) return psnr def _ssim_from_specs(est: torch.Tensor, target: torch.Tensor) -> float: if est.is_cuda: est_np = est.detach().cpu().numpy() tgt_np = target.detach().cpu().numpy() else: est_np = est.numpy() tgt_np = target.numpy() N, C, _, _ = est_np.shape acc, cnt = 0.0, 0 for n in range(N): for c in range(C): ref = tgt_np[n, c, ...] out = est_np[n, c, ...] rng = float(out.max() - out.min()) rng = 1.0 if rng == 0.0 else rng s = sk_ssim(out, ref, win_size=7, data_range=rng) acc += float(s); cnt += 1 return acc / max(cnt, 1) # ========================================================== # Streaming, DDP-friendly Audio FAD # (embeddings identical to official FrechetAudioDistance) # ========================================================== class _RunningGaussianStats: def __init__(self, feat_dim: int, device: torch.device): self.D = feat_dim self.device = device self.reset() def reset(self): D = self.D self.count = torch.zeros(1, device=self.device, dtype=torch.float64) self.sum_feat = torch.zeros(D, device=self.device, dtype=torch.float64) self.sum_outer = torch.zeros(D, D, device=self.device, dtype=torch.float64) @torch.no_grad() def update(self, feats: torch.Tensor): # [N, D] if feats is None or feats.numel() == 0: return f = feats.to(dtype=torch.float64) self.count += torch.tensor([f.shape[0]], device=self.device, dtype=torch.float64) self.sum_feat += f.sum(dim=0) self.sum_outer += f.t().mm(f) @torch.no_grad() def sync(self): if dist_torch.is_initialized(): for t in (self.count, self.sum_feat, self.sum_outer): dist_torch.all_reduce(t, op=dist_torch.ReduceOp.SUM) @torch.no_grad() def mean_cov(self, eps: float = 1e-6): n = int(self.count.item()) if n == 0: return None, None mean = self.sum_feat / self.count cov = self.sum_outer / self.count - torch.ger(mean, mean) cov = cov + torch.eye(self.D, device=self.device, dtype=torch.float64) * eps return mean, cov @torch.no_grad() def _frechet_distance_torch(mean1, cov1, mean2, cov2) -> float: diff = mean1 - mean2 diff2 = diff.dot(diff) evals1, evecs1 = torch.linalg.eigh(cov1) sqrt1 = evecs1 @ torch.diag(evals1.clamp(min=0).sqrt()) @ evecs1.t() prod = sqrt1 @ cov2 @ sqrt1 evals_prod = torch.linalg.eigvalsh(prod).clamp(min=0).sqrt() trace = torch.trace(cov1 + cov2) - 2.0 * evals_prod.sum() return float((diff2 + trace).item()) class StreamingFAD: """ Mono (downmix) FID-style streaming FAD: - update_from_wavs(paths, is_real=True/False) - compute() # does DDP all_reduce internally """ def __init__(self, fad_backend, pad_seconds: float = 0.96, batch_size: int = 16): self.fad = fad_backend self.device = self.fad.device self.bs = batch_size self.pad_len = int(round(self.fad.sample_rate * float(pad_seconds))) self.feat_dim = self._infer_feat_dim() self.real_stats = _RunningGaussianStats(self.feat_dim, self.device) self.fake_stats = _RunningGaussianStats(self.feat_dim, self.device) def _infer_feat_dim(self) -> int: sr = self.fad.sample_rate x = np.zeros((self.pad_len,), dtype=np.float32) emb = self.fad.get_embeddings([x], sr=sr) return int(emb.shape[-1]) if isinstance(emb, np.ndarray) else int(emb.shape[-1]) @torch.no_grad() def _load_and_resample(self, path: str): try: audio, sr = sf.read(path, dtype="float32", always_2d=False) except Exception as e: print(f"[StreamingFAD] read error: {path}: {e}") return None if audio is None or (isinstance(audio, np.ndarray) and audio.size == 0): return None if isinstance(audio, np.ndarray) and audio.ndim == 2: audio = audio.mean(axis=1) if sr != self.fad.sample_rate: try: audio = resampy.resample(audio, sr, self.fad.sample_rate) except Exception as e: print(f"[StreamingFAD] resample error: {path}: {e}") return None if audio.shape[0] < self.pad_len: pad = np.zeros((self.pad_len - audio.shape[0],), dtype=np.float32) audio = np.concatenate([audio, pad], axis=0) return audio.astype(np.float32, copy=False) @torch.no_grad() def update_from_wavs(self, wav_paths, is_real: bool): if not wav_paths: return xs = [] for p in wav_paths: a = self._load_and_resample(p) if a is not None: xs.append(a) if not xs: return feats_chunks = [] for i in range(0, len(xs), self.bs): chunk = xs[i:i+self.bs] emb_np = self.fad.get_embeddings(chunk, sr=self.fad.sample_rate) if isinstance(emb_np, np.ndarray): if emb_np.size == 0: continue feats_chunks.append(torch.from_numpy(emb_np).to(self.device)) else: if emb_np.numel() == 0: continue feats_chunks.append(emb_np.to(self.device)) if len(feats_chunks) == 0: return feats = torch.cat(feats_chunks, dim=0) (self.real_stats if is_real else self.fake_stats).update(feats) @torch.no_grad() def compute(self) -> float: self.real_stats.sync() self.fake_stats.sync() m1, c1 = self.real_stats.mean_cov() m2, c2 = self.fake_stats.mean_cov() if (m1 is None) or (m2 is None): raise RuntimeError("StreamingFAD: empty stats") return _frechet_distance_torch(m1, c1, m2, c2) class StereoStreamingFAD: def __init__(self, fad_backend, pad_seconds: float = 0.96, batch_size: int = 16): self.fad = fad_backend self.device = self.fad.device self.bs = batch_size self.pad_len = int(round(self.fad.sample_rate * float(pad_seconds))) self.feat_dim = self._infer_feat_dim() self.L_real = _RunningGaussianStats(self.feat_dim, self.device) self.L_fake = _RunningGaussianStats(self.feat_dim, self.device) self.R_real = _RunningGaussianStats(self.feat_dim, self.device) self.R_fake = _RunningGaussianStats(self.feat_dim, self.device) def _infer_feat_dim(self) -> int: sr = self.fad.sample_rate x = np.zeros((self.pad_len,), dtype=np.float32) emb = self.fad.get_embeddings([x], sr=sr) return int(emb.shape[-1]) if isinstance(emb, np.ndarray) else int(emb.shape[-1]) @torch.no_grad() def _load_lr_and_resample_pad(self, path: str): try: audio, sr = sf.read(path, dtype="float32", always_2d=True) # [T, C] except Exception as e: print(f"[StereoFAD] read error: {path}: {e}") return None, None if audio is None or audio.size == 0: return None, None C = audio.shape[1] if C == 1: L = audio[:, 0]; R = audio[:, 0] else: L = audio[:, 0]; R = audio[:, 1] if C >= 2 else audio[:, 0] if sr != self.fad.sample_rate: try: L = resampy.resample(L, sr, self.fad.sample_rate) R = resampy.resample(R, sr, self.fad.sample_rate) except Exception as e: print(f"[StereoFAD] resample error: {path}: {e}") return None, None def _pad_to_len(x: np.ndarray, n: int): if x.shape[0] >= n: return x.astype(np.float32, copy=False) pad = np.zeros((n - x.shape[0],), dtype=np.float32) return np.concatenate([x, pad], axis=0) L = _pad_to_len(L, self.pad_len) R = _pad_to_len(R, self.pad_len) return L, R @torch.no_grad() def update_from_wavs(self, wav_paths, is_real: bool): if not wav_paths: return L_list, R_list = [], [] for p in wav_paths: L, R = self._load_lr_and_resample_pad(p) if L is not None and R is not None: L_list.append(L); R_list.append(R) if not L_list: return def _embed_and_update(xs, stats_obj: _RunningGaussianStats): feats_chunks = [] for i in range(0, len(xs), self.bs): chunk = xs[i:i+self.bs] emb_np = self.fad.get_embeddings(chunk, sr=self.fad.sample_rate) if isinstance(emb_np, np.ndarray): if emb_np.size == 0: continue feats_chunks.append(torch.from_numpy(emb_np).to(self.device)) else: if emb_np.numel() == 0: continue feats_chunks.append(emb_np.to(self.device)) if len(feats_chunks) == 0: return feats = torch.cat(feats_chunks, dim=0) stats_obj.update(feats) if is_real: _embed_and_update(L_list, self.L_real) _embed_and_update(R_list, self.R_real) else: _embed_and_update(L_list, self.L_fake) _embed_and_update(R_list, self.R_fake) @torch.no_grad() def compute(self): for t in (self.L_real, self.L_fake, self.R_real, self.R_fake): t.sync() mL_r, cL_r = self.L_real.mean_cov() mL_f, cL_f = self.L_fake.mean_cov() mR_r, cR_r = self.R_real.mean_cov() mR_f, cR_f = self.R_fake.mean_cov() if (mL_r is None) or (mL_f is None) or (mR_r is None) or (mR_f is None): raise RuntimeError("StereoStreamingFAD: empty stats") fad_left = _frechet_distance_torch(mL_r, cL_r, mL_f, cL_f) fad_right = _frechet_distance_torch(mR_r, cR_r, mR_f, cR_f) fad_mean = 0.5 * (fad_left + fad_right) return float(fad_left), float(fad_right), float(fad_mean) # ----------------------------- # Stereo-friendly Audio Metrics (LSD/SSIM/MelCos/DRMS) # ----------------------------- def _load_librosa_stereo(path: str, sr: int) -> np.ndarray: y, _ = librosa.load(path, sr=sr, mono=False) y = _ensure_stereo_np(y) # (2, T) return y def _mel_cosine_single_channel(wav: np.ndarray, ref: np.ndarray, sr: int, mel_tf: MelScale) -> float: hop_length = 160; n_fft = 743 est_mag = np.abs(librosa.stft(wav, hop_length=hop_length, n_fft=n_fft)) # [F, T] ref_mag = np.abs(librosa.stft(ref, hop_length=hop_length, n_fft=n_fft)) est_mag_t = torch.tensor(est_mag, dtype=torch.float32) # [F,T] ref_mag_t = torch.tensor(ref_mag, dtype=torch.float32) # [F,T] est_mel = mel_tf(est_mag_t) # [80, T] ref_mel = mel_tf(ref_mag_t) # [80, T] sim = F.cosine_similarity(est_mel.flatten(), ref_mel.flatten(), dim=0) return float(sim.item()) # ----------------------------- # Evaluate # ----------------------------- def evaluate(args, dataset_name, eval_type, metric_logger, loss_fns, gt_dir, exp_dir, secs, device, rank, world_size, modals): lpips_loss_fn, dreamsim_loss_fn, fid_loss_fn = loss_fns if eval_type == 'rollout': eval_name = 'rollout' image_idxs = secs.copy() elif eval_type == 'time': eval_name = eval_type image_idxs = secs.copy() else: raise ValueError(f"Unknown eval_type {eval_type}") if 'v' in modals: for s in secs: metric_logger.meters[f'{dataset_name}_{eval_name}_fid_{int(s)}'].update(0.0, n=0) # Episodes split by rank all_eps = sorted([e for e in os.listdir(gt_dir) if os.path.isdir(os.path.join(gt_dir, e))]) eps = all_eps[rank::world_size] if len(eps) == 0: return to_tensor = transforms.ToTensor() fad_streams = {} stereo_mode = False if 'a' in modals: try: FADLib = safe_import_fad() except Exception as e: if rank == 0: print(f"[WARN] Fail to import frechet_audio_distanceļ¼š{e}") FADLib = None if FADLib is not None: base_fad = FADLib( model_name=args.fad_model, sample_rate=args.fad_sr, verbose=False ) if args.fad_model == 'vggish' and not args.mono: stereo_mode = True for sec in secs: fad_streams[sec] = StereoStreamingFAD(base_fad, pad_seconds=args.fad_pad_sec, batch_size=16) else: for sec in secs: fad_streams[sec] = StreamingFAD(base_fad, pad_seconds=args.fad_pad_sec, batch_size=16) mel_tf = MelScale(n_mels=80, sample_rate=16000, n_stft=372) for batch_start in tqdm(range(0, len(eps), args.batch_size), total=(len(eps) + args.batch_size - 1) // args.batch_size, disable=(rank != 0)): batch_eps = eps[batch_start:batch_start + args.batch_size] # per-sec containers (vision) gt_img_batch, exp_img_batch = {}, {} gt_img_paths_batch, exp_img_paths_batch = {}, {} denorm_pairs_by_sec = {} secs_py = [int(s) for s in secs] denorm_pairs_by_sec = {s: [] for s in secs_py} for sec in secs: gt_img_batch[sec], exp_img_batch[sec] = [], [] gt_img_paths_batch[sec], exp_img_paths_batch[sec] = [], [] # per-sec containers (audio paths) gt_wav_paths_batch, exp_wav_paths_batch = {}, {} for sec in secs: gt_wav_paths_batch[sec], exp_wav_paths_batch[sec] = [], [] for ep in batch_eps: gt_ep_dir = os.path.join(gt_dir, ep) exp_ep_dir = os.path.join(exp_dir, ep) if (not os.path.isdir(gt_ep_dir)) or (not os.path.isdir(exp_ep_dir)): continue gt_dist_p = os.path.join(gt_ep_dir, "distance.json") exp_dist_p = os.path.join(exp_ep_dir, "distance.json") try: if os.path.isfile(gt_dist_p) and os.path.isfile(exp_dist_p): with open(gt_dist_p, "r") as f: gt_list = json.load(f) with open(exp_dist_p, "r") as f: exp_list = json.load(f) gt_map = {int(it["sec"]): float(it["denorm_gt"]) for it in gt_list if "sec" in it and "denorm_gt" in it} exp_map = {int(it["sec"]): float(it["denorm_pred"]) for it in exp_list if "sec" in it and "denorm_pred" in it} for s in secs_py: if s in gt_map and s in exp_map: denorm_pairs_by_sec[s].append((gt_map[s], exp_map[s])) except Exception: pass for sec, image_idx in zip(secs, image_idxs): # ---- vision if 'v' in modals: gt_sec_img_path = os.path.join(gt_ep_dir, f'{int(image_idx)}.png') exp_sec_img_path = os.path.join(exp_ep_dir, f'{int(image_idx)}.png') if os.path.isfile(gt_sec_img_path) and os.path.isfile(exp_sec_img_path): try: gt_img = to_tensor(Image.open(gt_sec_img_path).convert("RGB")).unsqueeze(0).to(device) exp_img = to_tensor(Image.open(exp_sec_img_path).convert("RGB")).unsqueeze(0).to(device) if torch.isfinite(gt_img).all() and torch.isfinite(exp_img).all(): gt_img_batch[sec].append(gt_img) exp_img_batch[sec].append(exp_img) gt_img_paths_batch[sec].append(gt_sec_img_path) exp_img_paths_batch[sec].append(exp_sec_img_path) except Exception: pass # ---- audio if 'a' in modals: gt_sec_wav_path = os.path.join(gt_ep_dir, f'{int(image_idx)}.wav') exp_sec_wav_path = os.path.join(exp_ep_dir, f'{int(image_idx)}.wav') if os.path.isfile(gt_sec_wav_path) and os.path.isfile(exp_sec_wav_path): gt_wav_paths_batch[sec].append(gt_sec_wav_path) exp_wav_paths_batch[sec].append(exp_sec_wav_path) # ---- vision metric update per batch if 'v' in modals: for sec in secs: if (len(gt_img_batch[sec]) == 0) or (len(exp_img_batch[sec]) == 0): continue lpips_dists = lpips_loss_fn(gt_img_paths_batch[sec], exp_img_paths_batch[sec]) dreamsim_dists = dreamsim_loss_fn(gt_img_paths_batch[sec], exp_img_paths_batch[sec]) metric_logger.meters[f'{dataset_name}_{eval_name}_lpips_{sec}'].update(lpips_dists, n=1) metric_logger.meters[f'{dataset_name}_{eval_name}_dreamsim_{sec}'].update(dreamsim_dists, n=1) sec_gt_batch = torch.cat(gt_img_batch[sec], dim=0) sec_exp_batch = torch.cat(exp_img_batch[sec], dim=0) if torch.isfinite(sec_gt_batch).all() and torch.isfinite(sec_exp_batch).all(): fid_loss_fn[sec].update(images=sec_gt_batch, is_real=True) fid_loss_fn[sec].update(images=sec_exp_batch, is_real=False) psnr_vals = _psnr_from_tensors(sec_gt_batch, sec_exp_batch, data_range=1.0) # (N,) metric_logger.meters[f'{dataset_name}_{eval_name}_psnr_{sec}'].update(psnr_vals.mean(), n=1) # ---- audio metrics per batch if 'a' in modals: # FAD (streaming) if len(fad_streams) > 0: for sec in secs: if len(gt_wav_paths_batch[sec]) == 0 and len(exp_wav_paths_batch[sec]) == 0: continue fad_streams[sec].update_from_wavs(gt_wav_paths_batch[sec], is_real=True) fad_streams[sec].update_from_wavs(exp_wav_paths_batch[sec], is_real=False) # LSD / SSIM / MelCos / dRMS-db _AUDIO_SR = 16000 for sec in secs: gt_list = gt_wav_paths_batch[sec] exp_list = exp_wav_paths_batch[sec] if len(gt_list) == 0 or len(exp_list) == 0: continue pair_cnt = min(len(gt_list), len(exp_list)) if pair_cnt == 0: continue lsd_L, lsd_R, ssim_L, ssim_R = [], [], [], [] mel_L, mel_R = [], [] mel_lsd_L, mel_lsd_R = [], [] mel_ssim_L, mel_ssim_R = [], [] sispec_nl_L, sispec_nl_R = [], [] sispec_log_L, sispec_log_R = [], [] mel_sispec_nl_L, mel_sispec_n_R = [], [] mel_sispec_log_L, mel_sispec_log_R = [], [] for i in range(pair_cnt): gpath = gt_list[i] epath = exp_list[i] try: g_st = _load_librosa_stereo(gpath, _AUDIO_SR) # (2,T) e_st = _load_librosa_stereo(epath, _AUDIO_SR) # (2,T) if args.mono: g_mono = g_st.mean(axis=0) e_mono = e_st.mean(axis=0) # LSD/SSIM gt_sp = _wav_to_spectrogram(g_mono, rate=_AUDIO_SR) ex_sp = _wav_to_spectrogram(e_mono, rate=_AUDIO_SR) lsd_val = _lsd_from_specs(ex_sp.clone(), gt_sp.clone()) ssim_val = _ssim_from_specs(ex_sp.clone(), gt_sp.clone()) # MelCos mel_val = _mel_cosine_single_channel(e_mono, g_mono, _AUDIO_SR, mel_tf) # mel_lsd & mel_ssim mel_lsd_val, mel_ssim_val = _mel_lsd_ssim_single(e_mono, g_mono, mel_tf) # sispec sispec_nl = _sispec_from_specs(ex_sp.clone(), gt_sp.clone(), log_domain=False) sispec_log = _sispec_from_specs(ex_sp.clone(), gt_sp.clone(), log_domain=True) # Mel sispec mel_sispec_nl = _sispec_from_specs(ex_m.clone(), gt_m.clone(), log_domain=False) mel_sispec_log = _sispec_from_specs(ex_m.clone(), gt_m.clone(), log_domain=True) metric_logger.meters[f'{dataset_name}_{eval_name}_lsd_{sec}'].update(lsd_val, n=1) metric_logger.meters[f'{dataset_name}_{eval_name}_ssim_{sec}'].update( torch.tensor(ssim_val), n=1 ) metric_logger.meters[f'{dataset_name}_{eval_name}_melcos_{sec}'].update( torch.tensor(mel_val), n=1 ) metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_lsd_{sec}'].update( torch.tensor(float(mel_lsd_val)), n=1 ) metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_ssim_{sec}'].update( torch.tensor(float(mel_ssim_val)), n=1 ) metric_logger.meters[f'{dataset_name}_{eval_name}_non_log_sispec_{sec}'].update( torch.tensor(float(sispec_nl)), n=1 ) metric_logger.meters[f'{dataset_name}_{eval_name}_sispec_{sec}'].update( torch.tensor(float(sispec_log)), n=1 ) metric_logger.meters[f'{dataset_name}_{eval_name}_final_non_log_mel_sispec_{sec}'].update( torch.tensor(float(mel_sispec_nl)), n=1 ) metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_sispec_{sec}'].update( torch.tensor(float(mel_sispec_log)), n=1 ) else: for ch, (acc_lsd, acc_ssim, acc_mel, acc_mel_lsd, acc_mel_ssim, acc_sispec_nl, acc_sispec_log, acc_mel_sispec_nl, acc_mel_sispec_log) in enumerate([ (lsd_L, ssim_L, mel_L, mel_lsd_L, mel_ssim_L, sispec_nl_L, sispec_log_L, mel_sispec_nl_L, mel_sispec_log_L), (lsd_R, ssim_R, mel_R, mel_lsd_R, mel_ssim_R, sispec_nl_R, sispec_log_R, mel_sispec_n_R, mel_sispec_log_R), ]): g = g_st[ch]; e = e_st[ch] # LSD/SSIM gt_sp = _wav_to_spectrogram(g, rate=_AUDIO_SR) ex_sp = _wav_to_spectrogram(e, rate=_AUDIO_SR) acc_lsd.append(float(_lsd_from_specs(ex_sp.clone(), gt_sp.clone()))) acc_ssim.append(float(_ssim_from_specs(ex_sp.clone(), gt_sp.clone()))) # MelCos acc_mel.append(_mel_cosine_single_channel(e, g, _AUDIO_SR, mel_tf)) # mel_lsd & mel_ssim mel_lsd_val, mel_ssim_val = _mel_lsd_ssim_single(e, g, mel_tf) acc_mel_lsd.append(mel_lsd_val) acc_mel_ssim.append(mel_ssim_val) # sispec acc_sispec_nl.append( float(_sispec_from_specs(ex_sp.clone(), gt_sp.clone(), log_domain=False)) ) acc_sispec_log.append( float(_sispec_from_specs(ex_sp.clone(), gt_sp.clone(), log_domain=True)) ) # Mel est_mag = np.abs(librosa.stft(e, n_fft=743, hop_length=160)) ref_mag = np.abs(librosa.stft(g, n_fft=743, hop_length=160)) est_mel = mel_tf(torch.from_numpy(est_mag).float()) # [M,T] ref_mel = mel_tf(torch.from_numpy(ref_mag).float()) # [M,T] ex_m = est_mel.T.unsqueeze(0).unsqueeze(0) # [1,1,T,M] gt_m = ref_mel.T.unsqueeze(0).unsqueeze(0) # [1,1,T,M] # sispec(Mel, non_log / log) acc_mel_sispec_nl.append( float(_sispec_from_specs(ex_m.clone(), gt_m.clone(), log_domain=False)) ) acc_mel_sispec_log.append( float(_sispec_from_specs(ex_m.clone(), gt_m.clone(), log_domain=True)) ) except Exception: pass if not args.mono: def _maybe_mean(x): return float(np.mean(x)) if len(x) > 0 else None v = _maybe_mean(lsd_L); w = _maybe_mean(lsd_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_lsdL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_lsdR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_lsd_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) v = _maybe_mean(ssim_L); w = _maybe_mean(ssim_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_ssimL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_ssimR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_ssim_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) v = _maybe_mean(mel_L); w = _maybe_mean(mel_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_melcosL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_melcosR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_melcos_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) v = _maybe_mean(mel_lsd_L); w = _maybe_mean(mel_lsd_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_lsdL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_lsdR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_lsd_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) v = _maybe_mean(mel_ssim_L); w = _maybe_mean(mel_ssim_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_ssimL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_ssimR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_ssim_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) v = _maybe_mean(sispec_nl_L); w = _maybe_mean(sispec_nl_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_non_log_sispecL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_non_log_sispecR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_non_log_sispec_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) v = _maybe_mean(sispec_log_L); w = _maybe_mean(sispec_log_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_sispecL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_sispecR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_sispec_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) v = _maybe_mean(mel_sispec_nl_L); w = _maybe_mean(mel_sispec_n_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_non_log_mel_sispecL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_non_log_mel_sispecR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_non_log_mel_sispec_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) v = _maybe_mean(mel_sispec_log_L); w = _maybe_mean(mel_sispec_log_R) if v is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_sispecL_{sec}'].update(torch.tensor(v), n=1) if w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_sispecR_{sec}'].update(torch.tensor(w), n=1) if v is not None and w is not None: metric_logger.meters[f'{dataset_name}_{eval_name}_final_mel_sispec_{sec}'].update(torch.tensor(0.5*(v+w)), n=1) for s in secs_py: pairs = denorm_pairs_by_sec[s] if not pairs: continue arr = np.asarray(pairs, dtype=np.float32) mask = np.isfinite(arr).all(axis=1) if not np.any(mask): continue se_mean = float(np.mean((arr[mask, 1] - arr[mask, 0]) ** 2)) metric_logger.meters[f'{dataset_name}_{eval_name}_denorm_mse_{s}'].update( torch.tensor(se_mean), n=1 ) if 'v' in modals: feature_dim = 2048 sec_list = [int(s) for s in secs] tmp_dir = Path(os.path.join(args.exp_dir, ".fid_tmp")) if dist_torch.is_initialized(): if dist_torch.get_rank() == 0: tmp_dir.mkdir(parents=True, exist_ok=True) dist_torch.barrier() else: tmp_dir.mkdir(parents=True, exist_ok=True) if dist_torch.is_initialized(): my_rank = dist_torch.get_rank() world_size = dist_torch.get_world_size() else: my_rank = 0 world_size = 1 for s in sec_list: fid_m = fid_loss_fn[s] state = { "real_sum": fid_m.real_sum.detach().to("cpu", torch.float64), "real_cov_sum": fid_m.real_cov_sum.detach().to("cpu", torch.float64), "fake_sum": fid_m.fake_sum.detach().to("cpu", torch.float64), "fake_cov_sum": fid_m.fake_cov_sum.detach().to("cpu", torch.float64), "num_real_images": torch.tensor(int(fid_m.num_real_images.item()), dtype=torch.int64), "num_fake_images": torch.tensor(int(fid_m.num_fake_images.item()), dtype=torch.int64), } out_path = tmp_dir / f"fid_sec{s}_rank{my_rank}.pt" torch.save(state, out_path) if dist_torch.is_initialized(): dist_torch.barrier() if (not dist_torch.is_initialized()) or my_rank == 0: for s in sec_list: agg = { "real_sum": torch.zeros(feature_dim, dtype=torch.float64), "real_cov_sum": torch.zeros((feature_dim, feature_dim), dtype=torch.float64), "fake_sum": torch.zeros(feature_dim, dtype=torch.float64), "fake_cov_sum": torch.zeros((feature_dim, feature_dim), dtype=torch.float64), "num_real_images": torch.tensor(0, dtype=torch.int64), "num_fake_images": torch.tensor(0, dtype=torch.int64), } for r in range(world_size): p = tmp_dir / f"fid_sec{s}_rank{r}.pt" if not p.exists(): continue st = torch.load(p, map_location="cpu") agg["real_sum"] += st["real_sum"] agg["real_cov_sum"] += st["real_cov_sum"] agg["fake_sum"] += st["fake_sum"] agg["fake_cov_sum"] += st["fake_cov_sum"] agg["num_real_images"] += st["num_real_images"] agg["num_fake_images"] += st["num_fake_images"] fid_m = fid_loss_fn[s] fid_m.real_sum = agg["real_sum"].to(fid_m.device, fid_m.real_sum.dtype) fid_m.real_cov_sum = agg["real_cov_sum"].to(fid_m.device, fid_m.real_cov_sum.dtype) fid_m.fake_sum = agg["fake_sum"].to(fid_m.device, fid_m.fake_sum.dtype) fid_m.fake_cov_sum = agg["fake_cov_sum"].to(fid_m.device, fid_m.fake_cov_sum.dtype) fid_m.num_real_images = torch.tensor( int(agg["num_real_images"].item()), device=fid_m.device, dtype=fid_m.num_real_images.dtype ) fid_m.num_fake_images = torch.tensor( int(agg["num_fake_images"].item()), device=fid_m.device, dtype=fid_m.num_fake_images.dtype ) try: val = float(fid_m.compute().item()) metric_logger.meters[f'{dataset_name}_{eval_name}_fid_{s}'].update(val, n=1) except Exception as e: print(f"[WARN] FID compute failed at sec={s}: {e}") for s in sec_list: for r in range(world_size): p = tmp_dir / f"fid_sec{s}_rank{r}.pt" try: if p.exists(): p.unlink() except Exception: pass try: tmp_dir.rmdir() except Exception: pass if dist_torch.is_initialized(): dist_torch.barrier() if 'a' in modals and len(fad_streams) > 0: for sec in secs: try: if stereo_mode: fad_L, fad_R, fad_avg = fad_streams[sec].compute() metric_logger.meters[f'{dataset_name}_{eval_name}_fadL_{sec}'].update(fad_L, n=1) metric_logger.meters[f'{dataset_name}_{eval_name}_fadR_{sec}'].update(fad_R, n=1) metric_logger.meters[f'{dataset_name}_{eval_name}_fad_{sec}'].update(fad_avg, n=1) else: fad_val = float(fad_streams[sec].compute()) metric_logger.meters[f'{dataset_name}_{eval_name}_fad_{sec}'].update(fad_val, n=1) except Exception as e: if rank == 0: print(f"[WARN] FAD compute failed at sec={sec}: {e}") continue # ----------------------------- # Save # ----------------------------- def save_metric_to_disk(metric_logger, log_p, rank): if dist_torch.is_initialized(): metric_logger.synchronize_between_processes() if rank == 0: log_stats = {k: float(meter.global_avg) for k, meter in metric_logger.meters.items()} os.makedirs(os.path.dirname(log_p), exist_ok=True) with open(log_p, 'w') as json_file: json.dump(log_stats, json_file, indent=4) print(f"[OK] Metrics saved to: {log_p}") # ----------------------------- # Main # ----------------------------- def main(args): rank, world_size, local_rank = setup_distributed() device = f"cuda:{local_rank}" if world_size > 1 else ("cuda" if torch.cuda.is_available() else "cpu") torch.backends.cudnn.benchmark = True dataset_name = args.dataset secs = np.array([i for i in range(1, 17)], dtype=int) # vision metrics (will only be used if 'v' in modals) lpips_loss_fn = get_loss_fn('lpips', secs, device) dreamsim_loss_fn = get_loss_fn('dreamsim', secs, device) fid_metrics_vision = get_loss_fn('fid', secs, device) try: metric_logger = dist.MetricLogger(delimiter=" ") if rank == 0: print(f"Evaluating {args.eval_name} {dataset_name} | modals = {args.modals}") time_loss_fns = (lpips_loss_fn, dreamsim_loss_fn, fid_metrics_vision) with torch.no_grad(): evaluate( args=args, dataset_name=dataset_name, eval_type=args.eval_name, metric_logger=metric_logger, loss_fns=time_loss_fns, gt_dir=args.gt_dir, exp_dir=args.exp_dir, secs=secs, device=device, rank=rank, world_size=world_size, modals=args.modals ) output_fn = os.path.join(args.exp_dir, f'{dataset_name}_{args.eval_name}.json') save_metric_to_disk(metric_logger, output_fn, rank) except Exception as e: if rank == 0: print(e) finally: if dist_torch.is_initialized(): dist_torch.barrier() dist_torch.destroy_process_group() # ----------------------------- # CLI # ----------------------------- if __name__ == "__main__": parser = argparse.ArgumentParser(allow_abbrev=False) parser.add_argument("--batch_size", type=int, default=64, help="batch size") parser.add_argument("--gt_dir", type=str, required=True, help="gt directory") parser.add_argument("--exp_dir", type=str, required=True, help="experiment directory (also save json here)") parser.add_argument("--eval_name", type=str, default='time', choices=['time', 'rollout'], help="eval type") parser.add_argument("--dataset", type=str, required=True, help="dataset name (for metric keys & json name)") parser.add_argument("--modals", type=str, default="av", choices=["a", "v", "av"], help="a=audio only (wav), v= image only (png), av=both") # FAD options parser.add_argument("--fad_model", type=str, default="vggish", choices=["vggish", "pann", "clap", "encodec"], help="embedding model for FAD") parser.add_argument("--fad_sr", type=int, default=16000, help="sampling rate for FAD") # Stereo VGGish FAD options parser.add_argument("--mono", action="store_true", help="default as stereo, add --mono to mono") parser.add_argument("--fad_pad_sec", type=float, default=1.0, help="pad the input of VGGish to x seconds") args = parser.parse_args() main(args)