Update app_enhance.py

app_enhance.py

@@ -1,79 +1,131 @@
 import os
-import uuid
 import subprocess
-from pathlib import Path
-from typing import Tuple
+import spaces
+import torch
+import cv2
+import uuid
 import gradio as gr
+import numpy as np
+
+from PIL import Image
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from gfpgan.utils import GFPGANer
+from realesrgan.utils import RealESRGANer
+
+def runcmd(cmd, verbose = False):
+
+    process = subprocess.Popen(
+        cmd,
+        stdout = subprocess.PIPE,
+        stderr = subprocess.PIPE,
+        text = True,
+        shell = True
+    )
+    std_out, std_err = process.communicate()
+    if verbose:
+        print(std_out.strip(), std_err)
+    pass
+
+
+if not os.path.exists('GFPGANv1.4.pth'):
+    runcmd("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('realesr-general-x4v3.pth'):
+    runcmd("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+
+
+
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+model_path = 'realesr-general-x4v3.pth'
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
 
-# Constants
-TMP_DIR = "/tmp/gradio/output/"
-GIF_EXT = "gif"
-PALETTE_PATH = "/tmp/gradio/palette.png"
-
-def ensure_tmp_dir():
-    os.makedirs(TMP_DIR, exist_ok=True)
-
-def generate_interpolation_frames(img_a: str, img_b: str, exp: int = 4):
-    """Runs inference_img.py to generate interpolated frames"""
-    cmd = [
-        "python3", "inference_img.py",
-        "--img", img_a, img_b,
-        "--exp", str(exp)
-    ]
-    subprocess.run(cmd, check=True)
-
-def create_palette():
-    """Generates GIF palette from interpolated frames"""
-    cmd = [
-        "ffmpeg", "-y", "-r", "14", "-f", "image2",
-        "-i", f"{TMP_DIR}img%d.png",
-        "-vf", "palettegen=stats_mode=single",
-        PALETTE_PATH
-    ]
-    subprocess.run(cmd, check=True)
-
-def write_gif(gif_path: str):
-    """Creates final interpolated GIF using palette"""
-    cmd = [
-        "ffmpeg", "-y", "-r", "14", "-f", "image2",
-        "-i", f"{TMP_DIR}img%d.png",
-        "-i", PALETTE_PATH,
-        "-lavfi", "paletteuse",
-        gif_path
-    ]
-    subprocess.run(cmd, check=True)
-
-def enhance_image(img_a: str, img_b: str, mode: str) -> Tuple[str, str]:
-    ensure_tmp_dir()
-    gif_path = f"{TMP_DIR}{uuid.uuid4()}.{GIF_EXT}"
-
-    try:
-        generate_interpolation_frames(img_a, img_b)
-        create_palette()
-        write_gif(gif_path)
-        return gif_path, gif_path
-    except subprocess.CalledProcessError as e:
-        raise gr.Error(f"Interpolation failed: {e}")
-
-# Gradio UI
-def build_interface():
-    with gr.Blocks(title="RIFE Interpolation") as demo:
-        with gr.Row():
-            input_imageA = gr.Image(label="Image A", type="filepath")
-            input_imageB = gr.Image(label="Image B", type="filepath")
-        enhance_mode = gr.Dropdown(choices=["default"], value="default", label="Mode")
-        output_image = gr.Image(label="Result GIF", type="filepath")
-        output_path = gr.Textbox(label="GIF Path", interactive=False)
 
-        g_btn = gr.Button("Interpolate")
+@spaces.GPU(duration=15)
+def enhance_image(
+    input_image: Image,
+    scale: int,
+    enhance_mode: str,
+):
+    only_face = enhance_mode == "Only Face Enhance"
+    if enhance_mode == "Only Face Enhance":
+        face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=scale, arch='clean', channel_multiplier=2)
+    elif enhance_mode == "Only Image Enhance":
+        face_enhancer = None
+    else:
+        face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+    
+    img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
 
+    h, w = img.shape[0:2]
+    if h < 300:
+        img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+    
+    if face_enhancer is not None:
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=only_face, paste_back=True)
+    else:
+        output, _ = upsampler.enhance(img, outscale=scale)
+
+    # if scale != 2:
+    #     interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
+    #     h, w = img.shape[0:2]
+    #     output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
+    
+    h, w = output.shape[0:2]
+    max_size = 3480
+    if h > max_size:
+        w = int(w * max_size / h)
+        h = max_size
+    
+    if w > max_size:
+        h = int(h * max_size / w)
+        w = max_size
+    
+    output = cv2.resize(output, (w, h), interpolation=cv2.INTER_LANCZOS4)
+
+    enhanced_image = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
+    tmpPrefix = "/tmp/gradio/"
+
+    extension = 'png'
+
+    targetDir = f"{tmpPrefix}output/"
+    if not os.path.exists(targetDir):
+        os.makedirs(targetDir)
+
+    enhanced_path = f"{targetDir}{uuid.uuid4()}.{extension}"
+    enhanced_image.save(enhanced_path, quality=100)
+        
+    return enhanced_image, enhanced_path
+
+
+def create_demo() -> gr.Blocks:
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                scale = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Scale")
+            with gr.Column():
+                enhance_mode = gr.Dropdown(
+                    label="Enhance Mode",
+                    choices=[
+                        "Only Face Enhance",
+                        "Only Image Enhance",
+                        "Face Enhance + Image Enhance",
+                    ],
+                    value="Face Enhance + Image Enhance",
+                )
+                g_btn = gr.Button("Enhance Image")
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="pil")
+            with gr.Column():
+                output_image = gr.Image(label="Enhanced Image", type="pil", interactive=False)
+                enhance_image_path = gr.File(label="Download the Enhanced Image", interactive=False)
+                
+        
         g_btn.click(
             fn=enhance_image,
-            inputs=[input_imageA, input_imageB, enhance_mode],
-            outputs=[output_image, output_path]
+            inputs=[input_image, scale, enhance_mode],
+            outputs=[output_image, enhance_image_path],
         )
 
-    return demo
-
-demo = build_interface()
-demo.launch()
+    return demo