Update app.py

app.py

@@ -4,43 +4,33 @@ 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_DIR = os.path.join(os.path.expanduser("~"), ".imgen_models")
-os.makedirs(MODEL_DIR, exist_ok=True)
-
-def download_model(url, path):
-    if not os.path.exists(path):
-        print(f"Downloading model from {url}...")
-        r = requests.get(url, allow_redirects=True)
-        with open(path, 'wb') as f:
-            f.write(r.content)
+# Download GFPGAN model if not already present
+GFPGAN_MODEL = "GFPGANv1.4.pth"
+GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
 
-download_model(
-    "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth",
-    os.path.join(MODEL_DIR, "GFPGANv1.4.pth")
-)
+if not os.path.exists(GFPGAN_MODEL):
+    import requests
+    r = requests.get(GFPGAN_URL, allow_redirects=True)
+    open(GFPGAN_MODEL, 'wb').write(r.content)
 
-download_model(
-    "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
-    os.path.join(MODEL_DIR, "realesr-general-x4v3.pth")
-)
+# Download RealESRGAN model
+ESRGAN_MODEL = "realesr-general-x4v3.pth"
+ESRGAN_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
 
-GFPGAN_MODEL = os.path.join(MODEL_DIR, "GFPGANv1.4.pth")
-ESRGAN_MODEL = os.path.join(MODEL_DIR, "realesr-general-x4v3.pth")
-
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-print(f"Using device: {device}")
+if not os.path.exists(ESRGAN_MODEL):
+    import requests
+    r = requests.get(ESRGAN_URL, allow_redirects=True)
+    open(ESRGAN_MODEL, 'wb').write(r.content)
 
+# Initialize Real-ESRGAN as background upsampler
 esr_model = SRVGGNetCompact(
     num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32,
     upscale=4, act_type='prelu'
-).to(device)
+)
 
 bg_upsampler = RealESRGANer(
     scale=4,
@@ -49,45 +39,39 @@ bg_upsampler = RealESRGANer(
     tile=0,
     tile_pad=10,
     pre_pad=0,
-    half=(device=='cuda')
+    half=torch.cuda.is_available()
 )
 
+# Initialize GFPGAN with Real-ESRGAN as background upsampler
 restorer = GFPGANer(
     model_path=GFPGAN_MODEL,
     upscale=2,
     arch='clean',
     channel_multiplier=2,
-    bg_upsampler=bg_upsampler,
-    device=device
+    bg_upsampler=bg_upsampler
 )
 
+# Enhancement function
 def enhance(image):
-    try:
-        img_np = np.array(image.convert("RGB"))
-        img = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
-
-        _, _, restored_img = restorer.enhance(
-            img,
-            has_aligned=False,
-            only_center_face=False,
-            paste_back=True
-        )
-        restored_pil = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
+    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
-        temp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
-        restored_pil.save(temp_file.name, format="JPEG", quality=95)
+    _, _, restored_img = restorer.enhance(
+        img,
+        has_aligned=False,
+        only_center_face=False,
+        paste_back=True
+    )
 
-        return temp_file.name
-    except Exception as e:
-        return f"Error: {e}"
+    restored_pil = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
+    return restored_pil
 
+# Gradio UI
 iface = gr.Interface(
     fn=enhance,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.File(label="Download Enhanced JPG"),
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
     title="IMGEN - AI Photo Enhancer (Face + Outfit)",
-    description="Enhance face and outfit with GFPGAN + RealESRGAN.",
-    allow_flagging="never"
+    description="Upload your photo (ID, CV, profile) and enhance both the face and outfit with AI using GFPGAN + RealESRGAN."
 )
 
 if __name__ == "__main__":