| """Generate images from a trained checkpoint. |
| |
| Run: python diffusion/sample.py [--n 64] [--out samples.png] [--seed 0] |
| """ |
| import argparse |
| import os |
|
|
| import torch |
| from torchvision.utils import save_image |
|
|
| from train_diffusion import UNet, Diffusion, T, DDIM_STEPS |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", default=os.path.join(os.path.dirname(__file__), "out_big", "checkpoint.pt")) |
| ap.add_argument("--base", type=int, default=128, help="must match the trained model") |
| ap.add_argument("--n", type=int, default=64, help="number of images") |
| ap.add_argument("--steps", type=int, default=DDIM_STEPS) |
| ap.add_argument("--seed", type=int, default=None) |
| ap.add_argument("--out", default=os.path.join(os.path.dirname(__file__), "generated.png")) |
| args = ap.parse_args() |
|
|
| if args.seed is not None: |
| torch.manual_seed(args.seed) |
|
|
| model = UNet(base=args.base).cuda().eval() |
| ck = torch.load(args.ckpt, map_location="cuda", weights_only=True) |
| model.load_state_dict(ck["ema"]) |
|
|
| imgs = Diffusion(T, "cuda").ddim_sample(model, args.n, args.steps, "cuda") |
| save_image(imgs * 0.5 + 0.5, args.out, nrow=int(args.n ** 0.5)) |
| print(f"saved {args.n} images to {args.out} (checkpoint epoch {ck['epoch']})") |
|
|
