Update app.py

app.py

@@ -9,72 +9,187 @@ from basicsr.utils.download_util import load_file_from_url
 from realesrgan import RealESRGANer
 from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 
-
+# ────────────────────────────────────────────────────────
+# Globals
+# ────────────────────────────────────────────────────────
 last_file = None
 img_mode = "RGBA"
 
 
-def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
-    """Real-ESRGAN function to restore (and upscale) images.
+# ────────────────────────────────────────────────────────
+# 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
+
+
+def reset():
+    """Resets the Image components and deletes the last processed image."""
+    global last_file
+    if last_file:
+        try:
+            print(f"Deleting {last_file} ...")
+            os.remove(last_file)
+        except Exception as e:
+            print("Delete error:", e)
+        last_file = None
+    return gr.update(value=None), gr.update(value=None)
+
+
+def has_transparency(img):
     """
-    if not 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":
+        transparent = img.info.get("transparency", -1)
+        for _, index in img.getcolors():
+            if index == transparent:
+                return True
+    elif img.mode == "RGBA":
+        extrema = img.getextrema()
+        if extrema[3][0] < 255:
+            return True
+    return False
+
+
+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
+
+
+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). "
+            "Balanced detail recovery. Good default for Flux realism."
+        ),
+        "RealESRNet_x4plus": (
+            "**RealESRNet_x4plus (4×)** — Softer but great on noisy/compressed sources "
+            "(old JPEGs, screenshots)."
+        ),
+        "RealESRGAN_x4plus_anime_6B": (
+            "**RealESRGAN_x4plus_anime_6B (4×)** — For anime/illustrations/line art only. "
+            "Not recommended for real-life photos."
+        ),
+        "RealESRGAN_x2plus": (
+            "**RealESRGAN_x2plus (2×)** — Faster, lighter 2× cleanup when you don't need 4×."
+        ),
+        "realesr-general-x4v3": (
+            "**realesr-general-x4v3 (4×)** — Versatile mixed-content model with adjustable denoise. "
+            "**Denoise Strength** slider only affects this model (blends with the WDN variant). "
+            "Try 0.3–0.5 for slightly cleaner, sharper results."
+        ),
+    }
+    return tips.get(model_name, "")
+
+
+# ────────────────────────────────────────────────────────
+# 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
 
-    # Define model parameters
+    # ----- 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
         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
         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
         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
         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.
         file_url = [
-            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
-            '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-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth'
         ]
 
-    # Determine model paths
-    model_path = os.path.join('weights', model_name + '.pth')
-    if not os.path.isfile(model_path):
-        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
-        for url in file_url:
-            # model_path will be updated
-            model_path = load_file_from_url(
-                url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
-
-    # Use dni to control the denoise strength
-    dni_weight = None
-    if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
-        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
-        model_path = [model_path, wdn_model_path]
-        dni_weight = [denoise_strength, 1 - denoise_strength]
-
-    # Restorer Class
+    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.
+    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)
+        local_path = os.path.join(weights_dir, fname)
+        if not os.path.isfile(local_path):
+            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')
+        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]
+    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(
         scale=netscale,
         model_path=model_path,
         dni_weight=dni_weight,
         model=model,
-        tile=0,
+        tile=256,          # Safe VRAM default; increase if you have headroom
         tile_pad=10,
         pre_pad=10,
-        half=False,
-        gpu_id=None
+        half=bool(use_cuda),  # FP16 on GPU
+        gpu_id=gpu_id
     )
 
-    # Use GFPGAN for face enhancement
+    # ----- Optional face enhancement -----
+    face_enhancer = None
     if face_enhance:
         from gfpgan import GFPGANer
         face_enhancer = GFPGANer(
@@ -82,140 +197,105 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
             upscale=outscale,
             arch='clean',
             channel_multiplier=2,
-            bg_upsampler=upsampler)
+            bg_upsampler=upsampler
+        )
 
-    # Convert the input PIL image to cv2 image, so that it can be processed by realesrgan
+    # ----- Convert PIL -> cv2 (handle RGB/RGBA) -----
     cv_img = numpy.array(img)
-    img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
+    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)
 
-    # Apply restoration
+    # ----- Enhance -----
     try:
-        if face_enhance:
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+        if face_enhancer:
+            _, _, output = face_enhancer.enhance(cv_img, has_aligned=False, only_center_face=False, paste_back=True)
         else:
-            output, _ = upsampler.enhance(img, outscale=outscale)
+            output, _ = upsampler.enhance(cv_img, outscale=int(outscale))
     except RuntimeError as error:
         print('Error', error)
-        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
+        print('Tip: If you hit CUDA OOM, try a smaller tile size (e.g., 128).')
+        return None
+
+    # ----- cv2 -> RGBA/RGB for Gradio, also save -----
+    if output.ndim == 3 and output.shape[2] == 4:
+        display_img = cv2.cvtColor(output, cv2.COLOR_BGRA2RGBA)
+        extension = 'png'
     else:
-        # Save restored image and return it to the output Image component
-        if img_mode == 'RGBA':  # RGBA images should be saved in png format
-            extension = 'png'
-        else:
-            extension = 'jpg'
+        display_img = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+        extension = 'jpg'
 
-        out_filename = f"output_{rnd_string(8)}.{extension}"
+    out_filename = f"output_{rnd_string(8)}.{extension}"
+    try:
         cv2.imwrite(out_filename, output)
         global last_file
         last_file = out_filename
-        return out_filename
-
-
-def rnd_string(x):
-    """Returns a string of 'x' random characters
-    """
-    characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
-    result = "".join((random.choice(characters)) for i in range(x))
-    return result
-
-
-def reset():
-    """Resets the Image components of the Gradio interface and deletes
-    the last processed image
-    """
-    global last_file
-    if last_file:
-        print(f"Deleting {last_file} ...")
-        os.remove(last_file)
-        last_file = None
-    return gr.update(value=None), gr.update(value=None)
-
-
-def has_transparency(img):
-    """This function works by first checking to see if a "transparency" property is defined
-    in the image's info -- if so, we return "True". Then, if the image is using indexed colors
-    (such as in GIFs), it gets the index of the transparent color in the palette
-    (img.info.get("transparency", -1)) and checks if it's used anywhere in the canvas
-    (img.getcolors()). If the image is in RGBA mode, then presumably it has transparency in
-    it, but it double-checks by getting the minimum and maximum values of every color channel
-    (img.getextrema()), and checks if the alpha channel's smallest value falls below 255.
-    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":
-        transparent = img.info.get("transparency", -1)
-        for _, index in img.getcolors():
-            if index == transparent:
-                return True
-    elif img.mode == "RGBA":
-        extrema = img.getextrema()
-        if extrema[3][0] < 255:
-            return True
-    return False
+    except Exception as e:
+        print("Save error:", e)
 
-
-def image_properties(img):
-    """Returns the dimensions (width and height) and color mode of the input image and
-    also sets the global img_mode variable to be used by the realesrgan function
-    """
-    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
+    return display_img  # ndarray so Gradio displays immediately
 
 
+# ────────────────────────────────────────────────────────
+# UI
+# ────────────────────────────────────────────────────────
 def main():
-    # Gradio Interface
     with gr.Blocks(title="Real-ESRGAN Gradio Demo", theme="ParityError/Interstellar") as demo:
+        gr.Markdown("## Image Upscaler")
 
-        gr.Markdown(
-            """ Image Upscaler
-        """
-        )
-
-        with gr.Accordion("Upscaling option"):
+        with gr.Accordion("Upscaling options", open=True):
             with gr.Row():
-                model_name = gr.Dropdown(label="Upscaler model",
-                                         choices=["RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B",
-                                                  "RealESRGAN_x2plus", "realesr-general-x4v3"],
-                                         value="RealESRGAN_x4plus_anime_6B", show_label=True)
-                denoise_strength = gr.Slider(label="Denoise Strength",
-                                             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)",
+                model_name = gr.Dropdown(
+                    label="Upscaler model",
+                    choices=[
+                        "RealESRGAN_x4plus",
+                        "RealESRNet_x4plus",
+                        "RealESRGAN_x4plus_anime_6B",
+                        "RealESRGAN_x2plus",
+                        "realesr-general-x4v3",
+                    ],
+                    value="RealESRGAN_x4plus",  # photoreal default
+                    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)
+
+        # Model tips panel (auto-updates)
+        model_tips = gr.Markdown(model_tip_text("RealESRGAN_x4plus"))
+
         with gr.Row():
             with gr.Group():
                 input_image = gr.Image(label="Input Image", type="pil", image_mode="RGBA")
                 input_image_properties = gr.Textbox(label="Image Properties", max_lines=1)
             output_image = gr.Image(label="Output Image", image_mode="RGBA")
+
         with gr.Row():
             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)
-        restore_btn.click(fn=realesrgan,
-                          inputs=[input_image, model_name, denoise_strength, face_enhance, outscale],
-                          outputs=output_image)
-        reset_btn.click(fn=reset, inputs=[], outputs=[output_image, input_image])
-        # reset_btn.click(None, inputs=[], outputs=[input_image], _js="() => (null)\n")
-        # Undocumented method to clear a component's value using Javascript
+        model_name.change(fn=model_tip_text, inputs=model_name, outputs=model_tips)
 
-        gr.Markdown(
-            """
-        """
+        restore_btn.click(
+            fn=realesrgan,
+            inputs=[input_image, model_name, denoise_strength, face_enhance, outscale],
+            outputs=output_image
         )
+        reset_btn.click(fn=reset, inputs=[], outputs=[output_image, input_image])
+
+        gr.Markdown("")  # spacer
 
     demo.launch()
 
 
 if __name__ == "__main__":
-    main()
+    main()