"""Pure evaluation: no optimizers, training steps, or EMA updates.""" import argparse import json import hashlib import os from pathlib import Path import subprocess import sys import numpy as np import torch from torch.utils.data import DataLoader, Subset from torchvision.utils import save_image from tqdm import tqdm from dataset import ImageFolderDataset from omegaconf import OmegaConf import torch.nn.functional as F from models import Encoder, Decoder, ARModel from train import (seed_everything, patchify, img_unint8_to_norm, DiTiScheduler, reconstruction, sampling, euler_recon_loop, euler_sample_loop, img_norm_to_uint8, unpatchify) def state_dict(state): return {k.removeprefix('_orig_mod.'): v for k, v in state.items()} def load_models(config, checkpoint, device='cuda', ema=True, need_ar=True): state = torch.load(checkpoint, map_location='cpu', weights_only=True, mmap=True) enc, dec = Encoder(config), Decoder(config) enc.load_state_dict(state_dict(state['enc']), strict=True) dec.load_state_dict(state_dict(state['dec']), strict=True) ar = None if need_ar: ar = ARModel(config) key = 'ar_model_ema' if ema else 'ar_model' if state.get(key) is None: raise ValueError(f'Checkpoint has no {key}; select a matching AR checkpoint') ar.load_state_dict(state_dict(state[key]), strict=True) for model in (enc, dec, ar): if model is not None: model.to(device).eval().requires_grad_(False) return enc, dec, ar def scheduler(config): if config.get('time_stages') and config.get('k_per_stage'): return DiTiScheduler(1000, config.z_len, config.time_stages, config.k_per_stage, config.reverse_t) return None def labels_for(config, batch_size, device): if config.use_label: raise ValueError('The released CIFAR-10 protocols are unconditional') return F.one_hot(torch.full((batch_size,), config.num_classes-1, device=device, dtype=torch.long), config.num_classes).float() @torch.no_grad() def decode_batch(config, enc, dec, ar, x, labels, mode, diti=None): """x is a patchified batch; its shape also defines unconditional generation.""" t = torch.ones(x.shape[0], device=x.device) if config.use_noise_query else None # wARTok draws reconstruction noise before the generation branch, even for AR-only eval. noise_recon = torch.randn_like(x) if config.use_noise_query else None if mode == 'reconstruct': if config.use_noise_query and config.diffusion_decoder: result, _ = euler_recon_loop(enc, dec, x, labels, config=config, diti_scheduler=diti) else: result, _ = reconstruction(enc, dec, x, labels, query_tokens=noise_recon, t=t) if config.l2_v_target: result = noise_recon - result elif mode == 'generate': kwargs = dict(bz=x.shape[0], class_label=labels, temperature=config.temperature, topK=config.topK, topP=config.topP, cfg=config.cfg, cfg_schedule=config.cfg_schedule, cfg_power=config.cfg_power, cache_kv=config.cache_kv, config=config) if config.use_noise_query and config.diffusion_decoder: result = euler_sample_loop(enc, dec, ar, **kwargs, diti_scheduler=diti) else: noise = torch.randn_like(x) if config.use_noise_query else None result = sampling(enc, dec, ar, **kwargs, t=t, query_tokens=noise) if config.l2_v_target: result = noise - result else: raise ValueError(mode) return img_norm_to_uint8(unpatchify(result, config.image_size, config.patch_size)) def build_parser(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--config', required=True) parser.add_argument('--checkpoint', required=True) parser.add_argument('--mode', choices=['generate', 'reconstruct'], default='generate') parser.add_argument('--data', help='Required for reconstruction; use the CIFAR-10 training split') parser.add_argument('--output', required=True) parser.add_argument('--num-samples', type=int, default=50000) parser.add_argument('--batch-size', type=int, default=64, help='Per process batch size (original evaluation: 64)') parser.add_argument('--seed', type=int, default=42) parser.add_argument('--legacy', action=argparse.BooleanOptionalAction, default=True, help='Use the same seed on every rank (wARTok behavior); --no-legacy uses seed + rank') parser.add_argument('--device', default='cuda') parser.add_argument('--no-ema', action='store_true') parser.add_argument('--fid-reference') return parser def main(): parser = build_parser() args = parser.parse_args() if args.num_samples < 1 or args.batch_size < 1: parser.error('sample count and batch size must be positive') world = int(os.environ.get('WORLD_SIZE', 1)); rank = int(os.environ.get('RANK', 0)) if world > 1: if args.device != 'cuda': parser.error('distributed evaluation requires CUDA') torch.cuda.set_device(int(os.environ['LOCAL_RANK'])) # Ranks exchange files, not CUDA tensors; CPU barriers avoid requiring GPU P2P. torch.distributed.init_process_group('gloo') c = OmegaConf.load(args.config) torch.backends.cuda.matmul.allow_tf32 = bool(c.use_tf32) torch.backends.cudnn.allow_tf32 = bool(c.use_tf32) seed_everything(args.seed) enc, dec, ar = load_models(c, args.checkpoint, args.device, not args.no_ema, args.mode == 'generate') diti = scheduler(c) # Strided partition, without DistributedSampler padding: exactly N images globally. indices = list(range(rank, args.num_samples, world)) loader = None if args.mode == 'reconstruct': if not args.data: parser.error('--data is required for reconstruction') dataset = ImageFolderDataset(path=args.data, resolution=c.image_size, use_label=False, xflip=False, crop_type=c.eval_crop_type) if args.num_samples > len(dataset): parser.error('num-samples exceeds dataset size') loader = iter(DataLoader(Subset(dataset, indices), batch_size=args.batch_size, shuffle=False)) # A fixed evaluation seed is explicit; historical training-loop RNG was not checkpointed. seed_everything(args.seed + (0 if args.legacy else rank)) samples = [] with torch.inference_mode(): for start in tqdm(range(0, len(indices), args.batch_size), desc=args.mode, disable=rank != 0): n = min(args.batch_size, len(indices)-start) if loader is None: x = torch.empty(n, c.x_len, c.patch_dim, device=args.device) else: images, _ = next(loader) x = patchify(img_unint8_to_norm(images.to(args.device)), c.patch_size) labels = labels_for(c, n, args.device) images = decode_batch(c, enc, dec, ar, x, labels, args.mode, diti) samples.append(images.permute(0,2,3,1).cpu().numpy()) arr = np.concatenate(samples) if samples else np.empty((0,c.image_size,c.image_size,3),dtype=np.uint8) out = Path(args.output); out.mkdir(parents=True, exist_ok=True) if world > 1: np.savez(out/f'rank-{rank}.npz', images=arr, indices=np.asarray(indices)) torch.distributed.barrier() if rank == 0: if world > 1: arr = np.empty((args.num_samples,c.image_size,c.image_size,3),dtype=np.uint8) for r in range(world): with np.load(out/f'rank-{r}.npz') as shard: arr[shard['indices']]=shard['images'] assert len(arr) == args.num_samples path = out/'samples.npz'; np.savez_compressed(path, arr) save_image(torch.from_numpy(arr[:64]).permute(0,3,1,2).float()/255, out/'preview.png',nrow=8) report = dict(mode=args.mode, num_samples=len(arr), unique_images=len({hashlib.sha256(im.tobytes()).digest() for im in arr}), seed=args.seed, world_size=world, batch_size=args.batch_size, legacy=args.legacy, rank_seed_mode='shared' if args.legacy else 'independent', config=args.config, checkpoint=args.checkpoint, ar_ema=(not args.no_ema) if args.mode == 'generate' else None, decoder='euler25' if c.use_noise_query and c.diffusion_decoder else 'single_step', precision='32-true', fid=None, torch_version=torch.__version__, config_sha256=hashlib.sha256(Path(args.config).read_bytes()).hexdigest()) if args.fid_reference: report['reference_sha256']=hashlib.sha256(Path(args.fid_reference).read_bytes()).hexdigest() result = subprocess.run([sys.executable, str(Path(__file__).with_name('calc_fid.py')), '--ref_batch', args.fid_reference, '--sample_batch', str(path), '--batch_size','64'], text=True, check=True, stdout=subprocess.PIPE) print(result.stdout) report['fid']=float(next(line.split(':',1)[1] for line in result.stdout.splitlines() if line.startswith('FID_RESULT:'))) (out/'metrics.json').write_text(json.dumps(report,indent=2)+'\n') print(json.dumps(report, indent=2)) if world > 1: torch.distributed.barrier(); torch.distributed.destroy_process_group() if __name__ == '__main__': main()