Update app.py

app.py

@@ -1,94 +1,19 @@
-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
-
-# Create model directory
-MODEL_DIR = os.path.join(os.path.expanduser("~"), ".imgen_models")
-os.makedirs(MODEL_DIR, exist_ok=True)
-
-# Download GFPGAN model if missing
-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 RealESRGAN model if missing
-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(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)
-
-# Setup RealESRGAN background upsampler
-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()
-)
-
-# Setup GFPGAN restorer with bg upsampler
-restorer = GFPGANer(
-    model_path=GFPGAN_MODEL,
-    upscale=2,
-    arch='clean',
-    channel_multiplier=2,
-    bg_upsampler=bg_upsampler
-)
-
-def enhance(image):
-    # Convert PIL image to OpenCV BGR
-    img_np = np.array(image.convert("RGB"))
-    img = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
-
-    # Enhance faces + upscale background/outfit
-    _, _, restored_img = restorer.enhance(
-        img,
-        has_aligned=False,
-        only_center_face=False,  # Enhance all faces in image
-        paste_back=True          # Paste restored faces back on bg
-    )
-
-    # Convert back to PIL RGB
-    restored_pil = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
-
-    # Save to temporary JPG file for download
-    temp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
-    restored_pil.save(temp_file.name, format="JPEG", quality=95)
-
-    return temp_file.name
-
-iface = gr.Interface(
-    fn=enhance,
-    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 face and outfit with AI (GFPGAN + RealESRGAN). Output is a downloadable JPG file.",
-    allow_flagging="never"
+from model.face_enhancer import enhance_face
+from model.dress_enhancer import enhance_clothes
+
+def enhance_image(input_image):
+    face_enhanced = enhance_face(input_image)
+    final_image = enhance_clothes(face_enhanced)
+    return final_image
+
+demo = gr.Interface(
+    fn=enhance_image,
+    inputs=gr.Image(type="pil", label="Upload an image"),
+    outputs=gr.Image(type="pil", label="Enhanced Image"),
+    title="Face + Dress Image Enhancer",
+    description="Enhances facial details using GFPGAN and clothing using SwinIR.",
 )
 
 if __name__ == "__main__":
-    print("Running on GPU" if torch.cuda.is_available() else "Running on CPU")
-    iface.launch()
+    demo.launch()