Upload eval_fad.py with huggingface_hub

eval_fad.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+
+import numpy as np
+import soundfile as sf
+import torch
+from diffusers.pipelines.deprecated.audio_diffusion.mel import Mel
+from frechet_audio_distance import FrechetAudioDistance
+
+from sample import generate, image_from_tensor, load_clap, load_model
+
+PANN_EMBED_DIM = 2048
+
+def fad_score(frechet, background_dir, eval_dir):
+    audio_bg = frechet._FrechetAudioDistance__load_audio_files(background_dir, dtype="float32")
+    embds_bg = frechet.get_embeddings(audio_bg, sr=frechet.sample_rate).reshape(-1, PANN_EMBED_DIM)
+    audio_ev = frechet._FrechetAudioDistance__load_audio_files(eval_dir, dtype="float32")
+    embds_ev = frechet.get_embeddings(audio_ev, sr=frechet.sample_rate).reshape(-1, PANN_EMBED_DIM)
+    mu1, sigma1 = frechet.calculate_embd_statistics(embds_bg)
+    mu2, sigma2 = frechet.calculate_embd_statistics(embds_ev)
+    return frechet.calculate_frechet_distance(mu1, sigma1, mu2, sigma2)
+
+def first_caption_per_clip(meta):
+    seen, order = {}, []
+    for pair_idx, ci in enumerate(meta["pair_clip_idx"]):
+        if ci not in seen:
+            seen[ci] = meta["captions"][pair_idx]
+            order.append(ci)
+    return order, seen
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default="/root/data")
+    ap.add_argument("--audio-dir", default="/root/clotho_raw/evaluation")
+    ap.add_argument("--weights", default="/root/runs/audio_v1/model_best.safetensors")
+    ap.add_argument("--config", default="/root/runs/audio_v1/config.json")
+    ap.add_argument("--clap", default="laion/clap-htsat-unfused")
+    ap.add_argument("--n-eval", type=int, default=300)
+    ap.add_argument("--cfg", type=float, default=4.0)
+    ap.add_argument("--steps", type=int, default=50)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--work", default="/root/fad_work")
+    ap.add_argument("--out", default="/root/fad_results.json")
+    ap.add_argument("--pann-sr", type=int, default=32000)
+    args = ap.parse_args()
+    dev = "cuda"
+
+    meta = json.load(open(f"{args.data}/evaluation_meta.json"))
+    mel_arr = np.load(f"{args.data}/evaluation_mel.npy")
+    clip_order, caption_by_clip = first_caption_per_clip(meta)
+    rng = np.random.RandomState(args.seed)
+    n = min(args.n_eval, len(clip_order))
+    chosen = [clip_order[i] for i in rng.permutation(len(clip_order))[:n]]
+    prompts = [caption_by_clip[ci] for ci in chosen]
+    files = [meta["clip_files"][ci] for ci in chosen]
+    print(f"[fad] evaluating on {n} held-out clips from {args.audio_dir}", flush=True)
+
+    real_dir = os.path.join(args.work, "real")
+    gen_dir = os.path.join(args.work, "generated")
+    floor_dir = os.path.join(args.work, "griffinlim_floor")
+    uncond_dir = os.path.join(args.work, "no_prompt")
+    for d in (real_dir, gen_dir, floor_dir, uncond_dir):
+        shutil.rmtree(d, ignore_errors=True)
+        os.makedirs(d, exist_ok=True)
+
+    for fn in files:
+        shutil.copy(os.path.join(args.audio_dir, fn), os.path.join(real_dir, fn))
+
+    model, mel, _ = load_model(args.weights, args.config, dev)
+    enc = load_clap(args.clap, dev)
+    null_seq, null_pool = enc([""])
+
+    print("[fad] writing griffin-lim floor (real mel -> griffin-lim, isolates vocoder loss)", flush=True)
+    for ci, fn in zip(chosen, files):
+        img = image_from_tensor(torch.from_numpy(mel_arr[ci].astype(np.float32) / 127.5 - 1.0))
+        audio = mel.image_to_audio(img)
+        sf.write(os.path.join(floor_dir, fn), audio, mel.get_sample_rate())
+
+    print("[fad] generating from trained model", flush=True)
+    batch = 32
+    for i in range(0, n, batch):
+        p = prompts[i:i + batch]
+        seq, pool = enc(p)
+        x = generate(model, seq, pool, null_seq, null_pool, args.steps, args.cfg, dev)
+        for j, row in enumerate(x[:, 0]):
+            img = image_from_tensor(row)
+            audio = mel.image_to_audio(img)
+            sf.write(os.path.join(gen_dir, files[i + j]), audio, mel.get_sample_rate())
+        print(f"[fad] generated {min(i+batch,n)}/{n}", flush=True)
+
+    print("[fad] generating unconditional (no prompt) baseline", flush=True)
+    for i in range(0, n, batch):
+        m = min(batch, n - i)
+        ns = null_seq.expand(m, -1, -1)
+        npool = null_pool.expand(m, -1)
+        x = generate(model, ns, npool, null_seq, null_pool, args.steps, 1.0, dev)
+        for j, row in enumerate(x[:, 0]):
+            img = image_from_tensor(row)
+            audio = mel.image_to_audio(img)
+            sf.write(os.path.join(uncond_dir, files[i + j]), audio, mel.get_sample_rate())
+
+    frechet = FrechetAudioDistance(model_name="pann", sample_rate=args.pann_sr,
+                                    use_pca=False, use_activation=False, verbose=False)
+
+    results = {}
+    for name, d in [("generated_vs_real", gen_dir), ("griffinlim_floor_vs_real", floor_dir),
+                    ("no_prompt_vs_real", uncond_dir)]:
+        score = fad_score(frechet, real_dir, d)
+        results[name] = score
+        print(f"[fad] FAD {name:<26} = {score:.4f}", flush=True)
+
+    results["config"] = {"n_eval": n, "cfg": args.cfg, "steps": args.steps,
+                          "pann_sample_rate": args.pann_sr, "weights": args.weights}
+    json.dump(results, open(args.out, "w"), indent=2)
+    print(f"[fad] wrote {args.out}", flush=True)
+
+if __name__ == "__main__":
+    main()