Update app.py

app.py

@@ -22,31 +22,32 @@ DRIVE_IDS = {
     "hresnet_x4": "15xmXXZNH2wMyeQv4ie5hagT7eWK9MgP6",   # placeholder = ESRGAN x2
 }
 
-# SRCNN .pth lives in the same HF Space repo (uploaded by user)
+# srcnn_x4.pth must be in the Space repo root (same folder as app.py)
 SRCNN_PTH = os.path.join(os.path.dirname(__file__), "srcnn_x4.pth")
 
 MODEL_LABELS = {
     "esrgan_x4":  "Real-ESRGAN ×4",
     "srcnn_x4":   "SRCNN ×4",
-    "hresnet_x4": "HResNet ×4",        # placeholder until real weights available
+    "hresnet_x4": "HResNet ×4",
 }
 
 MODEL_SCALES = {
     "esrgan_x4":  4,
     "srcnn_x4":   4,
-    "hresnet_x4": 2,                   # underlying model is ESRGAN x2
+    "hresnet_x4": 2,   # underlying model is ESRGAN x2 (placeholder)
 }
 
 
 # ===========================================================================
-# SRCNN architecture  (3-layer, matches the .pth weights)
+# SRCNN architecture — 3 conv layers, 1-channel (Y / grayscale) input
+# Your .pth was trained on grayscale, so num_channels=1 here.
 # ===========================================================================
 class SRCNN(nn.Module):
-    def __init__(self, num_channels: int = 3):
+    def __init__(self, num_channels: int = 1):
         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.conv1 = nn.Conv2d(num_channels, 64, kernel_size=9, padding=4)
+        self.conv2 = nn.Conv2d(64,           32, kernel_size=5, padding=2)
+        self.conv3 = nn.Conv2d(32, num_channels, kernel_size=5, padding=2)
         self.relu  = nn.ReLU(inplace=True)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -66,53 +67,56 @@ 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)
+    """Download ESRGAN ONNX from Drive via gdown (handles confirmation pages)."""
+    dest = os.path.join(CACHE_DIR, f"{key}.onnx")
     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)
+        ONNX_SESSIONS[key] = (sess, sess.get_inputs()[0])
         print(f"Loaded {MODEL_LABELS[key]} ✓")
     else:
-        print(f"[ERROR] Could not load {key} — file missing after download.")
+        print(f"[ERROR] {key} — file missing after download attempt.")
 
 
 def _load_srcnn_pth():
-    """Load SRCNN directly from the .pth file in the Space repo."""
+    """
+    Load SRCNN from .pth in the Space repo root.
+    The weights use 1-channel (grayscale / Y) input — confirmed by the
+    conv1.weight shape torch.Size([64, 1, 9, 9]) in the checkpoint.
+    Inference will convert RGB → YCbCr, enhance Y with SRCNN,
+    bicubic-upsample CbCr, then recompose back to RGB.
+    """
     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.")
+        print(f"[WARN] srcnn_x4.pth not found at {SRCNN_PTH} — SRCNN skipped.")
         return
-    model = SRCNN(num_channels=3)
+    model = SRCNN(num_channels=1)
     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]
+    for wrap_key in ("model", "state_dict", "params"):
+        if isinstance(state, dict) and wrap_key in state:
+            state = state[wrap_key]
     state = {k.replace("module.", ""): v for k, v in state.items()}
-    try:
-        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.load_state_dict(state, strict=True)
     model.eval()
     SRCNN_MODEL = model
-    print("Loaded SRCNN ×4 from .pth ✓")
+    print("Loaded SRCNN ×4 from .pth ✓  (grayscale/Y-channel model)")
 
 
 # Boot-time loading
-for k in ("esrgan_x4", "hresnet_x4"):
+for _k in ("esrgan_x4", "hresnet_x4"):
     try:
-        _load_esrgan_onnx(k)
-    except Exception as e:
-        print(f"[ERROR] {k}: {e}")
+        _load_esrgan_onnx(_k)
+    except Exception as _e:
+        print(f"[ERROR] {_k}: {_e}")
 
-_load_srcnn_pth()
+try:
+    _load_srcnn_pth()
+except Exception as _e:
+    print(f"[ERROR] SRCNN: {_e}")
 
 
 # ===========================================================================
@@ -120,41 +124,58 @@ _load_srcnn_pth()
 # ===========================================================================
 
 def _onnx_tile(sess, meta, tile: np.ndarray) -> np.ndarray:
-    """HWC float32 [0,1] → HWC float32 [0,1]."""
+    """HWC float32 [0,1] in → HWC float32 out."""
     patch = tile.transpose(2, 0, 1)[None, ...]
     out   = sess.run(None, {meta.name: patch})[0]
     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)
+    """
+    Enhance a single RGB tile using the grayscale SRCNN model.
+    Strategy: split into YCbCr → SRCNN on Y → bicubic CbCr → recompose RGB.
+    tile: HWC float32 [0, 1]
+    returns: HWC float32 [0, 1] at scale× resolution
+    """
+    tile_uint8  = (np.clip(tile, 0, 1) * 255).round().astype(np.uint8)
+    tile_pil    = Image.fromarray(tile_uint8)
+    tile_ycbcr  = tile_pil.convert("YCbCr")
+    y_pil, cb_pil, cr_pil = tile_ycbcr.split()
+
+    orig_w, orig_h = tile_pil.size
+    up_w, up_h     = orig_w * scale, orig_h * scale
+
+    # Upsample CbCr channels with bicubic (no SRCNN needed there)
+    cb_up = cb_pil.resize((up_w, up_h), Image.BICUBIC)
+    cr_up = cr_pil.resize((up_w, up_h), Image.BICUBIC)
+
+    # Bicubic upsample Y, then refine with SRCNN
+    y_arr = np.array(y_pil).astype(np.float32) / 255.0          # (H, W)
+    y_t   = torch.from_numpy(y_arr).unsqueeze(0).unsqueeze(0)   # (1, 1, H, W)
+    y_up  = F.interpolate(y_t, size=(up_h, up_w), mode="bicubic", align_corners=False)
+
     with torch.no_grad():
-        out = SRCNN_MODEL(up)
-    return out.squeeze(0).permute(1, 2, 0).numpy()
+        y_refined = SRCNN_MODEL(y_up)                            # (1, 1, H*s, W*s)
+
+    y_out = (y_refined.squeeze().numpy() * 255.0).clip(0, 255).round().astype(np.uint8)
+    y_up_pil = Image.fromarray(y_out, mode="L")
+
+    # Recompose YCbCr → RGB
+    out_rgb = Image.merge("YCbCr", [y_up_pil, cb_up, cr_up]).convert("RGB")
+    return np.array(out_rgb).astype(np.float32) / 255.0
 
 
 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
+    """Tile-based upscale dispatcher for ONNX (ESRGAN) and torch (SRCNN)."""
     if model_key == "srcnn_x4" and SRCNN_MODEL is None:
-        raise RuntimeError("SRCNN model not loaded.")
+        raise RuntimeError("SRCNN model not loaded — check that srcnn_x4.pth is in the repo root.")
     if model_key in ("esrgan_x4", "hresnet_x4") and model_key not in ONNX_SESSIONS:
-        raise RuntimeError(f"{MODEL_LABELS[model_key]} not loaded.")
+        raise RuntimeError(f"{MODEL_LABELS[model_key]} failed to load at startup.")
 
     scale = MODEL_SCALES[model_key]
+    TILE  = 128   # LR tile size (consistent across all models)
 
-    # 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
+    # Cap input size to avoid OOM
     w, h = input_img.size
     if w > max_dim or h > max_dim:
         factor    = max_dim / float(max(w, h))
@@ -174,8 +195,8 @@ def upscale(input_img: Image.Image, model_key: str, max_dim: int = 1024) -> Imag
 
     for i in range(tiles_h):
         for j in range(tiles_w):
-            y0, x0  = i * TILE, j * TILE
-            tile     = arr_pad[y0:y0 + TILE, x0:x0 + TILE]
+            y0, x0 = i * TILE, j * TILE
+            tile    = arr_pad[y0:y0 + TILE, x0:x0 + TILE]
 
             if model_key == "srcnn_x4":
                 up_tile = _srcnn_tile(tile, scale=scale)
@@ -190,25 +211,6 @@ def upscale(input_img: Image.Image, model_key: str, max_dim: int = 1024) -> Imag
     return Image.fromarray((final * 255.0).round().astype(np.uint8))
 
 
-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
-
-
 # ===========================================================================
 # Gradio callback
 # ===========================================================================
@@ -217,18 +219,28 @@ def run_upscale(input_img: Image.Image, model_name: str):
     if input_img is None:
         return None, None, None
 
-    # Map display label → key
-    key = next(k for k, v in MODEL_LABELS.items() if v == model_name)
+    key    = next(k for k, v in MODEL_LABELS.items() if v == model_name)
+    result = upscale(input_img, key)
+
+    # Resize original to same dimensions as output for the slider
+    up_w, up_h   = result.size
+    orig_resized = input_img.resize((up_w, up_h), Image.LANCZOS).convert("RGB")
 
-    result   = upscale(input_img, key)
-    orig_path, up_path = make_comparison_png(input_img, result)
+    # Save both as temp files for ImageSlider
+    tmp_orig = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    orig_resized.save(tmp_orig.name)
+    tmp_orig.close()
 
-    # Also save download copy
-    dl_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-    result.save(dl_tmp.name, format="PNG")
-    dl_tmp.close()
+    tmp_up = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    result.save(tmp_up.name)
+    tmp_up.close()
 
-    return (orig_path, up_path), result, dl_tmp.name
+    # Separate download copy
+    tmp_dl = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    result.save(tmp_dl.name, format="PNG")
+    tmp_dl.close()
+
+    return (tmp_orig.name, tmp_up.name), result, tmp_dl.name
 
 
 # ===========================================================================
@@ -237,59 +249,28 @@ def run_upscale(input_img: Image.Image, model_name: str):
 
 css = """
 @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 { 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, #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;
+    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;
 }
 #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;
+    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;
+    font-size: 0.75rem; font-weight: 700; letter-spacing: 0.1em;
+    text-transform: uppercase; color: #999; margin-bottom: 6px;
 }
 """
 
-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 Upscaler") as demo:
@@ -303,7 +284,6 @@ with gr.Blocks(css=css, title="SpectraGAN Upscaler") as demo:
 
     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)
@@ -314,32 +294,22 @@ with gr.Blocks(css=css, title="SpectraGAN Upscaler") as demo:
                 value=dropdown_choices[0],
                 show_label=False,
             )
-
             run_btn = gr.Button("⚡ Upscale", elem_id="run-btn")
-
-            dl_btn = gr.DownloadButton(
+            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>')
+            gr.HTML('<div class="section-label">Before / After — drag to compare</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,
-            )
+            out_preview = gr.Image(type="pil", show_label=False, height=200)
 
     run_btn.click(
         fn=run_upscale,