File size: 4,020 Bytes
6c311ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | from __future__ import annotations
import argparse
import numpy as np
import torch
from PIL import Image
from diffusers import AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
from dit_v6 import MMDiT
PROMPTS = [
"a bowl of ramen with a soft boiled egg",
"a red fox sitting in a snowy forest",
"a lighthouse on a cliff at sunset",
"a wooden cabin in the mountains",
"a cup of coffee on a wooden table",
"a golden retriever running on a beach",
"a city street at night with neon signs",
"a bowl of fresh strawberries",
"a sailboat on a calm lake",
]
@torch.no_grad()
def sample(model, seq, mask, pool, null_seq, null_mask, null_pool, steps, cfg, dev):
B = seq.shape[0]
x = torch.randn(B, 4, 32, 32, device=dev)
ns, nm, npo = null_seq.expand(B, -1, -1), null_mask.expand(B, -1), null_pool.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, mask, pool)
vu = model(x, t, ns, nm, npo)
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/pm6/best.pt")
ap.add_argument("--vae", default="madebyollin/sdxl-vae-fp16-fix")
ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
ap.add_argument("--t5", default="google/flan-t5-base")
ap.add_argument("--steps", type=int, default=50)
ap.add_argument("--cfg", type=float, default=5.0)
ap.add_argument("--t5-len", type=int, default=32)
ap.add_argument("--clip-len", type=int, default=40)
ap.add_argument("--out", default="/root/preview.png")
args = ap.parse_args()
dev = "cuda"
ck = torch.load(args.ckpt, map_location=dev)
c = ck["cfg"]
model = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"],
t5_len=c["t5_len"]).to(dev).eval()
model.load_state_dict(ck["ema"])
print(f"[preview] loaded {args.ckpt} step {ck['step']}", flush=True)
vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval()
vae_scale = vae.config.scaling_factor
t5_tok = T5TokenizerFast.from_pretrained(args.t5)
t5 = T5EncoderModel.from_pretrained(args.t5).to(dev).eval()
clip_tok = CLIPTokenizer.from_pretrained(args.clip)
clip_txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval()
def enc(strings):
te = t5_tok(strings, padding="max_length", max_length=args.t5_len, truncation=True,
return_tensors="pt").to(dev)
seq = t5(input_ids=te["input_ids"], attention_mask=te["attention_mask"]).last_hidden_state.float()
ce = clip_tok(strings, padding="max_length", max_length=args.clip_len, truncation=True,
return_tensors="pt").to(dev)
pool = clip_txt(input_ids=ce["input_ids"]).pooler_output.float()
return seq, te["attention_mask"].float(), pool
null_seq, null_mask, null_pool = enc([""])
cell, pad, cols = 256, 8, 3
rows = (len(PROMPTS) + cols - 1) // cols
sheet = Image.new("RGB", (cols * cell + (cols + 1) * pad, rows * cell + (rows + 1) * pad), (245, 246, 248))
for i, prompt in enumerate(PROMPTS):
seq, mask, pool = enc([prompt])
z = sample(model, seq, mask, pool, null_seq, null_mask, null_pool, args.steps, args.cfg, dev)
img = vae.decode((z / vae_scale).half()).sample.float()
img = ((img.clamp(-1, 1) + 1) / 2)[0].permute(1, 2, 0).cpu().numpy()
a = (img * 255).round().astype(np.uint8)
r, cc = divmod(i, cols)
sheet.paste(Image.fromarray(a), (pad + cc * (cell + pad), pad + r * (cell + pad)))
print(f"[preview] {i+1}/{len(PROMPTS)}: {prompt}", flush=True)
sheet.save(args.out)
print(f"[preview] wrote {args.out}", flush=True)
if __name__ == "__main__":
main()
|