Update app.py

app.py

@@ -1,33 +1,78 @@
-import gradio as gr
 import os
 import cv2
 import torch
+import gradio as gr
 import numpy as np
 from PIL import Image
+import tempfile
+import requests
+
 from gfpgan import GFPGANer
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from realesrgan.utils import RealESRGANer
+
+# -------------------------------
+# Model Setup
+# -------------------------------
+MODEL_DIR = os.path.join(os.path.expanduser("~"), ".imgen_models")
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+GFPGAN_MODEL = os.path.join(MODEL_DIR, "GFPGANv1.4.pth")
+GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
+
+if not os.path.exists(GFPGAN_MODEL):
+    print("Downloading GFPGAN model...")
+    r = requests.get(GFPGAN_URL, allow_redirects=True)
+    with open(GFPGAN_MODEL, 'wb') as f:
+        f.write(r.content)
 
-# Download GFPGAN model if not already present
-model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
-model_path = "GFPGANv1.4.pth"
+ESRGAN_MODEL = os.path.join(MODEL_DIR, "realesr-general-x4v3.pth")
+ESRGAN_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
 
-if not os.path.exists(model_path):
-    import requests
-    r = requests.get(model_url, allow_redirects=True)
-    open(model_path, 'wb').write(r.content)
+if not os.path.exists(ESRGAN_MODEL):
+    print("Downloading Real-ESRGAN model...")
+    r = requests.get(ESRGAN_URL, allow_redirects=True)
+    with open(ESRGAN_MODEL, 'wb') as f:
+        f.write(r.content)
 
-# Initialize GFPGAN
+# -------------------------------
+# Background Upsampler Setup
+# -------------------------------
+esr_model = SRVGGNetCompact(
+    num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32,
+    upscale=4, act_type='prelu'
+)
+
+bg_upsampler = RealESRGANer(
+    scale=4,
+    model_path=ESRGAN_MODEL,
+    model=esr_model,
+    tile=0,
+    tile_pad=10,
+    pre_pad=0,
+    half=torch.cuda.is_available()
+)
+
+# -------------------------------
+# GFPGAN Restorer Setup
+# -------------------------------
 restorer = GFPGANer(
-    model_path=model_path,
+    model_path=GFPGAN_MODEL,
     upscale=2,
     arch='clean',
     channel_multiplier=2,
-    bg_upsampler=None
+    bg_upsampler=bg_upsampler
 )
 
+# -------------------------------
+# Enhancement Function
+# -------------------------------
 def enhance(image):
-    # Convert PIL to numpy
-    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    # Ensure RGB format
+    img_np = np.array(image.convert("RGB"))
+    img = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
 
+    # Enhance with GFPGAN
     _, _, restored_img = restorer.enhance(
         img,
         has_aligned=False,
@@ -35,16 +80,30 @@ def enhance(image):
         paste_back=True
     )
 
+    # Convert to PIL for saving
     restored_pil = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
-    return restored_pil
 
+    # Save as .jpg to a temporary file
+    temp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
+    restored_pil.save(temp_file.name, format="JPEG", quality=95)
+
+    return temp_file.name  # Path to JPG file
+
+# -------------------------------
+# Gradio Interface
+# -------------------------------
 iface = gr.Interface(
     fn=enhance,
-    inputs=gr.Image(type="pil"),
-    outputs="image",
-    title="IMGEN - AI Photo Enhancer",
-    description="Upload your photo (ID, CV, profile) and enhance it with AI using GFPGAN."
+    inputs=gr.Image(type="pil", label="Upload Image (ID, CV, Profile)"),
+    outputs=gr.File(label="Download Enhanced JPG"),
+    title="📸 IMGEN - AI Photo Enhancer (Face + Outfit)",
+    description="Upload your photo and enhance both the face and outfit using AI (GFPGAN + RealESRGAN). Output is a downloadable JPG file.",
+    allow_flagging="never"
 )
 
+# -------------------------------
+# Launch App
+# -------------------------------
 if __name__ == "__main__":
-    iface.launch()
+    print("Running on GPU" if torch.cuda.is_available() else "Running on CPU")
+    iface.launch()