import os
import torch
from gfpgan import GFPGANer
import numpy as np
import cv2
from PIL import Image

# Get path to GFPGANv1.4.pth relative to this script
current_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.dirname(current_dir)
model_path = os.path.join(root_dir, 'GFPGANv1.4.pth')

if not os.path.exists(model_path):
    raise FileNotFoundError(f"Model not found at: {model_path}")

# Load GFPGAN model (no URL needed)
gfpgan = GFPGANer(
    model_path=model_path,
    upscale=2,
    arch='clean',
    channel_multiplier=2,
    bg_upsampler=None
)

def enhance_face(input_image):
    image_np = np.array(input_image)
    image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    
    _, _, output_bgr = gfpgan.enhance(
        image_bgr,
        has_aligned=False,
        only_center_face=False,
        paste_back=True
    )
    
    output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(output_rgb)