Update app.py

app.py

@@ -5,309 +5,156 @@ import onnxruntime as ort
 from PIL import Image
 import gradio as gr
 import tempfile
-from huggingface_hub import hf_hub_download, HfApi
+import gdown
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
 
 # ---------------------------------------------------------------------------
-# CONFIG — only these two lines ever need changing
+# Paths & constants
 # ---------------------------------------------------------------------------
-HF_REPO_ID  = "ParamDev/spectraGAN"   # your HF model repo
-HF_TOKEN    = os.environ.get("HF_TOKEN")  # set this in Space → Settings → Secrets
+CACHE_DIR = "/tmp/spectragan"
+os.makedirs(CACHE_DIR, exist_ok=True)
 
-# ---------------------------------------------------------------------------
-# Filenames used in the HF repo
-# ---------------------------------------------------------------------------
-HF_FILENAMES = {
-    "esrgan_x2":  "Real-ESRGAN_x2plus.onnx",
-    "esrgan_x4":  "Real-ESRGAN-x4plus.onnx",
-    "srcnn_x4":   "SRCNN_x4.onnx",
-    "hresnet_x4": "HResNet_x4.onnx",
+# Google Drive IDs for ESRGAN ONNX files
+DRIVE_IDS = {
+    "esrgan_x4":  "1wDBHad9RCJgJDGsPdapLYl3cr8j-PMJ6",
+    "hresnet_x4": "15xmXXZNH2wMyeQv4ie5hagT7eWK9MgP6",   # placeholder = ESRGAN x2
 }
 
-MODEL_SCALES = {
-    "esrgan_x2":  2,
-    "esrgan_x4":  4,
-    "srcnn_x4":   4,
-    "hresnet_x4": 4,
-}
+# SRCNN .pth lives in the same HF Space repo (uploaded by user)
+SRCNN_PTH = os.path.join(os.path.dirname(__file__), "srcnn_x4.pth")
 
 MODEL_LABELS = {
-    "esrgan_x2":  "Real-ESRGAN ×2",
     "esrgan_x4":  "Real-ESRGAN ×4",
     "srcnn_x4":   "SRCNN ×4",
-    "hresnet_x4": "HResNet ×4",
+    "hresnet_x4": "HResNet ×4",        # placeholder until real weights available
 }
 
-# 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()
-
-    # 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}")
-
-
-def _export_hresnet_onnx(out_path: str):
-    import torch
-    import torch.nn as nn
-
-    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,
+MODEL_SCALES = {
+    "esrgan_x4":  4,
+    "srcnn_x4":   4,
+    "hresnet_x4": 2,                   # underlying model is ESRGAN x2
 }
 
 
-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
+# SRCNN architecture  (3-layer, matches the .pth weights)
 # ===========================================================================
+class SRCNN(nn.Module):
+    def __init__(self, num_channels: int = 3):
+        super().__init__()
+        self.conv1 = nn.Conv2d(num_channels, 64, kernel_size=9, padding=9 // 2)
+        self.conv2 = nn.Conv2d(64, 32, kernel_size=5, padding=5 // 2)
+        self.conv3 = nn.Conv2d(32, num_channels, kernel_size=5, padding=5 // 2)
+        self.relu  = nn.ReLU(inplace=True)
 
-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
-
-    # 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"  [WARN] {key} not on Hub and no exporter available — skipping.")
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv3(self.relu(self.conv2(self.relu(self.conv1(x)))))
 
 
 # ===========================================================================
-# STEP 3 — Load ONNX sessions
+# Model loading
 # ===========================================================================
 
 sess_opts = ort.SessionOptions()
 sess_opts.intra_op_num_threads = 2
 sess_opts.inter_op_num_threads = 2
 
-SESSIONS     = {}
-INPUT_SHAPES = {}
+ONNX_SESSIONS = {}   # key → (ort.InferenceSession, input_meta)
+SRCNN_MODEL   = None
+
+
+def _load_esrgan_onnx(key: str):
+    """Download ESRGAN ONNX from Drive (gdown handles confirmation pages)."""
+    filename = f"{key}.onnx"
+    dest     = os.path.join(CACHE_DIR, filename)
+    if not os.path.exists(dest):
+        print(f"Downloading {MODEL_LABELS[key]} from Drive …")
+        gdown.download(id=DRIVE_IDS[key], output=dest, quiet=False, fuzzy=True)
+    if os.path.exists(dest):
+        sess = ort.InferenceSession(dest, sess_options=sess_opts,
+                                    providers=["CPUExecutionProvider"])
+        meta = sess.get_inputs()[0]
+        ONNX_SESSIONS[key] = (sess, meta)
+        print(f"Loaded {MODEL_LABELS[key]} ✓")
+    else:
+        print(f"[ERROR] Could not load {key} — file missing after download.")
+
 
-for key, path in MODEL_PATHS.items():
+def _load_srcnn_pth():
+    """Load SRCNN directly from the .pth file in the Space repo."""
+    global SRCNN_MODEL
+    if not os.path.exists(SRCNN_PTH):
+        print(f"[WARN] srcnn_x4.pth not found at {SRCNN_PTH} — SRCNN will be skipped.")
+        return
+    model = SRCNN(num_channels=3)
+    state = torch.load(SRCNN_PTH, map_location="cpu")
+    # Unwrap common checkpoint wrappers
+    for key in ("model", "state_dict", "params"):
+        if isinstance(state, dict) and key in state:
+            state = state[key]
+    state = {k.replace("module.", ""): v for k, v in state.items()}
     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}")
+        model.load_state_dict(state, strict=True)
+    except RuntimeError:
+        model.load_state_dict(state, strict=False)
+        print("[WARN] SRCNN loaded with strict=False (minor key mismatch).")
+    model.eval()
+    SRCNN_MODEL = model
+    print("Loaded SRCNN ×4 from .pth ✓")
+
+
+# Boot-time loading
+for k in ("esrgan_x4", "hresnet_x4"):
+    try:
+        _load_esrgan_onnx(k)
     except Exception as e:
-        print(f"[ERROR] Could not load {key}: {e}")
+        print(f"[ERROR] {k}: {e}")
+
+_load_srcnn_pth()
 
 
 # ===========================================================================
-# STEP 4 — Inference helpers
+# Inference helpers
 # ===========================================================================
 
-def run_onnx_tile(sess, meta, tile_np: np.ndarray) -> np.ndarray:
-    patch = np.transpose(tile_np, (2, 0, 1))[None, ...]
+def _onnx_tile(sess, meta, tile: np.ndarray) -> np.ndarray:
+    """HWC float32 [0,1] → HWC float32 [0,1]."""
+    patch = tile.transpose(2, 0, 1)[None, ...]
     out   = sess.run(None, {meta.name: patch})[0]
-    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:
-    if key not in SESSIONS:
-        raise ValueError(f"Model '{key}' is not loaded.")
-
-    sess, meta        = SESSIONS[key]
-    H_in, W_in        = INPUT_SHAPES[key]
-    scale              = MODEL_SCALES[key]
-
+    return out.squeeze(0).transpose(1, 2, 0)
+
+
+def _srcnn_tile(tile: np.ndarray, scale: int = 4) -> np.ndarray:
+    """Bicubic upsample → SRCNN refine. HWC float32 [0,1] → HWC float32."""
+    t = torch.from_numpy(tile.transpose(2, 0, 1)).unsqueeze(0)   # NCHW
+    h, w = t.shape[2], t.shape[3]
+    up   = F.interpolate(t, size=(h * scale, w * scale),
+                         mode="bicubic", align_corners=False)
+    with torch.no_grad():
+        out = SRCNN_MODEL(up)
+    return out.squeeze(0).permute(1, 2, 0).numpy()
+
+
+def upscale(input_img: Image.Image, model_key: str, max_dim: int = 1024) -> Image.Image:
+    """
+    Tile-based upscale dispatcher.
+    Works for both ONNX sessions (ESRGAN) and the torch SRCNN model.
+    """
+    # Guard
+    if model_key == "srcnn_x4" and SRCNN_MODEL is None:
+        raise RuntimeError("SRCNN model not loaded.")
+    if model_key in ("esrgan_x4", "hresnet_x4") and model_key not in ONNX_SESSIONS:
+        raise RuntimeError(f"{MODEL_LABELS[model_key]} not loaded.")
+
+    scale = MODEL_SCALES[model_key]
+
+    # Tile size: ESRGAN uses fixed 128×128 LR tiles;
+    # SRCNN works on any size (we tile at 128 for consistency).
+    TILE = 128
+
+    # Cap input size
     w, h = input_img.size
     if w > max_dim or h > max_dim:
         factor    = max_dim / float(max(w, h))
@@ -315,149 +162,189 @@ def tile_upscale_model(input_img: Image.Image, key: str, max_dim: int = 1024) ->
 
     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 / TILE)
+    tiles_w           = math.ceil(w_orig / TILE)
 
-    arr_padded = np.pad(
+    arr_pad = np.pad(
         arr,
-        ((0, tiles_h * H_in - h_orig), (0, tiles_w * W_in - w_orig), (0, 0)),
+        ((0, tiles_h * TILE - h_orig), (0, tiles_w * TILE - w_orig), (0, 0)),
         mode="reflect",
     )
-    out_arr = np.zeros(
-        (tiles_h * H_in * scale, tiles_w * W_in * scale, 3), dtype=np.float32
-    )
+    out = np.zeros((tiles_h * TILE * scale, tiles_w * TILE * 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
-            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
+            y0, x0  = i * TILE, j * TILE
+            tile     = arr_pad[y0:y0 + TILE, x0:x0 + TILE]
 
-    final = np.clip(out_arr[:h_orig*scale, :w_orig*scale], 0.0, 1.0)
+            if model_key == "srcnn_x4":
+                up_tile = _srcnn_tile(tile, scale=scale)
+            else:
+                sess, meta = ONNX_SESSIONS[model_key]
+                up_tile    = _onnx_tile(sess, meta, tile)
+
+            oy0, ox0 = i * TILE * scale, j * TILE * scale
+            out[oy0:oy0 + TILE * scale, ox0:ox0 + TILE * scale] = up_tile
+
+    final = np.clip(out[:h_orig * scale, :w_orig * scale], 0.0, 1.0)
     return Image.fromarray((final * 255.0).round().astype(np.uint8))
 
 
-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)
+def make_comparison_png(original: Image.Image, upscaled: Image.Image) -> str:
+    """
+    Save a side-by-side PNG (original | upscaled, same display height)
+    that the ImageSlider widget will use.
+    """
+    up_w, up_h = upscaled.size
+    orig_resized = original.resize((up_w, up_h), Image.LANCZOS)
+
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    orig_resized.save(tmp.name)          # left image for slider
+    tmp.close()
+
+    tmp2 = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    upscaled.save(tmp2.name)             # right image for slider
+    tmp2.close()
+
+    return tmp.name, tmp2.name
 
 
 # ===========================================================================
-# STEP 5 — Gradio comparison callback
+# Gradio callback
 # ===========================================================================
 
-def compare_models(
-    input_img: Image.Image,
-    use_esrgan_x2: bool,
-    use_esrgan_x4: bool,
-    use_srcnn:     bool,
-    use_hresnet:   bool,
-    include_8x:    bool,
-):
+def run_upscale(input_img: Image.Image, model_name: str):
     if input_img is None:
-        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(None)
-            continue
-        try:
-            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:
-            print(f"[ERROR] {key}: {e}")
-            previews.append(None)
-            downloads.append(None)
+        return None, None, None
+
+    # Map display label → key
+    key = next(k for k, v in MODEL_LABELS.items() if v == model_name)
+
+    result   = upscale(input_img, key)
+    orig_path, up_path = make_comparison_png(input_img, result)
 
-    while len(previews)  < 4: previews.append(None)
-    while len(downloads) < 4: downloads.append(None)
-    return previews + downloads
+    # Also save download copy
+    dl_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    result.save(dl_tmp.name, format="PNG")
+    dl_tmp.close()
+
+    return (orig_path, up_path), result, dl_tmp.name
 
 
 # ===========================================================================
-# STEP 6 — Gradio UI
+# Gradio UI
 # ===========================================================================
 
 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;
+@import url('https://fonts.googleapis.com/css2?family=DM+Sans:wght@400;600;700&display=swap');
+
+body, .gradio-container { font-family: 'DM Sans', sans-serif !important; }
+
+#title {
+    text-align: center;
+    padding: 24px 0 8px;
 }
+#title h1 {
+    font-size: 2rem;
+    font-weight: 700;
+    letter-spacing: -0.5px;
+    margin: 0;
+}
+#title p { color: #666; margin: 4px 0 0; }
+
 #run-btn {
-    background: linear-gradient(135deg, #1a1a2e, #16213e) !important;
-    color: #e2e2e2 !important; font-size: 1rem !important;
-    font-weight: 600 !important; border-radius: 8px !important;
-    padding: 12px 28px !important;
+    background: linear-gradient(135deg, #0f0c29, #302b63, #24243e) !important;
+    color: #fff !important;
+    font-weight: 700 !important;
+    font-size: 1rem !important;
+    border-radius: 10px !important;
+    padding: 14px 0 !important;
+    width: 100%;
+    letter-spacing: 0.03em;
+    transition: opacity 0.2s;
 }
-#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%;
+#run-btn:hover { opacity: 0.85; }
+
+#dl-btn button {
+    background: #f4f4f4 !important;
+    border: 1px solid #ddd !important;
+    color: #333 !important;
+    border-radius: 8px !important;
+    width: 100%;
+    font-size: 0.85rem !important;
+}
+
+.section-label {
+    font-size: 0.75rem;
+    font-weight: 700;
+    letter-spacing: 0.1em;
+    text-transform: uppercase;
+    color: #999;
+    margin-bottom: 6px;
 }
-.model-toggle label { font-size: 0.9rem; }
 """
 
-PANEL_KEYS = ["esrgan_x2", "esrgan_x4", "srcnn_x4", "hresnet_x4"]
+available_models = [v for k, v in MODEL_LABELS.items()
+                    if k == "srcnn_x4" and SRCNN_MODEL is not None
+                    or k in ONNX_SESSIONS]
+
+# Always show all three in dropdown regardless of load state
+# (shows error in output if model failed to load)
+dropdown_choices = list(MODEL_LABELS.values())
 
-with gr.Blocks(css=css, title="SpectraGAN — Multi-Model Comparison") as demo:
+with gr.Blocks(css=css, title="SpectraGAN Upscaler") as demo:
 
-    gr.Markdown("""
-    # 🖼️ SpectraGAN — Multi-Model Upscaler Comparison
-    Upload an image, select models, and compare results side by side.
+    gr.HTML("""
+    <div id="title">
+        <h1>🖼️ SpectraGAN Upscaler</h1>
+        <p>Choose a model, upscale your image, and drag the slider to compare.</p>
+    </div>
     """)
 
-    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")
-            gr.Markdown("### Options")
-            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")
-
-    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>')
-                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
-                ))
+    with gr.Row(equal_height=True):
+
+        # ── Left panel: controls ───────────────────────────────────────────
+        with gr.Column(scale=1, min_width=260):
+            gr.HTML('<div class="section-label">Source Image</div>')
+            inp_image = gr.Image(type="pil", show_label=False, height=260)
+
+            gr.HTML('<div class="section-label" style="margin-top:16px">Model</div>')
+            model_dropdown = gr.Dropdown(
+                choices=dropdown_choices,
+                value=dropdown_choices[0],
+                show_label=False,
+            )
+
+            run_btn = gr.Button("⚡ Upscale", elem_id="run-btn")
+
+            dl_btn = gr.DownloadButton(
+                label="⬇ Download upscaled PNG",
+                elem_id="dl-btn",
+                visible=True,
+            )
+
+        # ── Right panel: results ───────────────────────────────────────────
+        with gr.Column(scale=2):
+
+            gr.HTML('<div class="section-label">Before / After</div>')
+            slider = gr.ImageSlider(
+                label="Drag to compare",
+                show_label=False,
+                height=420,
+                type="filepath",
+            )
+
+            gr.HTML('<div class="section-label" style="margin-top:16px">Upscaled Preview</div>')
+            out_preview = gr.Image(
+                type="pil",
+                show_label=False,
+                height=200,
+            )
 
     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,
+        fn=run_upscale,
+        inputs=[inp_image, model_dropdown],
+        outputs=[slider, out_preview, dl_btn],
     )
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch(server_name="0.0.0.0", server_port=7860)