Update app.py

app.py

@@ -1,89 +1,89 @@
-# app.py — minimal Gradio app that loads your GAN and outputs a 10x10 image grid per click
-import os
-import numpy as np
-import torch
-from torchvision.utils import make_grid
-from PIL import Image
-import gradio as gr
-from scipy.stats import truncnorm
-
-# Your repo modules
-import archs
-import cfg
-
-# ---- tweak these if your filenames differ ----
-ARCH     = "arch_cifar10"               # e.g. "arch_cifar10"
-DATASET  = "celeba"                     # used to build ckpt path below
-CKPT_PATH = os.path.join("model", f"{DATASET}.pth")
-GENOTYPE_PATH = os.path.join("nas", "latest_G.npy")
-N_SAMPLES = 50                         # 5x10
-N_COLS    = 10
-
-def _to_pil_grid(t):
-    # t: (3,H,W) in [0,1] (we'll pass normalize=True to make_grid)
-    arr = (t.clamp(0,1) * 255).byte().permute(1,2,0).cpu().numpy()
-    return Image.fromarray(arr)
-
-# Cache the model so it loads only once
-_GEN = None
-_DEVICE = "cpu"
-_LATENT_DIM = 120  # fallback if not present in args
-
-def load_model_once():
-    global _GEN, _DEVICE, _LATENT_DIM
-    if _GEN is not None:
-        return _GEN, _DEVICE, _LATENT_DIM
-
-    args = cfg.parse_args()
-    # space may or may not have GPU; handle both
-    has_cuda = torch.cuda.is_available()
-    args.gpu_ids = [0] if has_cuda else []
-
-    _DEVICE = "cuda" if has_cuda else "cpu"
-
-    # 1) Load genotype
-    genotype_G = np.load(GENOTYPE_PATH, allow_pickle=True)
-
-    # 2) Build generator
-    G_base = eval(f"archs.{ARCH}.Generator")(args, genotype_G)
-    if has_cuda:
-        CKPT_PATH = os.path.join("model", f"{DATASET}.pth")
-        G = torch.nn.DataParallel(G_base, device_ids=args.gpu_ids).cuda(args.gpu_ids[0])
-    else:
-        CKPT_PATH = os.path.join("model", f"{DATASET}_cpu.pth")
-        G = G_base
-
-    # 3) Load checkpoint (expects {'gen_state_dict': ...})
-    ckpt = torch.load(CKPT_PATH, map_location=_DEVICE)
-    state = ckpt["gen_state_dict"] if "gen_state_dict" in ckpt else ckpt
-    G.load_state_dict(state)
-    G.eval()
-
-    # 4) Latent dim from args, fallback to common default
-    _LATENT_DIM = int(getattr(args, "latent_dim", 120))
-
-    _GEN = G
-    return _GEN, _DEVICE, _LATENT_DIM
-
-def generate_grid():
-    G, device, latent_dim = load_model_once()
-
-    with torch.no_grad():
-        z_np = truncnorm.rvs(-1, 1, loc=0, scale=1, size=(N_SAMPLES, latent_dim)).astype(np.float32)
-        z = torch.from_numpy(z_np).to(device)
-        imgs = G(z)  # (B,3,H,W) in [-1,1] or [0,1]
-
-        # make a grid; normalize handles [-1,1]
-        grid = make_grid(imgs, nrow=N_COLS, normalize=True, scale_each=True)
-        return _to_pil_grid(grid)
-
-# ---------------- Gradio UI ----------------
-with gr.Blocks(title="MMD-PMish-NAS-GAN Image Generation") as demo:
-    gr.Markdown("## MMD-PMish-NAS-GAN Image Generation\nClick **Generate** to sample a fresh 10×10 grid from your model.")
-    generate_btn = gr.Button("Generate", variant="primary")
-    out = gr.Image(label="Output Grid", type="pil")
-
-    generate_btn.click(fn=generate_grid, inputs=None, outputs=out, queue=True)
-
-if __name__ == "__main__":
-    demo.queue().launch()
+# app.py — minimal Gradio app that loads your GAN and outputs a 10x10 image grid per click
+import os
+import numpy as np
+import torch
+from torchvision.utils import make_grid
+from PIL import Image
+import gradio as gr
+from scipy.stats import truncnorm
+
+# Your repo modules
+import archs
+import cfg
+
+# ---- tweak these if your filenames differ ----
+ARCH     = "arch_cifar10"               # e.g. "arch_cifar10"
+DATASET  = "celeba"                     # used to build ckpt path below
+CKPT_PATH = os.path.join("model", f"{DATASET}.pth")
+GENOTYPE_PATH = os.path.join("nas", "latest_G.npy")
+N_SAMPLES = 50                         # 5x10
+N_COLS    = 10
+
+def _to_pil_grid(t):
+    # t: (3,H,W) in [0,1] (we'll pass normalize=True to make_grid)
+    arr = (t.clamp(0,1) * 255).byte().permute(1,2,0).cpu().numpy()
+    return Image.fromarray(arr)
+
+# Cache the model so it loads only once
+_GEN = None
+_DEVICE = "cpu"
+_LATENT_DIM = 120  # fallback if not present in args
+
+def load_model_once():
+    global _GEN, _DEVICE, _LATENT_DIM
+    if _GEN is not None:
+        return _GEN, _DEVICE, _LATENT_DIM
+
+    args = cfg.parse_args()
+    # space may or may not have GPU; handle both
+    has_cuda = torch.cuda.is_available()
+    args.gpu_ids = [0] if has_cuda else []
+
+    _DEVICE = "cuda" if has_cuda else "cpu"
+
+    # 1) Load genotype
+    genotype_G = np.load(GENOTYPE_PATH, allow_pickle=True)
+
+    # 2) Build generator
+    G_base = eval(f"archs.{ARCH}.Generator")(args, genotype_G)
+    if has_cuda:
+        CKPT_PATH = os.path.join("model", f"{DATASET}.pth")
+        G = torch.nn.DataParallel(G_base, device_ids=args.gpu_ids).cuda(args.gpu_ids[0])
+    else:
+        CKPT_PATH = os.path.join("model", f"{DATASET}_cpu.pth")
+        G = G_base
+
+    # 3) Load checkpoint (expects {'gen_state_dict': ...})
+    ckpt = torch.load(CKPT_PATH, map_location=_DEVICE)
+    state = ckpt["gen_state_dict"] if "gen_state_dict" in ckpt else ckpt
+    G.load_state_dict(state)
+    G.eval()
+
+    # 4) Latent dim from args, fallback to common default
+    _LATENT_DIM = int(getattr(args, "latent_dim", 120))
+
+    _GEN = G
+    return _GEN, _DEVICE, _LATENT_DIM
+
+def generate_grid():
+    G, device, latent_dim = load_model_once()
+
+    with torch.no_grad():
+        z_np = truncnorm.rvs(-1, 1, loc=0, scale=1, size=(N_SAMPLES, latent_dim)).astype(np.float32)
+        z = torch.from_numpy(z_np).to(device)
+        imgs = G(z)  # (B,3,H,W) in [-1,1] or [0,1]
+
+        # make a grid; normalize handles [-1,1]
+        grid = make_grid(imgs, nrow=N_COLS, normalize=True, scale_each=True)
+        return _to_pil_grid(grid)
+
+# ---------------- Gradio UI ----------------
+with gr.Blocks(title="MMD-PMish-NAS-GAN Image Generation") as demo:
+    gr.Markdown("## MMD-PMish-NAS-GAN Image Generation\nClick **Generate** to sample a fresh grid of images using the model.")
+    generate_btn = gr.Button("Generate", variant="primary")
+    out = gr.Image(label="Output Grid", type="pil")
+
+    generate_btn.click(fn=generate_grid, inputs=None, outputs=out, queue=True)
+
+if __name__ == "__main__":
+    demo.queue().launch()