| """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) | |