Update app.py

convert to app to API

app.py

@@ -1,162 +1,90 @@
-import os
+#######################################################################################
+#
+# MIT License
+#
+# Copyright (c) [2025] [leonelhs@gmail.com]
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#######################################################################################
+
+# This file implements an API endpoint RealESRGAN system.
+# It provides functionality to enhances an image by upscaling it 4 times its original size .
+
+
+# Source code is based on or inspired by several projects.
+# For more details and proper attribution, please refer to the following resources:
+#
+# - [Real-ESRGAN] - [https://github.com/xinntao/Real-ESRGAN]
+
 
-import gradio as gr
 import torch
+import gradio as gr
 from basicsr.archs.rrdbnet_arch import RRDBNet
-from basicsr.archs.srvgg_arch import SRVGGNetCompact
-from gfpgan.utils import GFPGANer
 from huggingface_hub import hf_hub_download
-from realesrgan.utils import RealESRGANer
+from realesrgan import RealESRGANer
 
 REALESRGAN_REPO_ID = 'leonelhs/realesrgan'
-GFPGAN_REPO_ID = 'leonelhs/gfpgan'
-
-os.system("pip freeze")
-
-
-def showGPU():
-    if torch.cuda.is_available():
-        devices = torch.cuda.device_count()
-        current = torch.cuda.current_device()
-        return f"Running on GPU:{current} of {devices} total devices"
-    return "Running on CPU"
-
 
-def download_model_gfpgan(file):
-    return hf_hub_download(repo_id=GFPGAN_REPO_ID, filename=file)
+MODELS = {
+    'RealESRGAN X4+':'RealESRGAN_x4plus.pth',
+    'RealESRNet X4+':'RealESRNet_x4plus.pth',
+    'AI-Forever X4+':'AI-Forever_x4plus.pth'
+}
 
-
-def download_model_realesrgan(file):
+def download_model(file):
     return hf_hub_download(repo_id=REALESRGAN_REPO_ID, filename=file)
 
+def predict(img, method="RealESRGAN_x4plus.pth"):
+    """
+        Enhances an image by upscaling it 4 times its original size.
+        Parameters:
+            img (string):  Path to the input image file.
+            method (string): Upscaling method to use (RealESRGAN, RealESRNet, AI-Forever).
+        Returns:
+            path: File path to the enhanced, upscaled image.
+    """
 
-def select_upsampler(version, netscale=4):
-    model = None
-    dni_weight = None
-
-    version = version + ".pth"
-    model_path = download_model_realesrgan(version)
-
-    if version == 'RealESRGAN_x4plus.pth':  # 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)
-
-    if version == 'RealESRNet_x4plus.pth':  # 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)
-
-    if version == 'AI-Forever_x4plus.pth':  # 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)
-
-    if version == 'RealESRGAN_x4plus_anime_6B.pth':  # 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)
+    model_path = download_model(method)
 
-    if version == 'RealESRGAN_x2plus.pth':  # 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  # This is
+    # 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)
+    half =  True if torch.cuda.is_available() else False
 
-    if version == 'AI-Forever_x2plus.pth':  # 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  # This is
-
-    if version == 'realesr-animevideov3.pth':  # x4 VGG-style model (XS size)
-        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
-
-    if version == 'realesr-general-x4v3.pth':  # 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')
-        model_path = [
-            download_model_realesrgan("realesr-general-x4v3.pth"),
-            download_model_realesrgan("realesr-general-wdn-x4v3.pth")
-        ]
-        dni_weight = [0.2, 0.8]
-
-    half = True if torch.cuda.is_available() else False
-
-    return RealESRGANer(
-        scale=netscale,
+    output , _ = RealESRGANer(
+        scale=4,
         model_path=model_path,
-        dni_weight=dni_weight,
         model=model,
-        tile=0,
-        tile_pad=10,
-        pre_pad=0,
-        half=half,
-        gpu_id=0)
-
-
-def select_face_enhancer(version, scale, upsampler):
-    if 'v1.2' in version:
-        model_path = download_model_gfpgan('GFPGANv1.2.pth')
-        return GFPGANer(
-            model_path=model_path, upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
-    elif 'v1.3' in version:
-        model_path = download_model_gfpgan('GFPGANv1.3.pth')
-        return GFPGANer(
-            model_path=model_path, upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
-    elif 'v1.4' in version:
-        model_path = download_model_gfpgan('GFPGANv1.4.pth')
-        return GFPGANer(
-            model_path=model_path, upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
-    elif 'RestoreFormer' in version:
-        model_path = download_model_gfpgan('RestoreFormer.pth')
-        return GFPGANer(
-            model_path=model_path, upscale=scale, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
-
-
-def predict(image, version_upsampler, version_enhancer, scale):
-    scale = int(scale)
-
-    upsampler = select_upsampler(version_upsampler)
-
-    if "No additional" not in version_enhancer:
-        face_enhancer = select_face_enhancer(version_enhancer, scale, upsampler)
-        _, _, output = face_enhancer.enhance(image, has_aligned=False, only_center_face=False, paste_back=True)
-    else:
-        output, _ = upsampler.enhance(image, outscale=scale)
-
-    log = f"General enhance version: {version_upsampler}\n " \
-          f"Face enhance version: {version_enhancer} \n " \
-          f"Scale:{scale} \n {showGPU()}"
-
-    return output, log
-
+        half=half).enhance(img)
 
-title = "Super Face"
-description = r"""
-<b>Practical Image Restoration Algorithm based on Real-ESRGAN, GFPGAN</b>
-"""
-article = r"""
-<center><span>xintao.wang@outlook.com or xintaowang@tencent.com</span></center>
-</br>
-<center><a href='https://github.com/TencentARC/GFPGAN' target='_blank'>Github Repo ⭐ </a> are welcome</center>
-"""
+    return output
 
-demo = gr.Interface(
+app = gr.Interface(
     predict, [
-        gr.Image(type="numpy", label="Input"),
-        gr.Dropdown([
-            'RealESRGAN_x2plus',
-            'RealESRGAN_x4plus',
-            'RealESRNet_x4plus',
-            'AI-Forever_x2plus',
-            'AI-Forever_x4plus',
-            'RealESRGAN_x4plus_anime_6B',
-            'realesr-animevideov3',
-            'realesr-general-x4v3'],
-            type="value", value='RealESRGAN_x4plus', label='General restoration algorithm', info="version"),
-        gr.Dropdown([
-            'No additional face process',
-            'GFPGANv1.2',
-            'GFPGANv1.3',
-            'GFPGANv1.4',
-            'RestoreFormer'],
-            type="value", value='No additional face process', label='Special face restoration algorithm',
-            info="version"),
-        gr.Dropdown(["1", "2", "3", "4"], value="2", label="Rescaling factor")
+        gr.Image(type="numpy", label="Image input"),
+        gr.Dropdown(list(MODELS.items()),
+                type="value", value='RealESRGAN_x4plus.pth', label='Enhancer model'),
     ], [
-        gr.Image(type="numpy", label="Output", interactive=False),
-        gr.Textbox(label="log info")
+        gr.Image(type="numpy", label="Image enhanced")
     ],
-    title=title,
-    description=description,
-    article=article)
+    title="Image enhancer",
+    description="RealESRGAN X4 plus")
 
-demo.queue().launch(share=False, debug=True)
+app.launch(share=False, debug=True, show_error=True, mcp_server=True)
+app.queue()