app.py

import os
import cv2
import gradio as gr
import torch
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from gfpgan.utils import GFPGANer
from realesrgan.utils import RealESRGANer

# Create output directory
os.makedirs("output", exist_ok=True)

# Load background upsampler (RealESRGAN)
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
upsampler = RealESRGANer(
    scale=4,
    model_path='realesr-general-x4v3.pth',
    model=model,
    tile=0,
    tile_pad=10,
    pre_pad=0,
    half=torch.cuda.is_available()
)

def inference(img_path, version, scale):
    extension = os.path.splitext(os.path.basename(str(img_path)))[1]
    img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)

    if img is None:
        return None, None

    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    elif img.shape[2] == 4:
        img = img[:, :, :3]  # Remove alpha

    if version == 'v1.2':
        model_path = 'GFPGANv1.2.pth'
        arch = 'clean'
    elif version == 'v1.3':
        model_path = 'GFPGANv1.3.pth'
        arch = 'clean'
    elif version == 'v1.4':
        model_path = 'GFPGANv1.4.pth'
        arch = 'clean'
    elif version == 'RestoreFormer':
        model_path = 'RestoreFormer.pth'
        arch = 'RestoreFormer'
    else:
        return None, None

    face_enhancer = GFPGANer(
        model_path=model_path,
        upscale=2,
        arch=arch,
        channel_multiplier=2,
        bg_upsampler=upsampler
    )

    _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)

    # Rescale
    if scale != 2:
        h, w = output.shape[:2]
        output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=cv2.INTER_LANCZOS4)

    save_path = f"output/restored_{version}.jpg"
    cv2.imwrite(save_path, output)
    output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)

    return output, save_path

demo = gr.Interface(
    fn=inference,
    inputs=[
        gr.Image(type="filepath", label="Input Image"),
        gr.Radio(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer'], label="GFPGAN Version", value="v1.4"),
        gr.Slider(1, 4, value=2, label="Rescaling Factor")
    ],
    outputs=[
        gr.Image(type="numpy", label="Restored Image"),
        gr.File(label="Download")
    ],
    title="GFPGAN Face Restoration on Hugging Face",
    description="Restore old or AI-generated faces using GFPGAN."
)

demo.queue().launch()