File size: 1,052 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 | """Minimal round trip: image -> latents -> image."""
import sys
import numpy as np
import torch
from PIL import Image
from ae import AutoEncoder
src = sys.argv[1] if len(sys.argv) > 1 else "input.png"
dst = sys.argv[2] if len(sys.argv) > 2 else "recon.png"
model = AutoEncoder.from_pretrained("Muinez/f8c32-ae") # HF repo id or local dir
img = Image.open(src).convert("RGB")
w, h = (img.width // 8) * 8, (img.height // 8) * 8 # sides must be multiples of 8
img = img.crop((0, 0, w, h))
x = torch.from_numpy(np.asarray(img, dtype=np.float32)).permute(2, 0, 1)[None] / 127.5 - 1
z = model.encode(x) # (1, 32, h/8, w/8), normalized
print("latent", tuple(z.shape), f"mean {z.mean():.3f} std {z.std():.3f}")
gen = torch.Generator(model.device).manual_seed(0)
y = model.decode(z, steps=4, generator=gen).clamp(-1, 1)
mse = (y.cpu() - x).pow(2).mean()
print(f"PSNR {10 * torch.log10(4 / mse):.2f} dB")
Image.fromarray(((y[0] + 1) * 127.5).round().byte().permute(1, 2, 0).cpu().numpy()).save(dst)
print("saved", dst)
|