| |
| """Сэмпл из transformer/ (или из .pt чекпойнта обучения). |
| |
| python3 src/sample.py --prompt "A photograph of a red fox in fresh snow." --out fox.png |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from PIL import Image |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| sys.path[:0] = [str(ROOT / "src"), str(ROOT / "vae")] |
|
|
| from model import DEV, UDiT |
| from text import EttinText |
| from ae.model import AutoEncoder |
| from export import load_exported |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default=str(ROOT / "transformer"), help="папка transformer/ или .pt") |
| ap.add_argument("--prompt", required=True) |
| ap.add_argument("--out", default=str(ROOT / "tmp" / "sample.png")) |
| ap.add_argument("--size", type=int, default=256, help="сторона кадра, кратна 64") |
| ap.add_argument("--steps", type=int, default=16, help="шагов тела") |
| ap.add_argument("--tail-steps", type=int, default=8, help="шагов декодера (хвоста)") |
| ap.add_argument("--cfg", type=float, default=2.5) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--n", type=int, default=1, help="сколько картинок подряд от --seed") |
| ap.add_argument("--text-len", type=int, default=192) |
| args = ap.parse_args() |
|
|
| ae = AutoEncoder.from_pretrained(str(ROOT / "vae"), device=DEV) |
| enc = EttinText(str(ROOT / "text_encoder"), DEV, max_len=args.text_len) |
| m = UDiT(ae, text_dim=enc.dim).to(DEV) |
|
|
| src = Path(args.model) |
| if src.is_dir(): |
| step = load_exported(m, enc, ae, src, DEV).get("step") |
| else: |
| sd = torch.load(src, map_location=DEV, weights_only=False) |
| own = m.body.state_dict() |
| m.body.load_state_dict({k: v for k, v in sd["body"].items() if k in own}, strict=False) |
| enc.mixer.load_state_dict(sd["mixer"]) |
| if sd.get("decoder") is not None: |
| ae.decoder.load_state_dict(sd["decoder"]) |
| step = sd.get("step") |
| m.eval() |
| print(f"[model] {src}, шаг {step}", flush=True) |
|
|
| null_hs, null_mk = enc.slices([""], pad_to_len=args.text_len) |
| null_t = enc.mix(null_hs.float()) |
| hs, mk = enc.slices([args.prompt], pad_to_len=args.text_len) |
| txt = enc.mix(hs.float()) |
|
|
| out = Path(args.out) |
| out.parent.mkdir(parents=True, exist_ok=True) |
| gh = gw = args.size // 8 |
| for i in range(args.n): |
| y = m.sample(txt, mk, gh, gw, steps=args.steps, tail_steps=args.tail_steps, cfg=args.cfg, |
| null_text=null_t, null_mask=null_mk, |
| generator=torch.Generator(DEV).manual_seed(args.seed + i)) |
| im = Image.fromarray(((y[0].clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0).cpu().numpy()) |
| path = out if args.n == 1 else out.with_name(f"{out.stem}_{args.seed + i}{out.suffix}") |
| im.save(path) |
| print(f"[ok] {path} {im.size[0]}x{im.size[1]}, seed {args.seed + i}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
