Spaces:
Runtime error
Runtime error
File size: 2,315 Bytes
0587b57 b124220 0587b57 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """
Gradio demo for Aging-GAN: upload a face, choose direction, and get an aged or rejuvenated output.
"""
import gradio as gr
import torch
from pathlib import Path
from PIL import Image
import torchvision.transforms as T
from model import initialize_models
# Utils
def get_device() -> torch.device:
"""Return CUDA device if available else CPU."""
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Transforms
preprocess = T.Compose(
[
T.Resize((256 + 50, 256 + 50), antialias=True),
T.CenterCrop(256),
T.ToTensor(),
T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
)
postprocess = T.Compose([T.Normalize(mean=[-1, -1, -1], std=[2, 2, 2]), T.ToPILImage()])
# Load models & checkpoint once
device = get_device()
# initialize G (young→old) and F (old→young)
G, F, _, _ = initialize_models()
ckpt_path = Path("outputs/checkpoints/best.pth")
ckpt = torch.load(ckpt_path, map_location=device)
G.load_state_dict(ckpt["G"])
F.load_state_dict(ckpt["F"])
G.eval().to(device)
F.eval().to(device)
# Inference function
def infer(image: Image.Image, direction: str) -> Image.Image:
"""
Run a single forward pass through the chosen generator.
"""
# preprocess
x = preprocess(image).unsqueeze(0).to(device) # (1,3,256,256)
# generate
with torch.inference_mode():
if direction == "young2old":
y_hat = G(x)
else:
y_hat = F(x)
y_hat = torch.clamp(y_hat, -1, 1)
# postprocess & return PIL image
out = postprocess(y_hat.squeeze(0).cpu())
return out
# Launch Gradio
demo = gr.Interface(
fn=infer,
inputs=[
gr.Image(type="pil", label="Input Face"),
gr.Radio(
choices=["young2old", "old2young"],
value="young2old",
label="Transformation Direction",
),
],
outputs=gr.Image(type="pil", label="Output Face"),
title="Aging-GAN Demo",
description=(
"Upload a portrait, select “young2old” to age it or “old2young” to rejuvenate. "
"Powered by a ResNet-style CycleGAN generator. "
"TIP: Upload close-up photos of the face similar to ones in the Github README examples."
),
allow_flagging="never",
)
if __name__ == "__main__":
demo.launch()
|