Update app.py

app.py

@@ -46,6 +46,71 @@ from basicsr.utils.download_util import load_file_from_url
 from realesrgan import RealESRGANer
 from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 
+# Flag so UI can know if realesrgan is importable
+HAVE_REALESRGAN = True
+
+def build_rrdb(scale: int, num_block: int):
+    # simple wrapper to the imported RRDBNet class
+    return _RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=num_block, num_grow_ch=32, scale=scale)
+
+def _weights_dir() -> str:
+    ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+    wdir = os.path.join(ROOT_DIR, "weights")
+    os.makedirs(wdir, exist_ok=True)
+    return wdir
+
+def get_realesrganer(model_id: str, scale: int, tile: int, half: bool, device: str = "cpu") -> RealESRGANer:
+    """
+    model_id: one of {"x4plus", "x4plus-anime", "x2plus"}
+    Downloads weights if missing, builds the proper arch, and returns a RealESRGANer.
+    """
+    wdir = _weights_dir()
+
+    if model_id == "x4plus":
+        model = build_rrdb(scale=4, num_block=23)
+        netscale = 4
+        urls = ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"]
+        model_path = os.path.join(wdir, "RealESRGAN_x4plus.pth")
+        dni_weight = None
+    elif model_id == "x4plus-anime":
+        model = build_rrdb(scale=4, num_block=6)
+        netscale = 4
+        urls = ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"]
+        model_path = os.path.join(wdir, "RealESRGAN_x4plus_anime_6B.pth")
+        dni_weight = None
+    elif model_id == "x2plus":
+        model = build_rrdb(scale=2, num_block=23)
+        netscale = 2
+        urls = ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"]
+        model_path = os.path.join(wdir, "RealESRGAN_x2plus.pth")
+        dni_weight = None
+    else:
+        raise ValueError(f"Unknown model_id: {model_id}")
+
+    # ensure weights on disk
+    for url in urls:
+        fname = os.path.basename(url)
+        local_path = os.path.join(wdir, fname)
+        if not os.path.isfile(local_path):
+            load_file_from_url(url=url, model_dir=wdir, progress=True)
+
+    # Use GPU if visible; otherwise CPU
+    gpu_id = 0 if (device == "cuda") else None
+
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=tile or 256,
+        tile_pad=10,
+        pre_pad=10,
+        half=bool(half and device == "cuda"),
+        gpu_id=gpu_id
+    )
+    return upsampler
+
+
 # ─────────────────────────────────────────────────────────────
 # System checks & deps
 # ─────────────────────────────────────────────────────────────
@@ -58,12 +123,12 @@ FFMPEG = _which("ffmpeg")
 FFPROBE = _which("ffprobe")
 
 if not FFMPEG or not FFPROBE:
-MISSING_MSG = (
-    "⚠️ FFmpeg not found. Add a 'packages.txt' with exactly:\n"
-    "ffmpeg\n"
-    "libsm6\n"
-    "libxext6\n"
-    "Then restart the Space."
+    MISSING_MSG = (
+        "⚠️ FFmpeg not found. Add a 'packages.txt' with exactly:\n"
+        "ffmpeg\n"
+        "libsm6\n"
+        "libxext6\n"
+        "Then restart the Space."
 )
 else:
     MISSING_MSG = ""
@@ -299,9 +364,13 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
             channel_multiplier=2,
             bg_upsampler=upsampler
         )
+import random, string
+
+def rnd_string(n: int = 8) -> str:
+    return "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(n))
 
     # ----- PIL -> cv2 -----
-    cv_img = numpy.array(img)
+    cv_img = np.array(img)
     if cv_img.ndim == 3 and cv_img.shape[2] == 4:
         cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
     else:
@@ -417,9 +486,7 @@ def step1_extract(
             err = proc.stderr.read() if proc.stderr else ""
         except Exception:
             err = ""
-        return None, None, f"FFmpeg error or no frames produced.
-
-{err}", cmd_preview, last_html, None, None, None
+        return None, None, f"FFmpeg error or no frames produced.\\n\\n{err}", cmd_preview, last_html, None, None, None
 
     gallery = sample_paths(frames, 30)
     zip_path = work / "frames.zip"
@@ -475,7 +542,7 @@ def step2_upscale(
     
     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)
+    upsampler = get_realesrganer(model_id, scale, tile, half, device=device)
 
     work = Path(tempfile.mkdtemp(prefix="up_"))
     out_dir = work / "upscaled"; out_dir.mkdir(parents=True, exist_ok=True)
@@ -595,9 +662,8 @@ def step3_encode(frames_dir_state: str | None, prefix_state: str | None, orig_vi
             err = proc.stderr.read() if proc.stderr else ""
         except Exception:
             err = ""
-        return None, None, f"Encoding failed.
-
-{err}", last_html
+        return None, None, f"Encoding failed.\\n\\n{err}", last_html
+        
     return str(out_file), f"Video created: {out_file.name}", render_progress(100.0, "Encoding complete")
 
 # ───────────────── Quick Mode — one click: All frames → Upscale ×4 → MP4 (audio)