| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from safetensors.torch import load_file |
| from diffusers import AutoencoderKL |
| from transformers import CLIPTextModel, CLIPTokenizer |
|
|
| from dit import DiT |
|
|
| SCALE = 0.18215 |
|
|
| @torch.no_grad() |
| def sample(model, seq, pool, null_seq, null_pool, steps, cfg, dev, seed=None): |
| B = seq.shape[0] |
| g = None |
| if seed is not None: |
| g = torch.Generator(device=dev).manual_seed(seed) |
| x = torch.randn(B, 4, 32, 32, device=dev, generator=g) |
| ns, npool = null_seq.expand(B, -1, -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, pool) |
| vu = model(x, t, ns, npool) |
| x = x + (vu + cfg * (vc - vu)).float() * dt |
| return x |
|
|
| @torch.no_grad() |
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("prompts", nargs="+") |
| ap.add_argument("--out", default="out.png") |
| ap.add_argument("--cfg", type=float, default=5.0) |
| ap.add_argument("--steps", type=int, default=50) |
| ap.add_argument("--seed", type=int, default=None) |
| ap.add_argument("--device", default="cuda") |
| ap.add_argument("--weights", default="model.safetensors") |
| ap.add_argument("--config", default="config.json") |
| ap.add_argument("--vae", default="stabilityai/sd-vae-ft-mse") |
| ap.add_argument("--clip", default="openai/clip-vit-base-patch32") |
| ap.add_argument("--max-tokens", type=int, default=40) |
| args = ap.parse_args() |
| dev = args.device |
|
|
| d = json.load(open(args.config))["dit"] if os.path.exists(args.config) else { |
| "dim": 384, "depth": 12, "heads": 6} |
| model = DiT(dim=d["dim"], depth=d["depth"], heads=d["heads"]).to(dev).eval() |
| sd = load_file(args.weights) |
| model.load_state_dict({k[len("dit."):]: v for k, v in sd.items() if k.startswith("dit.")}) |
|
|
| vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval() |
| tok = CLIPTokenizer.from_pretrained(args.clip) |
| txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval() |
|
|
| def enc(strings): |
| t = tok(strings, padding="max_length", max_length=args.max_tokens, |
| truncation=True, return_tensors="pt").to(dev) |
| o = txt(**t) |
| return o.last_hidden_state.float(), o.pooler_output.float() |
|
|
| seq, pool = enc(args.prompts) |
| null_seq, null_pool = enc([""]) |
| z = sample(model, seq, pool, null_seq, null_pool, args.steps, args.cfg, dev, args.seed) |
| img = vae.decode((z / SCALE).half()).sample.float() |
| img = ((img.clamp(-1, 1) + 1) / 2).permute(0, 2, 3, 1).cpu().numpy() |
|
|
| n = len(args.prompts) |
| if n == 1: |
| Image.fromarray((img[0] * 255).round().astype(np.uint8)).save(args.out) |
| paths = [args.out] |
| else: |
| root, ext = os.path.splitext(args.out) |
| paths = [] |
| for i in range(n): |
| p = f"{root}_{i}{ext}" |
| Image.fromarray((img[i] * 255).round().astype(np.uint8)).save(p) |
| paths.append(p) |
| for p, s in zip(paths, args.prompts): |
| print(f'[pixelmodel] "{s}" -> {p} (cfg {args.cfg}, {args.steps} steps)') |
|
|
| if __name__ == "__main__": |
| main() |
|
|
