File size: 3,171 Bytes
b4f5f26 | 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 | #!/usr/bin/env python3
"""Сэмпл из 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 # noqa: E402
from text import EttinText # noqa: E402
from ae.model import AutoEncoder # noqa: E402
from export import load_exported # noqa: E402
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()
|