Upload app.py

app.py

@@ -36,7 +36,7 @@ from transformers import T5EncoderModel, CLIPModel, CLIPProcessor
 
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from realesrgan import RealESRGANer
-from codeformer.app import inference_app as codeformer_infer
+from gfpgan import GFPGANer
 
 # Optional quality metrics imports
 try:
@@ -89,13 +89,14 @@ print("FLUX.1-Fill loaded.")
 
 
 # ================================================================
-# 3. LOAD CODEFORMER & REAL-ESRGAN (Photo & Face Restoration)
+# 3. LOAD GFPGAN & REAL-ESRGAN (Photo & Face Restoration)
 # ================================================================
-print("Setting up CodeFormer & Real-ESRGAN...")
+print("Setting up GFPGAN & Real-ESRGAN...")
 weights_dir = "weights"
 os.makedirs(weights_dir, exist_ok=True)
-esrgan_path = f"{weights_dir}/RealESRGAN_x4plus.pth"
 
+# 1. Real-ESRGAN Weights Setup
+esrgan_path = f"{weights_dir}/RealESRGAN_x4plus.pth"
 if not os.path.exists(esrgan_path):
     url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
     urllib.request.urlretrieve(url, esrgan_path)
@@ -111,7 +112,22 @@ enhancer = RealESRGANer(
     device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
 )
 
-print("CodeFormer & Real-ESRGAN ready.")
+# 2. GFPGAN Weights Setup
+gfpgan_path = f"{weights_dir}/GFPGANv1.3.pth"
+if not os.path.exists(gfpgan_path):
+    url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
+    urllib.request.urlretrieve(url, gfpgan_path)
+
+gfpgan_restorer = GFPGANer(
+    model_path=gfpgan_path,
+    upscale=4,
+    arch="clean",
+    channel_multiplier=2,
+    bg_upsampler=enhancer,
+    device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+)
+
+print("GFPGAN & Real-ESRGAN ready.")
 
 
 # ================================================================
@@ -184,11 +200,10 @@ def _colorfulness(image):
 
 
 # ================================================================
-# 7. ENHANCEMENT FUNCTION (Deblur & Old Photo Restoration)
+# 7. ENHANCEMENT FUNCTION (GFPGAN + Real-ESRGAN)
 # ================================================================
-import tempfile
 @spaces.GPU(duration=45)
-def enhance_image(image, scale_factor, fidelity_weight=0.5, compute_metrics=False):
+def enhance_image(image, scale_factor, weight=0.5, compute_metrics=False):
     if image is None:
         raise gr.Error("Please upload an image first.")
 
@@ -196,47 +211,44 @@ def enhance_image(image, scale_factor, fidelity_weight=0.5, compute_metrics=Fals
     gc.collect()
     torch.cuda.empty_cache()
 
-    # 2. Cap image size to avoid out-of-memory errors
+    # 2. Cap image size
     if max(image.size) > 1200:
         image = cap_image(image, 1200)
 
-    upscale = 4 if scale_factor == "4x" else 2
+    target_scale = 4 if scale_factor == "4x" else 2
 
-    # 3. Save to a temporary file path to prevent OpenCV string conversion errors
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
-        tmp_path = tmp_file.name
-        image.save(tmp_path)
+    # 3. Format numpy BGR array for OpenCV/GFPGAN
+    rgb_np = np.array(image.convert("RGB"), dtype=np.uint8)
+    input_bgr = cv2.cvtColor(rgb_np, cv2.COLOR_RGB2BGR)
 
     try:
-        # Pass the valid file path to codeformer_infer
-        restored_bgr = codeformer_infer(
-            image=tmp_path,
-            background_enhance=True,
-            face_upsample=True,
-            upscale=upscale,
-            codeformer_fidelity=float(fidelity_weight)
+        # Update target scale for GFPGAN restorer
+        gfpgan_restorer.upscale = target_scale
+
+        # Run GFPGAN enhancement (Restores face + upscales background via Real-ESRGAN)
+        cropped_faces, restored_faces, restored_bgr = gfpgan_restorer.enhance(
+            input_bgr,
+            has_aligned=False,
+            only_center_face=False,
+            paste_back=True,
+            weight=float(weight)
         )
 
         if restored_bgr is not None:
             output_rgb = cv2.cvtColor(restored_bgr, cv2.COLOR_BGR2RGB)
             output_image = Image.fromarray(output_rgb)
         else:
-            raise ValueError("CodeFormer returned None")
+            raise ValueError("GFPGAN returned empty array")
 
     except Exception as e:
-        # Fallback to direct Real-ESRGAN upscaling if face processing fails
+        # Fallback directly to Real-ESRGAN if face parsing fails
         try:
-            rgb_np = np.array(image.convert("RGB"), dtype=np.uint8)
-            input_bgr = cv2.cvtColor(rgb_np, cv2.COLOR_RGB2BGR)
-            output_bgr, _ = enhancer.enhance(input_bgr, outscale=upscale)
+            output_bgr, _ = enhancer.enhance(input_bgr, outscale=target_scale)
             output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
             output_image = Image.fromarray(output_rgb)
         except Exception as fallback_err:
             raise gr.Error(f"Enhancement error: {fallback_err}")
     finally:
-        # Clean up temporary file & clear CUDA memory
-        if os.path.exists(tmp_path):
-            os.remove(tmp_path)
         gc.collect()
         torch.cuda.empty_cache()
 
@@ -467,7 +479,7 @@ with gr.Blocks(
 ) as demo:
 
     gr.Markdown("# CanvasAI\n*AI-powered image enhancement, colorization, and outpainting*")
-    gr.Markdown("**Enhance** with CodeFormer & Real-ESRGAN  |  **Colorize** B&W photos  |  **Outpaint** with FLUX.1-Fill")
+    gr.Markdown("**Enhance** with GFPGAN & Real-ESRGAN  |  **Colorize** B&W photos  |  **Outpaint** with FLUX.1-Fill")
     gr.Markdown("---")
 
     with gr.Tabs() as tabs:
@@ -539,14 +551,14 @@ with gr.Blocks(
 
         # Enhance Tab UI
         with gr.Tab("✨ Enhance", id="tab_enhance"):
-            gr.Markdown("## Restore old photos, deblur faces, and upscale up to 4x (CodeFormer + Real-ESRGAN).")
+            gr.Markdown("## Restore old photos, deblur faces, and upscale up to 4x (GFPGAN + Real-ESRGAN).")
             with gr.Row(equal_height=True):
                 with gr.Column(scale=1):
                     enh_input = gr.Image(label="Upload Image", type="pil", height=350)
                     enh_scale = gr.Radio(choices=["2x", "4x"], value="4x", label="Upscale Factor")
-                    enh_fidelity = gr.Slider(
+                    enh_weight = gr.Slider(
                         minimum=0.0, maximum=1.0, value=0.5, step=0.1, 
-                        label="Restoration Fidelity (0 = Max Deblur/Restoration, 1 = Max Original Identity)"
+                        label="Face Restoration Weight (Higher = stronger facial reconstruction)"
                     )
                     enh_metrics = gr.Checkbox(label="📊 Compute quality metrics", value=False)
                     enh_btn = gr.Button("Enhance & Deblur", variant="primary", size="lg")
@@ -556,7 +568,7 @@ with gr.Blocks(
             
             enh_btn.click(
                 fn=enhance_image, 
-                inputs=[enh_input, enh_scale, enh_fidelity, enh_metrics], 
+                inputs=[enh_input, enh_scale, enh_weight, enh_metrics], 
                 outputs=[enh_output, enh_info]
             )