Update app_enhance.py

app_enhance.py

@@ -1,131 +1,79 @@
 import os
-import subprocess
-import spaces
-import torch
-import cv2
 import uuid
+import subprocess
+from pathlib import Path
+from typing import Tuple
 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)
-
-
-@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
 
+# 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)
 
-def create_demo() -> gr.Blocks:
+        g_btn = gr.Button("Interpolate")
 
-    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_image, scale, enhance_mode],
-            outputs=[output_image, enhance_image_path],
+            inputs=[input_imageA, input_imageB, enhance_mode],
+            outputs=[output_image, output_path]
         )
 
-    return demo
+    return demo
+
+demo = build_interface()
+demo.launch()