import basicsr.utils.misc as m
print("Loaded misc.py from:", m.__file__)
import gradio as gr
import os
import torch
import cv2
from CodeFormer.inference_codeformer import main as codeformer_infer

def restore_face(image):
    input_path = "input.jpg"
    output_path = "output.png"

    cv2.imwrite(input_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))

    args = type('', (), {})()
    args.input_path = input_path
    args.output_path = output_path
    args.background_enhance = True
    args.face_upsample = True
    args.upscale = 2
    args.weight = 0.5
    args.has_aligned = False
    args.only_center_face = False
    args.device = "cpu"

    codeformer_infer(args)

    output = cv2.imread(output_path)
    output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)

    return output

demo = gr.Interface(
    fn=restore_face,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="numpy"),
    title="CodeFormer Face Restore",
)

demo.launch()