Update app.py

app.py

@@ -9,8 +9,22 @@
 # - Prefix defaults to input video filename if left blank
 # =============================
 
+
+import sys, types
+try:
+    import torchvision.transforms.functional_tensor as _ft  # noqa: F401
+except Exception:
+    from torchvision.transforms import functional as _F
+    _mod = types.ModuleType("torchvision.transforms.functional_tensor")
+    _mod.rgb_to_grayscale = _F.rgb_to_grayscale
+    sys.modules["torchvision.transforms.functional_tensor"] = _mod
+
+# ────────────────────────────────────────────────────────
+# Standard imports
+# ────────────────────────────────────────────────────────
 import os
 import re
+import cv2
 import json
 import math
 import time
@@ -18,6 +32,7 @@ import shutil
 import zipfile
 import tempfile
 import subprocess
+import inspect
 from pathlib import Path
 from typing import List, Optional, Tuple
 
@@ -25,6 +40,12 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
+from basicsr.archs.rrdbnet_arch import RRDBNet as _RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+
+from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+
 # ─────────────────────────────────────────────────────────────
 # System checks & deps
 # ─────────────────────────────────────────────────────────────
@@ -47,19 +68,30 @@ MISSING_MSG = (
 else:
     MISSING_MSG = ""
 
-# Try to import Real-ESRGAN stack
-try:
-    from realesrgan import RealESRGANer
-    from basicsr.archs.rrdbnet_arch import RRDBNet
-    HAVE_REALESRGAN = True
-except Exception as e:
-    HAVE_REALESRGAN = False
-    REAL_ERR = str(e)
-
 # ─────────────────────────────────────────────────────────────
 # Helpers
 # ─────────────────────────────────────────────────────────────
 
+    # Map UI model names (demo) to our internal model IDs
+def map_ui_model_to_internal(ui_name: str) -> str:
+    mapping = {
+        "RealESRGAN_x4plus": "x4plus",
+        "RealESRGAN_x4plus_anime_6B": "x4plus-anime",
+        "RealESRGAN_x2plus": "x2plus",
+        # Unsupported in our current RRDBNet wiring – fallback:
+        "RealESRNet_x4plus": "x4plus",
+        "realesr-general-x4v3": "x4plus",
+    }
+    return mapping.get(ui_name, "x4plus")
+
+def clamp_scale_for_model(outscale: int, model_id: str) -> int:
+    # Our current models are ×2 or ×4 only.
+    if model_id == "x2plus":
+        return 2
+    # For x4plus / x4plus-anime, force 4 (ignore 5–6)
+    return 4
+
+
 def sample_paths(paths: List[Path] | List[str], n: int = 30) -> List[str]:
     """Return up to n items sampled evenly across the list, preserving order (as strings)."""
     if not paths:
@@ -185,30 +217,124 @@ def build_ffmpeg_extract(
     return cmd
 
 
-def get_realesrganer(model_name: str, scale: int, tile: int, half: bool, device: str = "cpu"):
-    if not HAVE_REALESRGAN:
-        raise RuntimeError("realesrgan is not installed. See requirements.txt (realesrgan, basicsr, torch, numpy, scipy, scikit-image).")
-    if model_name in ("x4plus", "x4plus-anime"):
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-        model_scale = 4
-    elif model_name == "x2plus":
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
-        model_scale = 2
+def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
+    if img is None:
+        return
+
+    # ----- Select backbone + weights -----
+    if model_name == 'RealESRGAN_x4plus':
+        model = build_rrdb(scale=4, num_block=23); netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
+
+    elif model_name == 'RealESRNet_x4plus':
+        model = build_rrdb(scale=4, num_block=23); netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
+
+    elif model_name == 'RealESRGAN_x4plus_anime_6B':
+        model = build_rrdb(scale=4, num_block=6); netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+
+    elif model_name == 'RealESRGAN_x2plus':
+        model = build_rrdb(scale=2, num_block=23); netscale = 2
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
+
+    elif model_name == 'realesr-general-x4v3':
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu'); netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth'
+        ]
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+
+    # ----- Ensure weights on disk -----
+    ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+    weights_dir = os.path.join(ROOT_DIR, 'weights')
+    os.makedirs(weights_dir, exist_ok=True)
+
+    for url in file_url:
+        fname = os.path.basename(url)
+        local_path = os.path.join(weights_dir, fname)
+        if not os.path.isfile(local_path):
+            load_file_from_url(url=url, model_dir=weights_dir, progress=True)
+
+    if model_name == 'realesr-general-x4v3':
+        base_path = os.path.join(weights_dir, 'realesr-general-x4v3.pth')
+        wdn_path  = os.path.join(weights_dir, 'realesr-general-wdn-x4v3.pth')
+        model_path = [base_path, wdn_path]
+        denoise_strength = float(denoise_strength)
+        dni_weight = [1.0 - denoise_strength, denoise_strength]  # base, WDN
     else:
-        raise ValueError("Unknown Real-ESRGAN model")
-    if scale not in (2, 4):
-        scale = model_scale
+        model_path = os.path.join(weights_dir, f"{model_name}.pth")
+        dni_weight = None
+
+    # ----- CUDA / precision / tiling -----
+    use_cuda = False
+    try:
+        use_cuda = hasattr(cv2, "cuda") and cv2.cuda.getCudaEnabledDeviceCount() > 0
+    except Exception:
+        use_cuda = False
+    gpu_id = 0 if use_cuda else None
+
     upsampler = RealESRGANer(
-        scale=model_scale,
-        model_path=None,
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
         model=model,
-        tile=tile or 0,
+        tile=256,       # VRAM-safe default; lower to 128 if OOM
         tile_pad=10,
-        pre_pad=0,
-        half=half,
-        device=device,
+        pre_pad=10,
+        half=bool(use_cuda),
+        gpu_id=gpu_id
     )
-    return upsampler
+
+    # ----- Optional face enhancement -----
+    face_enhancer = None
+    if face_enhance:
+        from gfpgan import GFPGANer
+        face_enhancer = GFPGANer(
+            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+            upscale=outscale,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=upsampler
+        )
+
+    # ----- PIL -> cv2 -----
+    cv_img = numpy.array(img)
+    if cv_img.ndim == 3 and cv_img.shape[2] == 4:
+        cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
+    else:
+        cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGB2BGR)
+
+    # ----- Enhance -----
+    try:
+        if face_enhancer:
+            _, _, output = face_enhancer.enhance(cv_img, has_aligned=False, only_center_face=False, paste_back=True)
+        else:
+            output, _ = upsampler.enhance(cv_img, outscale=int(outscale))
+    except RuntimeError as error:
+        print('Error', error)
+        print('Tip: If you hit CUDA OOM, try a smaller tile size (e.g., 128).')
+        return None
+
+    # ----- cv2 -> display ndarray, also save -----
+    if output.ndim == 3 and output.shape[2] == 4:
+        display_img = cv2.cvtColor(output, cv2.COLOR_BGRA2RGBA)
+        extension = 'png'
+    else:
+        display_img = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+        extension = 'jpg'
+
+    out_filename = f"output_{rnd_string(8)}.{extension}"
+    try:
+        cv2.imwrite(out_filename, output)
+        global last_file
+        last_file = out_filename
+    except Exception as e:
+        print("Save error:", e)
+
+    return display_img
 
 
 def render_progress(pct: float, label: str = "") -> str:
@@ -323,16 +449,15 @@ def save_uploaded_images(files: List[gr.File] | None, prefix: str = "upload") ->
 
 def step2_upscale(
     frames_list: List[str] | None,
-    model_name: str,
-    scale: int,
+    ui_model_name: str,
+    outscale: int,
     tile: int,
     precision: str,
     prog_html: str,
     uploaded_imgs: List[gr.File] | None,
+    denoise_strength: float = 0.5, # NEW: accepted but currently unused
+    face_enhance: bool = False,    # NEW: accepted but currently unused (needs GFPGAN)
 ):
-    if not HAVE_REALESRGAN:
-        msg = "Real-ESRGAN not available. Ensure requirements.txt includes: --prefer-binary, numpy==1.26.4, scipy==1.11.4, scikit-image==0.22.0, opencv-python-headless, torch==2.2.2, realesrgan==0.3.0, basicsr==1.4.2, pillow, gradio."
-        return None, None, msg, prog_html
 
     # decide source: uploaded images take priority, else frames from step 1
     if uploaded_imgs and len(uploaded_imgs) > 0:
@@ -344,6 +469,10 @@ def step2_upscale(
     if not src_paths:
         return None, None, "No images provided. Upload files or run Step 1 first.", prog_html
 
+    # Map demo model -> internal, clamp scale to supported values
+    model_id = map_ui_model_to_internal(ui_model_name)
+    scale = clamp_scale_for_model(int(outscale or 4), model_id)
+    
     device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
     half = (precision == "half") and (device == "cuda")
     upsampler = get_realesrganer(model_name, scale, tile, half, device=device)
@@ -377,6 +506,9 @@ def step2_upscale(
         for p in up_paths:
             zf.write(p, p.name)
 
+    detail = (f"Upscaled: {len(up_paths)} | Model: {ui_model_name}→{model_id} | "
+              f"Scale: x{scale} | Tile: {tile} | Precision: {precision}")
+    
     return gallery, str(zip_path), f"Upscaled: {len(up_paths)}", render_progress(100.0, "Upscaling complete")
 
 # ───────────────── Encode (Step 3) — supports uploaded frames/ZIP & optional audio source
@@ -649,16 +781,46 @@ def build_ui():
             # STEP 2 — Upscale
             with gr.Tab("Step 2 · Upscale Frames"):
                 if not HAVE_REALESRGAN:
-                    gr.Markdown("⚠️ Upscaling disabled. Install dependencies in requirements.txt (see notes in code). Error: " + (REAL_ERR if 'REAL_ERR' in globals() else ""))
+                    gr.Markdown("⚠️ Upscaling disabled. Install dependencies in requirements.txt (realesrgan, basicsr, torch, etc.).")
+                
                 gr.Markdown("Use frames from Step 1 **or** upload images below.")
-                imgs_override = gr.Files(label="Upload images to upscale (JPG/PNG)", file_types=[".jpg", ".jpeg", ".png"], type="filepath")
+                imgs_override = gr.Files(
+                    label="Upload images to upscale (JPG/PNG)",
+                    file_types=[".jpg", ".jpeg", ".png"],
+                    type="filepath"
+                )
+                
+                with gr.Accordion("Upscaling options", open=True):
+                    with gr.Row():
+                        ui_model_name = gr.Dropdown(
+                            label="Upscaler model",
+                            choices=[
+                                "RealESRGAN_x4plus",
+                                "RealESRNet_x4plus",
+                                "RealESRGAN_x4plus_anime_6B",
+                                "RealESRGAN_x2plus",
+                                "realesr-general-x4v3",
+                            ],
+                            value="RealESRGAN_x4plus",
+                            show_label=True
+                        )
+                        denoise_strength = gr.Slider(
+                            label="Denoise Strength (only for realesr-general-x4v3)",
+                            minimum=0, maximum=1, step=0.1, value=0.5
+                        )
+                        outscale = gr.Slider(
+                            label="Resolution upscale",
+                            minimum=1, maximum=6, step=1, value=4, show_label=True
+                        )
+                        face_enhance = gr.Checkbox(label="Face Enhancement (GFPGAN)", value=False)
+                
                 with gr.Row():
-                    model_name = gr.Dropdown(["x4plus", "x4plus-anime", "x2plus"], value="x4plus", label="Model")
-                    scale = gr.Dropdown([2, 4], value=4, label="Output scale")
                     tile = gr.Number(value=0, label="Tile size (0 = auto)")
                     precision = gr.Dropdown(["auto", "half", "full"], value="auto", label="Precision (GPU=half, CPU=full)")
+                
                 with gr.Row():
                     btn_upscale = gr.Button("Step 2: Upscale", variant="primary")
+                
                 prog2 = gr.HTML(render_progress(0.0, "Idle"))
                 gallery_up = gr.Gallery(label="Upscaled preview (30 sampled)", columns=6, height=480)
                 zip_up = gr.File(label="Download upscaled ZIP")
@@ -666,7 +828,7 @@ def build_ui():
 
                 btn_upscale.click(
                     step2_upscale,
-                    inputs=[frames_state, model_name, scale, tile, precision, prog2, imgs_override],
+                    inputs=[frames_state, ui_model_name, outscale, tile, precision, prog2, imgs_override, denoise_strength, face_enhance],
                     outputs=[gallery_up, zip_up, details2, prog2],
                 )