| import argparse, json, os, sys, numpy as np, torch |
| from diffusers import AutoencoderKL |
| from transformers import CLIPTextModel, CLIPTokenizer, CLIPModel |
| sys.path.insert(0, '/root') |
| from dit import DiT |
|
|
| SCALE = 0.18215 |
|
|
| def _img_feats(self, pixel_values=None, **kw): |
| return self.visual_projection(self.vision_model(pixel_values=pixel_values).pooler_output) |
|
|
| def _txt_feats(self, input_ids=None, attention_mask=None, **kw): |
| return self.text_projection(self.text_model(input_ids=input_ids, |
| attention_mask=attention_mask).pooler_output) |
|
|
| CLIPModel.get_image_features = _img_feats |
| CLIPModel.get_text_features = _txt_feats |
|
|
| @torch.no_grad() |
| def sample(model, seq, pool, ns, npool, steps, cfg, dev): |
| B = seq.shape[0] |
| x = torch.randn(B, 4, 32, 32, device=dev) |
| a, b = ns.expand(B, -1, -1), npool.expand(B, -1) |
| dt = 1.0 / steps |
| for i in range(steps): |
| t = torch.full((B,), i * dt, device=dev) |
| with torch.autocast('cuda', dtype=torch.bfloat16): |
| vc = model(x, t, seq, pool) |
| vu = model(x, t, a, b) |
| x = x + (vu + cfg * (vc - vu)).float() * dt |
| return x |
|
|
| @torch.no_grad() |
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--ckpt', default='/root/runs/pm5/best.pt') |
| ap.add_argument('--data', default='/root/pm5eval/eval_256.npz') |
| ap.add_argument('--n', type=int, default=5000) |
| ap.add_argument('--batch', type=int, default=100) |
| ap.add_argument('--steps', type=int, default=50) |
| ap.add_argument('--cfgs', default='2,3,4,5,6') |
| ap.add_argument('--out', default='/root/pm5_eval.json') |
| args = ap.parse_args() |
| dev = 'cuda' |
|
|
| ck = torch.load(args.ckpt, map_location=dev) |
| model = DiT(dim=384, depth=12, heads=6).to(dev).eval() |
| model.load_state_dict(ck['ema']) |
| print(f"[eval] ckpt step {ck.get('step')} val {ck.get('val')}", flush=True) |
|
|
| vae = AutoencoderKL.from_pretrained('stabilityai/sd-vae-ft-mse').to(dev).half().eval() |
| tok = CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32') |
| txt = CLIPTextModel.from_pretrained('openai/clip-vit-base-patch32').to(dev).half().eval() |
| e = tok([''], padding='max_length', max_length=40, truncation=True, return_tensors='pt').to(dev) |
| o = txt(**e) |
| ns, npool = o.last_hidden_state.float(), o.pooler_output.float() |
|
|
| d = np.load(args.data, allow_pickle=True) |
| real = d['images'][:args.n] |
| caps = [str(x) for x in d['captions'][:args.n]] |
| n = len(caps) |
| print(f'[eval] {n} real images and captions', flush=True) |
|
|
| from torchmetrics.image.fid import FrechetInceptionDistance |
| from torchmetrics.multimodal.clip_score import CLIPScore |
|
|
| ref = CLIPScore(model_name_or_path='openai/clip-vit-base-patch32').to(dev) |
| for i in range(0, n, args.batch): |
| rb = torch.from_numpy(real[i:i + args.batch]).permute(0, 3, 1, 2).to(dev) |
| ref.update(rb, caps[i:i + args.batch]) |
| real_clip = float(ref.compute().item()) |
| print(f'[eval] REAL images CLIP score {real_clip:.2f} (metric sanity check / ceiling)', flush=True) |
| del ref |
| torch.cuda.empty_cache() |
|
|
| results = [] |
| for cfg in [float(x) for x in args.cfgs.split(',')]: |
| fid = FrechetInceptionDistance(feature=2048, normalize=True).to(dev) |
| clip = CLIPScore(model_name_or_path='openai/clip-vit-base-patch32').to(dev) |
| for i in range(0, n, args.batch): |
| rb = torch.from_numpy(real[i:i + args.batch].astype(np.float32) / 255.0).permute(0, 3, 1, 2).to(dev) |
| fid.update(rb, real=True) |
| for i in range(0, n, args.batch): |
| cb = caps[i:i + args.batch] |
| t = tok(cb, padding='max_length', max_length=40, truncation=True, return_tensors='pt').to(dev) |
| oo = txt(**t) |
| z = sample(model, oo.last_hidden_state.float(), oo.pooler_output.float(), |
| ns, npool, args.steps, cfg, dev) |
| img = vae.decode((z / SCALE).half()).sample.float() |
| img = (img.clamp(-1, 1) + 1) / 2 |
| fid.update(img, real=False) |
| clip.update((img * 255).to(torch.uint8), cb) |
| f = float(fid.compute().item()); c = float(clip.compute().item()) |
| results.append({'cfg': cfg, 'fid': round(f, 2), 'clip_score': round(c, 2), |
| 'n': n, 'steps': args.steps}) |
| print(f'[eval] cfg {cfg}: FID {f:.2f} CLIP {c:.2f}', flush=True) |
| del fid, clip |
| torch.cuda.empty_cache() |
|
|
| best = min(results, key=lambda r: r['fid']) |
| json.dump({'ckpt_step': ck.get('step'), 'val_loss': ck.get('val'), |
| 'real_image_clip_score': round(real_clip, 2), |
| 'results': results, 'best_fid': best}, open(args.out, 'w'), indent=2) |
| print(f"[eval] BEST FID {best['fid']} at cfg {best['cfg']} (CLIP {best['clip_score']})", flush=True) |
| print('EVALDONE', flush=True) |
|
|
| if __name__ == '__main__': |
| main() |
|
|
