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

# Download GFPGAN model if not already present
GFPGAN_MODEL = "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):
    import requests
    r = requests.get(GFPGAN_URL, allow_redirects=True)
    open(GFPGAN_MODEL, 'wb').write(r.content)

# Download RealESRGAN model
ESRGAN_MODEL = "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):
    import requests
    r = requests.get(ESRGAN_URL, allow_redirects=True)
    open(ESRGAN_MODEL, 'wb').write(r.content)

# Initialize Real-ESRGAN as background upsampler
esr_model = SRVGGNetCompact(
    num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32,
    upscale=4, act_type='prelu'
)

bg_upsampler = RealESRGANer(
    scale=4,
    model_path=ESRGAN_MODEL,
    model=esr_model,
    tile=0,
    tile_pad=10,
    pre_pad=0,
    half=torch.cuda.is_available()
)

# Initialize GFPGAN with Real-ESRGAN as background upsampler
restorer = GFPGANer(
    model_path=GFPGAN_MODEL,
    upscale=2,
    arch='clean',
    channel_multiplier=2,
    bg_upsampler=bg_upsampler
)

# Enhancement function
def enhance(image):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    _, _, 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))
    return restored_pil

# Gradio UI
iface = gr.Interface(
    fn=enhance,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="IMGEN - AI Photo Enhancer (Face + Outfit)",
    description="Upload your photo (ID, CV, profile) and enhance both the face and outfit with AI using GFPGAN + RealESRGAN."
)

if __name__ == "__main__":
    iface.launch()