Update app.py

app.py

@@ -11,33 +11,36 @@ 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
+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_model(
+    "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth",
+    os.path.join(MODEL_DIR, "GFPGANv1.4.pth")
+)
+
+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")
+)
+
 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
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Using device: {device}")
+
 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,
@@ -46,49 +49,46 @@ bg_upsampler = RealESRGANer(
     tile=0,
     tile_pad=10,
     pre_pad=0,
-    half=torch.cuda.is_available()
+    half=(device=='cuda')
 )
 
-# Setup GFPGAN restorer with bg upsampler
 restorer = GFPGANer(
     model_path=GFPGAN_MODEL,
     upscale=2,
     arch='clean',
     channel_multiplier=2,
-    bg_upsampler=bg_upsampler
+    bg_upsampler=bg_upsampler,
+    device=device
 )
 
 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
-    )
+    try:
+        img_np = np.array(image.convert("RGB"))
+        img = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
 
-    # Convert back to PIL RGB
-    restored_pil = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
+        _, _, 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))
 
-    # 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)
+        temp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
+        restored_pil.save(temp_file.name, format="JPEG", quality=95)
 
-    return temp_file.name
+        return temp_file.name
+    except Exception as e:
+        return f"Error: {e}"
 
 iface = gr.Interface(
     fn=enhance,
-    inputs=gr.Image(type="pil", label="Upload Image (ID, CV, Profile)"),
+    inputs=gr.Image(type="pil", label="Upload Image"),
     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.",
+    title="IMGEN - AI Photo Enhancer (Face + Outfit)",
+    description="Enhance face and outfit with GFPGAN + RealESRGAN.",
     allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    print("Running on GPU" if torch.cuda.is_available() else "Running on CPU")
     iface.launch()