Create app.py

app.py

@@ -0,0 +1,161 @@
+import os
+import cv2
+import time
+import gc
+import torch
+from gfpgan import GFPGANer
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+from realesrgan import RealESRGANer
+
+# === GPU MEMORY MONITORING ===
+def print_simple_gpu_memory():
+    allocated = torch.cuda.memory_allocated(0) / 1024**3
+    reserved = torch.cuda.memory_reserved(0) / 1024**3
+    print(f"[GPU Memory] Allocated: {allocated:.2f} GB | Reserved: {reserved:.2f} GB")
+
+# === GFPGAN STEP ===
+def run_gfpgan(image_path, output_path, model_path='GFPGAN\GFPGAN\experiments\pretrained_models\GFPGANv1.4.pth'):
+    gfpganer = GFPGANer(
+        model_path=model_path,
+        upscale=2,
+        arch='clean',
+        channel_multiplier=2,
+        bg_upsampler=None
+    )
+
+    img = cv2.imread(image_path)
+    if img is None:
+        raise FileNotFoundError(f"Input image not found: {image_path}")
+
+    _, _, restored_img = gfpganer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+    cv2.imwrite(output_path, restored_img)
+    print(f"[+] GFPGAN output saved: {output_path}")
+    return output_path
+
+# === RealESRGAN STEP ===
+def run_realesrgan(input_path, output_path, model_name='RealESRGAN_x4plus', outscale=4, gpu_id=None):
+    if model_name == 'RealESRGAN_x4plus':
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
+    else:
+        raise NotImplementedError(f'Model {model_name} not implemented')
+
+    model_path = os.path.join('weights', model_name + '.pth')
+    if not os.path.isfile(model_path):
+        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+        for url in file_url:
+            model_path = load_file_from_url(url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
+
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=None,
+        model=model,
+        tile=512,  # <<< IMPORTANT for 4 GB VRAM!
+        tile_pad=10,
+        pre_pad=0,
+        half=False,
+        gpu_id=gpu_id
+    )
+
+    img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)
+    if img is None:
+        raise FileNotFoundError(f'Input image not found: {input_path}')
+
+    output, _ = upsampler.enhance(img, outscale=outscale)
+
+    output_dir = os.path.dirname(output_path)
+    if output_dir and not os.path.exists(output_dir):
+        os.makedirs(output_dir, exist_ok=True)
+
+    cv2.imwrite(output_path, output)
+    print(f"[+] RealESRGAN output saved: {output_path}")
+    return output_path
+
+# === Combined Workflow ===
+def combined_enhance(input_img_path, rescale_factor, output_dir):
+    start_time = time.perf_counter()
+
+    base_name = os.path.splitext(os.path.basename(input_img_path))[0]
+
+    # Step 1: GFPGAN enhancement
+    gfpgan_out = os.path.join(output_dir, f"{base_name}_GFPGAN.png")
+    run_gfpgan(input_img_path, gfpgan_out)
+
+    print("[GPU Memory after GFPGAN]")
+    print_simple_gpu_memory()
+
+    # Free memory before loading RealESRGAN
+    gc.collect()
+    torch.cuda.empty_cache()
+    print("[*] Cleared GPU cache before RealESRGAN")
+
+    # Step 2: RealESRGAN enhancement using GFPGAN result
+    realesrgan_out = os.path.join(output_dir, f"{base_name}_GFPGAN_RealESRGAN_combined.png")
+    run_realesrgan(gfpgan_out, realesrgan_out, outscale=rescale_factor)
+
+    print("[GPU Memory after RealESRGAN]")
+    print_simple_gpu_memory()
+
+    print(f"[***] Final combined image saved at: {realesrgan_out}")
+
+    # Get resolution and file size info
+    img = cv2.imread(realesrgan_out)
+    height, width = img.shape[:2]
+    file_size_bytes = os.path.getsize(realesrgan_out)
+    file_size_mb = file_size_bytes / (1024 * 1024)
+
+    end_time = time.perf_counter()
+    elapsed = end_time - start_time
+    mins, secs = divmod(elapsed, 60)
+    hours, mins = divmod(mins, 60)
+
+    print(f"[***] Image resolution: {width}x{height}, file size: {file_size_mb:.2f} MB")
+    print(f"[***] Total processing time: {int(hours):02d}h:{int(mins):02d}m:{secs:05.2f}s")
+
+    return realesrgan_out
+
+import gradio as gr
+import os
+
+# Your core functions like combined_enhance must be defined/imported before this
+def wrapped_combined_enhance(image, scale_factor):
+    temp_input_path = "temp_input.png"
+    temp_output_dir = "output"
+    os.makedirs(temp_output_dir, exist_ok=True)
+    image.save(temp_input_path)
+
+    try:
+        final_path = combined_enhance(temp_input_path, int(scale_factor), temp_output_dir)
+        return final_path, "Enhancement succeeded!"
+    except Exception as e:
+        return None, f"[ERROR]: {str(e)}"
+
+with gr.Blocks(title="AI Face Enhancer (GFPGAN + RealESRGAN)") as demo:
+    gr.Markdown("# ✨ Face Enhancer Pro")
+
+    with gr.Row():
+        with gr.Column():
+            input_img   = gr.Image(type="pil", label="Upload Image")
+            scale_slider= gr.Slider(1, 4, value=2, step=1, label="Upscale Factor")
+            enhance_btn = gr.Button("Enhance")
+            status_box  = gr.Textbox(label="Status / Logs", lines=6, interactive=False)
+        with gr.Column():
+            output_img  = gr.Image(label="Enhanced Output")
+
+    def ui_callback(image, scale):
+        if image is None:
+            return None, "⚠️ Please upload an image first."
+        out_path, status = wrapped_combined_enhance(image, scale)
+        # return filepath (or None) and status string
+        return out_path, status
+
+    enhance_btn.click(
+        fn=ui_callback,
+        inputs=[input_img, scale_slider],
+        outputs=[output_img, status_box]
+    )
+
+demo.launch()