added flag to change image size in inference; added gradio

Author: codinglabsong

app.py

@@ -0,0 +1,89 @@
+"""
+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 aging_gan.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/epoch_0030.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()

requirements.txt

@@ -10,4 +10,5 @@ torchmetrics[image]
 wandb
 numpy
 python-dotenv
-boto3
+boto3
+gradio

setup.py

@@ -9,7 +9,7 @@
         "python-dotenv",
         "numpy",
         "wandb",
-        # "gradio",
+        "gradio",
         "ipykernel",
         "matplotlib",
         "torch",

src/aging_gan/inference.py

@@ -36,28 +36,37 @@ def parse_args() -> argparse.Namespace:
         default="young2old",
         help="'young2old' uses generator G, 'old2young' uses generator F",
     )
+    p.add_argument(
+        "--train_img_size",
+        type=int,
+        default=256,
+        help="The same img_size you used for training and evaluating.",
+    )
     return p.parse_args()
 
 
-# image helpers
-preprocess = T.Compose(
-    [
-        T.Resize(256, interpolation=Image.BICUBIC),
-        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()])
-
-
 @torch.inference_mode()
 def main() -> None:
     """Load a checkpoint and generate an aged face from ``--input``."""
     cfg = parse_args()
     device = get_device()
 
+    # image helpers
+    preprocess = T.Compose(
+        [
+            T.Resize(
+                (cfg.train_img_size + 50, cfg.train_img_size + 50), antialias=True
+            ),
+            T.CenterCrop(cfg.train_img_size),
+            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 generators and checkpoint
     G, F, *_ = initialize_models()  # returns G, F, DX, DY
     ckpt = torch.load(cfg.ckpt, map_location=device)  # same keys as used in train.py