File size: 2,119 Bytes
63a1291 | 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 | """PixelModel v3 inference from model.safetensors (standalone).
Needs only: model.safetensors, model.py, vocab.json, torch. No model.png,
no config.json - the architecture config is read from the safetensors header.
Produces output identical to main.py for the same prompt and resolution.
python convert_to_safetensors.py
python INFERENCE.py "a red double decker bus" --out bus.png
"""
from __future__ import annotations
import argparse
import json
import numpy as np
import torch
from PIL import Image
from safetensors import safe_open
from model import (
ModelConfig, PixelModelV3, load_vocab, encode_caption, make_coord_grid,
)
def load_from_safetensors(path, device):
with safe_open(path, framework="pt", device="cpu") as f:
meta = f.metadata() or {}
cfg = ModelConfig.from_json(json.loads(meta["config"]))
model = PixelModelV3(cfg)
state = {k: f.get_tensor(k) for k in f.keys()}
model.load_state_dict(state)
model.to(device).eval()
return model, cfg
def main():
ap = argparse.ArgumentParser()
ap.add_argument("prompt")
ap.add_argument("--out", default="out.png")
ap.add_argument("--res", type=int, default=128)
ap.add_argument("--safetensors", default="model.safetensors")
ap.add_argument("--vocab", default="vocab.json")
ap.add_argument("--device", default="cpu")
args = ap.parse_args()
model, cfg = load_from_safetensors(args.safetensors, args.device)
vocab = load_vocab(args.vocab)
tokens = encode_caption(args.prompt, vocab, cfg.max_tokens)
tokens = torch.from_numpy(tokens).long().unsqueeze(0).to(args.device)
coords = make_coord_grid(args.res, args.res, device=args.device,
dtype=torch.float32).unsqueeze(0)
with torch.no_grad():
rgb = model(tokens, coords)
img = (rgb.clamp(0, 1).reshape(args.res, args.res, 3).cpu().numpy() * 255.0
).round().astype(np.uint8)
Image.fromarray(img, "RGB").save(args.out)
print(f'[inference] "{args.prompt}" @ {args.res}x{args.res} -> {args.out}')
if __name__ == "__main__":
main()
|