Update app.py

app.py

@@ -11,33 +11,29 @@ from gfpgan import GFPGANer
 from basicsr.archs.srvgg_arch import SRVGGNetCompact
 from realesrgan.utils import RealESRGANer
 
-# -------------------------------
-# Model Setup
-# -------------------------------
+# 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)
 
-# -------------------------------
-# Background Upsampler Setup
-# -------------------------------
+# 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'
@@ -53,9 +49,7 @@ bg_upsampler = RealESRGANer(
     half=torch.cuda.is_available()
 )
 
-# -------------------------------
-# GFPGAN Restorer Setup
-# -------------------------------
+# Setup GFPGAN restorer with bg upsampler
 restorer = GFPGANer(
     model_path=GFPGAN_MODEL,
     upscale=2,
@@ -64,46 +58,37 @@ restorer = GFPGANer(
     bg_upsampler=bg_upsampler
 )
 
-# -------------------------------
-# Enhancement Function
-# -------------------------------
 def enhance(image):
-    # Ensure RGB format
+    # Convert PIL image to OpenCV BGR
     img_np = np.array(image.convert("RGB"))
     img = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
 
-    # Enhance with GFPGAN
+    # Enhance faces + upscale background/outfit
     _, _, restored_img = restorer.enhance(
         img,
         has_aligned=False,
-        only_center_face=False,
-        paste_back=True
+        only_center_face=False,  # Enhance all faces in image
+        paste_back=True          # Paste restored faces back on bg
     )
 
-    # Convert to PIL for saving
+    # Convert back to PIL RGB
     restored_pil = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
 
-    # Save as .jpg to a temporary file
+    # 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  # Path to JPG file
+    return temp_file.name
 
-# -------------------------------
-# Gradio Interface
-# -------------------------------
 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 the face and outfit using AI (GFPGAN + RealESRGAN). Output is a downloadable JPG file.",
+    description="Upload your photo and enhance both face and outfit with AI (GFPGAN + RealESRGAN). Output is a downloadable JPG file.",
     allow_flagging="never"
 )
 
-# -------------------------------
-# Launch App
-# -------------------------------
 if __name__ == "__main__":
     print("Running on GPU" if torch.cuda.is_available() else "Running on CPU")
     iface.launch()