Update app.py

app.py

@@ -1,8 +1,27 @@
+# ────────────────────────────────────────────────────────
+# TorchVision compat shim (MUST be before importing basicsr)
+# Fixes: ModuleNotFoundError: torchvision.transforms.functional_tensor
+# ────────────────────────────────────────────────────────
+import sys, types
+try:
+    # If old path exists, do nothing
+    import torchvision.transforms.functional_tensor as _ft  # noqa: F401
+except Exception:
+    # Map to the new API location
+    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 gradio as gr
 import cv2
 import numpy
 import os
 import random
+
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.download_util import load_file_from_url
 
@@ -20,14 +39,11 @@ img_mode = "RGBA"
 # Utilities
 # ────────────────────────────────────────────────────────
 def rnd_string(x: int) -> str:
-    """Returns a string of 'x' random characters."""
     characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
-    result = "".join((random.choice(characters)) for _ in range(x))
-    return result
+    return "".join((random.choice(characters)) for _ in range(x))
 
 
 def reset():
-    """Resets the Image components and deletes the last processed image."""
     global last_file
     if last_file:
         try:
@@ -40,10 +56,6 @@ def reset():
 
 
 def has_transparency(img):
-    """
-    Check for transparency in a PIL image.
-    https://stackoverflow.com/questions/43864101/python-pil-check-if-image-is-transparent
-    """
     if img.info.get("transparency", None) is not None:
         return True
     if img.mode == "P":
@@ -59,19 +71,13 @@ def has_transparency(img):
 
 
 def image_properties(img):
-    """Return resolution & color mode of the input image; set global img_mode."""
     global img_mode
     if img:
-        if has_transparency(img):
-            img_mode = "RGBA"
-        else:
-            img_mode = "RGB"
-        properties = f"Resolution: Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
-        return properties
+        img_mode = "RGBA" if has_transparency(img) else "RGB"
+        return f"Resolution: Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
 
 
 def model_tip_text(model_name: str) -> str:
-    """Return human-friendly guidance for the chosen model."""
     tips = {
         "RealESRGAN_x4plus": (
             "**RealESRGAN_x4plus (4×)** — Best for photoreal images (portraits, landscapes). "
@@ -101,50 +107,40 @@ def model_tip_text(model_name: str) -> str:
 # Core upscaling
 # ────────────────────────────────────────────────────────
 def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
-    """Real-ESRGAN function to restore (and upscale) images with robust defaults."""
     if img is None:
         return
 
     # ----- Select backbone + weights -----
-    if model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-        netscale = 4
+    if model_name == 'RealESRGAN_x4plus':
+        model = RRDBNet(3, 3, 64, 23, 32, scale=4); netscale = 4
         file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
 
-    elif model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-        netscale = 4
+    elif model_name == 'RealESRNet_x4plus':
+        model = RRDBNet(3, 3, 64, 23, 32, scale=4); 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':  # x4 RRDBNet model with 6 blocks
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
-        netscale = 4
+    elif model_name == 'RealESRGAN_x4plus_anime_6B':
+        model = RRDBNet(3, 3, 64, 6, 32, scale=4); 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':  # x2 RRDBNet model
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
-        netscale = 2
+    elif model_name == 'RealESRGAN_x2plus':
+        model = RRDBNet(3, 3, 64, 23, 32, scale=2); netscale = 2
         file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
 
-    elif model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)
-        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
-        netscale = 4
-        # We'll ensure BOTH base and WDN weights exist; order matters for DNI.
+    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 are on disk -----
-    # For the general-x4v3 case we download both; for others single file is fine.
+    # ----- 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)
 
-    # Track model paths
     local_paths = []
     for url in file_url:
         fname = os.path.basename(url)
@@ -153,27 +149,22 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
             local_path = load_file_from_url(url=url, model_dir=weights_dir, progress=True)
         local_paths.append(local_path)
 
-    # Default path(s)
     if model_name == 'realesr-general-x4v3':
-        # Order: [base, wdn] then set DNI weights accordingly
         base_path = os.path.join(weights_dir, 'realesr-general-x4v3.pth')
-        wdn_path = os.path.join(weights_dir, 'realesr-general-wdn-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)
-        # Weight for WDN equals denoise_strength (cleaner); base gets the remainder
-        dni_weight = [1.0 - denoise_strength, denoise_strength]
+        dni_weight = [1.0 - denoise_strength, denoise_strength]  # base, WDN
     else:
         model_path = os.path.join(weights_dir, f"{model_name}.pth")
         dni_weight = None
 
     # ----- CUDA / precision / tiling -----
-    # Be defensive: cv2.cuda may not exist in CPU-only builds.
     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(
@@ -181,10 +172,10 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
         model_path=model_path,
         dni_weight=dni_weight,
         model=model,
-        tile=256,          # Safe VRAM default; increase if you have headroom
+        tile=256,
         tile_pad=10,
         pre_pad=10,
-        half=bool(use_cuda),  # FP16 on GPU
+        half=bool(use_cuda),
         gpu_id=gpu_id
     )
 
@@ -200,7 +191,7 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
             bg_upsampler=upsampler
         )
 
-    # ----- Convert PIL -> cv2 (handle RGB/RGBA) -----
+    # ----- 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)
@@ -218,7 +209,7 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
         print('Tip: If you hit CUDA OOM, try a smaller tile size (e.g., 128).')
         return None
 
-    # ----- cv2 -> RGBA/RGB for Gradio, also save -----
+    # ----- cv2 -> display ndarray, also save -----
     if output.ndim == 3 and output.shape[2] == 4:
         display_img = cv2.cvtColor(output, cv2.COLOR_BGRA2RGBA)
         extension = 'png'
@@ -234,7 +225,7 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
     except Exception as e:
         print("Save error:", e)
 
-    return display_img  # ndarray so Gradio displays immediately
+    return display_img
 
 
 # ────────────────────────────────────────────────────────
@@ -255,7 +246,7 @@ def main():
                         "RealESRGAN_x2plus",
                         "realesr-general-x4v3",
                     ],
-                    value="RealESRGAN_x4plus",  # photoreal default
+                    value="RealESRGAN_x4plus",
                     show_label=True
                 )
                 denoise_strength = gr.Slider(
@@ -268,7 +259,6 @@ def main():
                 )
                 face_enhance = gr.Checkbox(label="Face Enhancement (GFPGAN)", value=False)
 
-        # Model tips panel (auto-updates)
         model_tips = gr.Markdown(model_tip_text("RealESRGAN_x4plus"))
 
         with gr.Row():
@@ -281,7 +271,6 @@ def main():
             reset_btn = gr.Button("Remove images")
             restore_btn = gr.Button("Upscale")
 
-        # Event listeners:
         input_image.change(fn=image_properties, inputs=input_image, outputs=input_image_properties)
         model_name.change(fn=model_tip_text, inputs=model_name, outputs=model_tips)