Update app.py

app.py

@@ -1,46 +1,28 @@
 import os
 import math
-import uuid
 import numpy as np
 import onnxruntime as ort
 from PIL import Image
 import gradio as gr
 import tempfile
-import requests
+from huggingface_hub import hf_hub_download, HfApi
 
 # ---------------------------------------------------------------------------
-# Directory & model paths
+# CONFIG — only these two lines ever need changing
 # ---------------------------------------------------------------------------
-MODEL_DIR = "model"
-os.makedirs(MODEL_DIR, exist_ok=True)
-
-MODEL_PATHS = {
-    "esrgan_x2":  os.path.join(MODEL_DIR, "Real-ESRGAN_x2plus.onnx"),
-    "esrgan_x4":  os.path.join(MODEL_DIR, "Real-ESRGAN-x4plus.onnx"),
-    "srcnn_x4":   os.path.join(MODEL_DIR, "SRCNN_x4.onnx"),
-    "hresnet_x4": os.path.join(MODEL_DIR, "HResNet_x4.onnx"),
-}
+HF_REPO_ID  = "ParamDev/spectraGAN"   # your HF model repo
+HF_TOKEN    = os.environ.get("HF_TOKEN")  # set this in Space → Settings → Secrets
 
 # ---------------------------------------------------------------------------
-# Google Drive file IDs
-# TODO: Replace the SRCNN / HResNet IDs with your own uploaded ONNX exports.
-#       Steps to export SRCNN to ONNX:
-#         import torch; from srcnn import SRCNN
-#         model = SRCNN(); model.load_state_dict(torch.load("srcnn.pth"))
-#         dummy = torch.randn(1, 3, 128, 128)
-#         torch.onnx.export(model, dummy, "SRCNN_x4.onnx",
-#                           input_names=["input"], output_names=["output"],
-#                           dynamic_axes={"input":{2:"H",3:"W"},"output":{2:"H",3:"W"}})
-#       Same pattern applies for HResNet / EDSR.
+# Filenames used in the HF repo
 # ---------------------------------------------------------------------------
-DRIVE_IDS = {
-    "esrgan_x2":  "15xmXXZNH2wMyeQv4ie5hagT7eWK9MgP6",
-    "esrgan_x4":  "1wDBHad9RCJgJDGsPdapLYl3cr8j-PMJ6",
-    "srcnn_x4":   "YOUR_SRCNN_DRIVE_FILE_ID_HERE",    # <-- replace
-    "hresnet_x4": "YOUR_HRESNET_DRIVE_FILE_ID_HERE",  # <-- replace
+HF_FILENAMES = {
+    "esrgan_x2":  "Real-ESRGAN_x2plus.onnx",
+    "esrgan_x4":  "Real-ESRGAN-x4plus.onnx",
+    "srcnn_x4":   "SRCNN_x4.onnx",
+    "hresnet_x4": "HResNet_x4.onnx",
 }
 
-# Scale factor each model produces
 MODEL_SCALES = {
     "esrgan_x2":  2,
     "esrgan_x4":  4,
@@ -48,7 +30,6 @@ MODEL_SCALES = {
     "hresnet_x4": 4,
 }
 
-# Human-readable labels shown in the UI
 MODEL_LABELS = {
     "esrgan_x2":  "Real-ESRGAN ×2",
     "esrgan_x4":  "Real-ESRGAN ×4",
@@ -56,167 +37,316 @@ MODEL_LABELS = {
     "hresnet_x4": "HResNet ×4",
 }
 
-# ---------------------------------------------------------------------------
-# SR3 NOTE
-# ---------------------------------------------------------------------------
-# SR3 (Super-Resolution via Repeated Refinement, Ho et al. 2022) is a
-# *diffusion model* and cannot be exported to ONNX in the same way as
-# feed-forward CNNs.  To integrate SR3:
-#   1. Clone the official repo: https://github.com/Janspiry/Image-Super-Resolution-via-Iterative-Refinement
-#   2. Place your checkpoint in model/sr3_x4.pth
-#   3. Add a `run_sr3(image: Image) -> Image` function that runs the
-#      denoising loop using PyTorch directly (add `torch` to requirements).
-#   4. Wire `run_sr3` into `upscale_one_model` for key "sr3_x4".
-# SR3 is intentionally omitted from the ONNX pipeline to avoid misleading
-# model shape assumptions.
-# ---------------------------------------------------------------------------
+# Local scratch directory (writable in HF Spaces)
+CACHE_DIR = "/tmp/spectragan_models"
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+
+# ===========================================================================
+# STEP 1 — ONNX export for SRCNN and HResNet
+# Only runs if the file is missing from both local cache and the HF repo.
+# Uses torch which is available in HF Spaces by default.
+# ===========================================================================
+
+def _export_srcnn_onnx(out_path: str):
+    import torch
+    import torch.nn as nn
+
+    class SRCNN(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)
+            self.conv2 = nn.Conv2d(64, 32, kernel_size=5, padding=2)
+            self.conv3 = nn.Conv2d(32, 3,  kernel_size=5, padding=2)
+            self.relu  = nn.ReLU(inplace=True)
+
+        def forward(self, x):
+            return self.conv3(self.relu(self.conv2(self.relu(self.conv1(x)))))
+
+    class SRCNNx4(nn.Module):
+        """Bicubic ×4 upsample → SRCNN refinement."""
+        def __init__(self):
+            super().__init__()
+            self.srcnn = SRCNN()
+
+        def forward(self, x):
+            up = torch.nn.functional.interpolate(
+                x, scale_factor=4, mode="bicubic", align_corners=False
+            )
+            return self.srcnn(up)
 
+    model = SRCNNx4()
 
-# ---------------------------------------------------------------------------
-# Google Drive downloader
-# ---------------------------------------------------------------------------
-def download_from_drive(file_id: str, dest_path: str):
-    URL = "https://drive.google.com/uc?export=download"
-    session = requests.Session()
-    response = session.get(URL, params={"id": file_id}, stream=True)
-    token = None
-    for key, value in response.cookies.items():
-        if key.startswith("download_warning"):
-            token = value
-            break
-    if token:
-        response = session.get(URL, params={"id": file_id, "confirm": token}, stream=True)
-    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
-    with open(dest_path, "wb") as f:
-        for chunk in response.iter_content(chunk_size=32768):
-            if chunk:
-                f.write(chunk)
-    print(f"Downloaded → {dest_path}")
+    # Try to load pretrained weights; fall back to random init if unavailable
+    weights_url = (
+        "https://github.com/yjn870/SRCNN-pytorch/raw/master/models/srcnn_x4.pth"
+    )
+    weights_path = os.path.join(CACHE_DIR, "srcnn_x4.pth")
+    if not os.path.exists(weights_path):
+        try:
+            import urllib.request
+            print("  Downloading SRCNN pretrained weights …")
+            urllib.request.urlretrieve(weights_url, weights_path)
+        except Exception as e:
+            print(f"  [WARN] Could not download SRCNN weights ({e}). Using random init.")
+            weights_path = None
 
+    if weights_path and os.path.exists(weights_path):
+        state = torch.load(weights_path, map_location="cpu")
+        try:
+            model.srcnn.load_state_dict(state, strict=True)
+            print("  SRCNN pretrained weights loaded ✓")
+        except RuntimeError:
+            model.srcnn.load_state_dict(state, strict=False)
+            print("  SRCNN weights loaded (partial) ✓")
+
+    model.eval()
+    dummy = torch.zeros(1, 3, 128, 128)
+    torch.onnx.export(
+        model, dummy, out_path,
+        opset_version=17,
+        input_names=["input"], output_names=["output"],
+        dynamic_axes={"input": {2: "H", 3: "W"}, "output": {2: "H", 3: "W"}},
+    )
+    print(f"  Exported SRCNN → {out_path}")
 
-# ---------------------------------------------------------------------------
-# LOCAL_ONLY mode
-# ---------------------------------------------------------------------------
-# Set to True after running export_models.py to load ONNX files straight
-# from the model/ directory without needing any Google Drive IDs.
-# Set to False (and fill in DRIVE_IDS) to download from Drive as before.
-LOCAL_ONLY = True
 
-# ---------------------------------------------------------------------------
-# Download models if missing
-# ---------------------------------------------------------------------------
-for key, path in MODEL_PATHS.items():
-    if os.path.isfile(path):
-        continue  # already on disk — nothing to do
+def _export_hresnet_onnx(out_path: str):
+    import torch
+    import torch.nn as nn
 
-    if LOCAL_ONLY:
-        print(f"[INFO] {key}: file not found at {path}. "
-              "Run export_models.py first, or set LOCAL_ONLY=False to use Drive.")
+    class ResBlock(nn.Module):
+        def __init__(self, n_feats):
+            super().__init__()
+            self.body = nn.Sequential(
+                nn.Conv2d(n_feats, n_feats, 3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(n_feats, n_feats, 3, padding=1),
+            )
+        def forward(self, x):
+            return x + self.body(x)
+
+    class EDSR(nn.Module):
+        """EDSR-baseline: 16 residual blocks, 64 channels, ×4."""
+        def __init__(self):
+            super().__init__()
+            n_feats = 64
+            self.head = nn.Conv2d(3, n_feats, 3, padding=1)
+            self.body = nn.Sequential(*[ResBlock(n_feats) for _ in range(16)])
+            self.body_tail = nn.Conv2d(n_feats, n_feats, 3, padding=1)
+            # ×4 via two ×2 pixel-shuffle stages
+            self.tail = nn.Sequential(
+                nn.Conv2d(n_feats, n_feats * 4, 3, padding=1),
+                nn.PixelShuffle(2),
+                nn.Conv2d(n_feats, n_feats * 4, 3, padding=1),
+                nn.PixelShuffle(2),
+                nn.Conv2d(n_feats, 3, 3, padding=1),
+            )
+
+        def forward(self, x):
+            x = self.head(x)
+            res = self.body_tail(self.body(x))
+            return self.tail(x + res)
+
+    model = EDSR()
+
+    # Try to load pretrained weights from HuggingFace
+    weights_url = (
+        "https://huggingface.co/eugenesiow/edsr-base/resolve/main/pytorch_model_4x.pt"
+    )
+    weights_path = os.path.join(CACHE_DIR, "edsr_x4.pt")
+    if not os.path.exists(weights_path):
+        try:
+            import urllib.request
+            print("  Downloading EDSR pretrained weights from HuggingFace …")
+            urllib.request.urlretrieve(weights_url, weights_path)
+        except Exception as e:
+            print(f"  [WARN] Could not download EDSR weights ({e}). Using random init.")
+            weights_path = None
+
+    if weights_path and os.path.exists(weights_path):
+        state = torch.load(weights_path, map_location="cpu")
+        if "model" in state:       state = state["model"]
+        if "state_dict" in state:  state = state["state_dict"]
+        state = {k.replace("module.", ""): v for k, v in state.items()}
+        try:
+            model.load_state_dict(state, strict=True)
+            print("  EDSR pretrained weights loaded ✓")
+        except RuntimeError:
+            model.load_state_dict(state, strict=False)
+            print("  EDSR weights loaded (partial) ✓")
+
+    model.eval()
+    dummy = torch.zeros(1, 3, 128, 128)
+    torch.onnx.export(
+        model, dummy, out_path,
+        opset_version=17,
+        input_names=["input"], output_names=["output"],
+        dynamic_axes={"input": {2: "H", 3: "W"}, "output": {2: "H", 3: "W"}},
+    )
+    print(f"  Exported HResNet → {out_path}")
+
+
+# Map model keys that need local export to their export functions
+EXPORTERS = {
+    "srcnn_x4":   _export_srcnn_onnx,
+    "hresnet_x4": _export_hresnet_onnx,
+}
+
+
+def _upload_to_hub(local_path: str, filename: str):
+    """Upload a file to the HF repo so future Space restarts skip the export."""
+    if not HF_TOKEN:
+        print(f"  [INFO] HF_TOKEN not set — skipping upload of {filename}.")
+        print(f"         Set HF_TOKEN in Space Secrets to persist generated models.")
+        return
+    try:
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=local_path,
+            path_in_repo=filename,
+            repo_id=HF_REPO_ID,
+            repo_type="model",
+            token=HF_TOKEN,
+        )
+        print(f"  Uploaded {filename} → {HF_REPO_ID} ✓")
+    except Exception as e:
+        print(f"  [WARN] Upload failed for {filename}: {e}")
+
+
+# ===========================================================================
+# STEP 2 — Fetch or generate all models
+# Priority: HF Hub cache → local export → upload back to Hub
+# ===========================================================================
+
+MODEL_PATHS = {}   # key → local file path ready for onnxruntime
+
+for key, filename in HF_FILENAMES.items():
+    local_path = os.path.join(CACHE_DIR, filename)
+
+    # Already cached locally from a previous run this session
+    if os.path.exists(local_path):
+        print(f"[CACHE] {MODEL_LABELS[key]} already in /tmp cache.")
+        MODEL_PATHS[key] = local_path
         continue
 
-    file_id = DRIVE_IDS[key]
-    if file_id.startswith("YOUR_"):
-        print(f"[WARN] Skipping {key}: Google Drive ID not set. "
-              "Run export_models.py or update DRIVE_IDS in app.py.")
+    # Try downloading from HF Hub first
+    try:
+        print(f"Fetching {MODEL_LABELS[key]} from HF Hub …")
+        dl_path = hf_hub_download(
+            repo_id=HF_REPO_ID,
+            filename=filename,
+            repo_type="model",
+            token=HF_TOKEN,
+            local_dir=CACHE_DIR,
+        )
+        MODEL_PATHS[key] = dl_path
+        print(f"  ✓  {dl_path}")
+        continue
+    except Exception as e:
+        print(f"  Not found on Hub ({e})")
+
+    # Not on Hub — export it locally if we have an exporter for it
+    if key in EXPORTERS:
+        print(f"Generating {MODEL_LABELS[key]} via local export …")
+        try:
+            EXPORTERS[key](local_path)
+            MODEL_PATHS[key] = local_path
+            # Upload back to Hub so next restart skips this step
+            print(f"  Uploading to Hub for future restarts …")
+            _upload_to_hub(local_path, filename)
+        except Exception as e:
+            print(f"  [ERROR] Export failed for {key}: {e}")
     else:
-        print(f"Downloading {MODEL_LABELS[key]} …")
-        download_from_drive(file_id, path)
+        print(f"  [WARN] {key} not on Hub and no exporter available — skipping.")
 
 
-# ---------------------------------------------------------------------------
-# Load ONNX sessions (only for models that have been downloaded)
-# ---------------------------------------------------------------------------
+# ===========================================================================
+# STEP 3 — Load ONNX sessions
+# ===========================================================================
+
 sess_opts = ort.SessionOptions()
 sess_opts.intra_op_num_threads = 2
 sess_opts.inter_op_num_threads = 2
 
-SESSIONS = {}      # key → ort.InferenceSession
-INPUT_SHAPES = {}  # key → (H_in, W_in)
+SESSIONS     = {}
+INPUT_SHAPES = {}
 
 for key, path in MODEL_PATHS.items():
-    if os.path.isfile(path):
-        try:
-            sess = ort.InferenceSession(
-                path,
-                sess_options=sess_opts,
-                providers=["CPUExecutionProvider"]
-            )
-            meta = sess.get_inputs()[0]
-            shape = tuple(meta.shape)
-            # shape is (1, 3, H, W) for fixed-size models, or (1,3,None,None) for dynamic
-            h = int(shape[2]) if shape[2] is not None and str(shape[2]).isdigit() else 128
-            w = int(shape[3]) if shape[3] is not None and str(shape[3]).isdigit() else 128
-            SESSIONS[key] = (sess, meta)
-            INPUT_SHAPES[key] = (h, w)
-            print(f"Loaded {MODEL_LABELS[key]}  tile={h}×{w}")
-        except Exception as e:
-            print(f"[ERROR] Could not load {key}: {e}")
-    else:
-        print(f"[INFO] {key} not available — will be skipped in comparisons.")
-
+    try:
+        sess  = ort.InferenceSession(
+            path,
+            sess_options=sess_opts,
+            providers=["CPUExecutionProvider"],
+        )
+        meta  = sess.get_inputs()[0]
+        shape = tuple(meta.shape)
+        h = int(shape[2]) if isinstance(shape[2], int) else 128
+        w = int(shape[3]) if isinstance(shape[3], int) else 128
+        SESSIONS[key]     = (sess, meta)
+        INPUT_SHAPES[key] = (h, w)
+        print(f"Loaded  {MODEL_LABELS[key]}  tile={h}×{w}")
+    except Exception as e:
+        print(f"[ERROR] Could not load {key}: {e}")
+
+
+# ===========================================================================
+# STEP 4 — Inference helpers
+# ===========================================================================
 
-# ---------------------------------------------------------------------------
-# Tile-based upscale for a single ONNX model
-# ---------------------------------------------------------------------------
 def run_onnx_tile(sess, meta, tile_np: np.ndarray) -> np.ndarray:
-    """Run one tile through any ONNX session. tile_np is HWC float32 [0,1]."""
-    patch = np.transpose(tile_np, (2, 0, 1))[None, ...]  # NCHW
+    patch = np.transpose(tile_np, (2, 0, 1))[None, ...]
     out   = sess.run(None, {meta.name: patch})[0]
-    out   = np.squeeze(out, axis=0)
-    return np.transpose(out, (1, 2, 0))                   # back to HWC
+    return np.transpose(np.squeeze(out, 0), (1, 2, 0))
 
 
 def tile_upscale_model(input_img: Image.Image, key: str, max_dim: int = 1024) -> Image.Image:
-    """
-    Upscale *input_img* using the ONNX model identified by *key*.
-    Returns a PIL Image at the upscaled resolution.
-    """
     if key not in SESSIONS:
-        raise ValueError(f"Model '{key}' is not loaded. Check the Drive ID / path.")
+        raise ValueError(f"Model '{key}' is not loaded.")
 
-    sess, meta = SESSIONS[key]
-    H_in, W_in = INPUT_SHAPES[key]
-    scale       = MODEL_SCALES[key]
+    sess, meta        = SESSIONS[key]
+    H_in, W_in        = INPUT_SHAPES[key]
+    scale              = MODEL_SCALES[key]
 
-    # Optionally cap input size to avoid OOM on large images
     w, h = input_img.size
     if w > max_dim or h > max_dim:
-        factor = max_dim / float(max(w, h))
+        factor    = max_dim / float(max(w, h))
         input_img = input_img.resize((int(w * factor), int(h * factor)), Image.LANCZOS)
 
-    arr = np.array(input_img.convert("RGB")).astype(np.float32) / 255.0
+    arr               = np.array(input_img.convert("RGB")).astype(np.float32) / 255.0
     h_orig, w_orig, _ = arr.shape
+    tiles_h           = math.ceil(h_orig / H_in)
+    tiles_w           = math.ceil(w_orig / W_in)
 
-    tiles_h = math.ceil(h_orig / H_in)
-    tiles_w = math.ceil(w_orig / W_in)
-    pad_h   = tiles_h * H_in - h_orig
-    pad_w   = tiles_w * W_in - w_orig
-
-    arr_padded = np.pad(arr, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
-    out_arr    = np.zeros((tiles_h * H_in * scale, tiles_w * W_in * scale, 3), dtype=np.float32)
-
+    arr_padded = np.pad(
+        arr,
+        ((0, tiles_h * H_in - h_orig), (0, tiles_w * W_in - w_orig), (0, 0)),
+        mode="reflect",
+    )
+    out_arr = np.zeros(
+        (tiles_h * H_in * scale, tiles_w * W_in * scale, 3), dtype=np.float32
+    )
     for i in range(tiles_h):
         for j in range(tiles_w):
             y0, x0   = i * H_in, j * W_in
-            tile      = arr_padded[y0:y0 + H_in, x0:x0 + W_in, :]
-            up_tile   = run_onnx_tile(sess, meta, tile)
-            oy0, ox0  = i * H_in * scale, j * W_in * scale
-            out_arr[oy0:oy0 + H_in * scale, ox0:ox0 + W_in * scale, :] = up_tile
+            up_tile  = run_onnx_tile(sess, meta, arr_padded[y0:y0+H_in, x0:x0+W_in])
+            oy0, ox0 = i * H_in * scale, j * W_in * scale
+            out_arr[oy0:oy0+H_in*scale, ox0:ox0+W_in*scale] = up_tile
 
-    final = np.clip(out_arr[0:h_orig * scale, 0:w_orig * scale, :], 0.0, 1.0)
+    final = np.clip(out_arr[:h_orig*scale, :w_orig*scale], 0.0, 1.0)
     return Image.fromarray((final * 255.0).round().astype(np.uint8))
 
 
-def upscale_8x_from_4x(input_img: Image.Image, key: str) -> Image.Image:
-    """Run ×4 model, then bicubic ×2 to reach ×8."""
+def upscale_8x(input_img: Image.Image, key: str) -> Image.Image:
     img_4x = tile_upscale_model(input_img, key)
     w, h   = input_img.size
     return img_4x.resize((w * 8, h * 8), Image.LANCZOS)
 
 
-# ---------------------------------------------------------------------------
-# Core comparison function  (called by the Gradio button)
-# ---------------------------------------------------------------------------
+# ===========================================================================
+# STEP 5 — Gradio comparison callback
+# ===========================================================================
+
 def compare_models(
     input_img: Image.Image,
     use_esrgan_x2: bool,
@@ -226,36 +356,32 @@ def compare_models(
     include_8x:    bool,
 ):
     if input_img is None:
-        return [None] * 8   # 4 preview + 4 download slots
-
-    selection = []
-    if use_esrgan_x2: selection.append("esrgan_x2")
-    if use_esrgan_x4: selection.append("esrgan_x4")
-    if use_srcnn:     selection.append("srcnn_x4")
-    if use_hresnet:   selection.append("hresnet_x4")
-
-    previews   = []
-    downloads  = []
+        return [None] * 8
+
+    wanted = {
+        "esrgan_x2":  use_esrgan_x2,
+        "esrgan_x4":  use_esrgan_x4,
+        "srcnn_x4":   use_srcnn,
+        "hresnet_x4": use_hresnet,
+    }
+    selection = [k for k, v in wanted.items() if v]
+    previews  = []
+    downloads = []
 
     for key in selection:
         if key not in SESSIONS:
             previews.append(None)
-            downloads.append(gr.DownloadButton(label=f"{MODEL_LABELS[key]} – not loaded",
-                                               visible=True, value=None))
+            downloads.append(None)
             continue
-
         try:
-            if include_8x and MODEL_SCALES[key] == 4:
-                result = upscale_8x_from_4x(input_img, key)
-                suffix = "×8"
-            else:
-                result = tile_upscale_model(input_img, key)
-                suffix = f"×{MODEL_SCALES[key]}"
-
+            result = (
+                upscale_8x(input_img, key)
+                if include_8x and MODEL_SCALES[key] == 4
+                else tile_upscale_model(input_img, key)
+            )
             tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
             result.save(tmp.name, format="PNG")
             tmp.close()
-
             previews.append(result)
             downloads.append(tmp.name)
         except Exception as e:
@@ -263,55 +389,38 @@ def compare_models(
             previews.append(None)
             downloads.append(None)
 
-    # Pad to always return exactly 4 preview + 4 download values
     while len(previews)  < 4: previews.append(None)
     while len(downloads) < 4: downloads.append(None)
+    return previews + downloads
 
-    return previews + downloads   # 8-element list
 
+# ===========================================================================
+# STEP 6 — Gradio UI
+# ===========================================================================
 
-# ---------------------------------------------------------------------------
-# Gradio UI – side-by-side comparison layout
-# ---------------------------------------------------------------------------
 css = """
 body { font-family: 'Segoe UI', sans-serif; }
-
 .panel-title {
-    text-align: center;
-    font-weight: 700;
-    font-size: 0.85rem;
-    letter-spacing: 0.08em;
-    text-transform: uppercase;
-    margin-bottom: 4px;
-    color: #555;
+    text-align: center; font-weight: 700; font-size: 0.85rem;
+    letter-spacing: 0.08em; text-transform: uppercase;
+    margin-bottom: 4px; color: #555;
 }
-
 #run-btn {
     background: linear-gradient(135deg, #1a1a2e, #16213e) !important;
-    color: #e2e2e2 !important;
-    font-size: 1rem !important;
-    font-weight: 600 !important;
-    border-radius: 8px !important;
+    color: #e2e2e2 !important; font-size: 1rem !important;
+    font-weight: 600 !important; border-radius: 8px !important;
     padding: 12px 28px !important;
 }
-#run-btn:hover {
-    background: linear-gradient(135deg, #0f3460, #533483) !important;
-}
-
+#run-btn:hover { background: linear-gradient(135deg, #0f3460, #533483) !important; }
 .dl-btn button {
-    background: #f0f4ff !important;
-    border: 1px solid #c5d0f5 !important;
-    color: #333 !important;
-    font-size: 0.78rem !important;
-    border-radius: 6px !important;
-    width: 100%;
+    background: #f0f4ff !important; border: 1px solid #c5d0f5 !important;
+    color: #333 !important; font-size: 0.78rem !important;
+    border-radius: 6px !important; width: 100%;
 }
-
 .model-toggle label { font-size: 0.9rem; }
 """
 
-ALL_KEYS    = ["esrgan_x2", "esrgan_x4", "srcnn_x4", "hresnet_x4"]
-PANEL_KEYS  = ALL_KEYS   # order for the 4 comparison panels
+PANEL_KEYS = ["esrgan_x2", "esrgan_x4", "srcnn_x4", "hresnet_x4"]
 
 with gr.Blocks(css=css, title="SpectraGAN — Multi-Model Comparison") as demo:
 
@@ -320,59 +429,35 @@ with gr.Blocks(css=css, title="SpectraGAN — Multi-Model Comparison") as demo:
     Upload an image, select models, and compare results side by side.
     """)
 
-    # ── Input row ──────────────────────────────────────────────────────────
     with gr.Row():
         inp_image = gr.Image(type="pil", label="Source Image", scale=2)
-
         with gr.Column(scale=1):
             gr.Markdown("### Models to compare")
-            chk_esrgan_x2  = gr.Checkbox(label="Real-ESRGAN ×2",  value=True,  elem_classes="model-toggle")
-            chk_esrgan_x4  = gr.Checkbox(label="Real-ESRGAN ×4",  value=True,  elem_classes="model-toggle")
-            chk_srcnn      = gr.Checkbox(label="SRCNN ×4",        value=True,  elem_classes="model-toggle")
-            chk_hresnet    = gr.Checkbox(label="HResNet ×4",      value=True,  elem_classes="model-toggle")
-
+            chk_esrgan_x2 = gr.Checkbox(label="Real-ESRGAN ×2",  value=True, elem_classes="model-toggle")
+            chk_esrgan_x4 = gr.Checkbox(label="Real-ESRGAN ×4",  value=True, elem_classes="model-toggle")
+            chk_srcnn     = gr.Checkbox(label="SRCNN ×4",         value=True, elem_classes="model-toggle")
+            chk_hresnet   = gr.Checkbox(label="HResNet ×4",       value=True, elem_classes="model-toggle")
             gr.Markdown("### Options")
-            chk_8x = gr.Checkbox(label="Also apply ×8 post-resize on ×4 models", value=False)
-
+            chk_8x  = gr.Checkbox(label="Also apply ×8 post-resize on ×4 models", value=False)
             run_btn = gr.Button("⚡ Run Comparison", elem_id="run-btn")
 
-    # ── Comparison grid ────────────────────────────────────────────────────
-    gr.Markdown("---")
-    gr.Markdown("## Results")
-
-    previews  = []
-    dl_btns   = []
+    gr.Markdown("---\n## Results")
+    previews = []
+    dl_btns  = []
 
     with gr.Row():
         for key in PANEL_KEYS:
             with gr.Column():
                 gr.HTML(f'<div class="panel-title">{MODEL_LABELS[key]}</div>')
-                img_out = gr.Image(
-                    type="pil",
-                    label=MODEL_LABELS[key],
-                    show_label=False,
-                    height=320,
-                )
-                dl_out = gr.DownloadButton(
-                    label="⬇ Download PNG",
-                    elem_classes="dl-btn",
-                    visible=True,
-                )
-                previews.append(img_out)
-                dl_btns.append(dl_out)
-
-    # ── Wire up ─────────────────────────────────────────────────────────────
+                previews.append(gr.Image(type="pil", show_label=False, height=320))
+                dl_btns.append(gr.DownloadButton(
+                    label="⬇ Download PNG", elem_classes="dl-btn", visible=True
+                ))
+
     run_btn.click(
         fn=compare_models,
-        inputs=[
-            inp_image,
-            chk_esrgan_x2,
-            chk_esrgan_x4,
-            chk_srcnn,
-            chk_hresnet,
-            chk_8x,
-        ],
-        outputs=previews + dl_btns,   # 8 outputs: 4 images + 4 download buttons
+        inputs=[inp_image, chk_esrgan_x2, chk_esrgan_x4, chk_srcnn, chk_hresnet, chk_8x],
+        outputs=previews + dl_btns,
     )
 
 demo.launch(server_name="0.0.0.0", server_port=7860)