"""
This file is used for deploying hugging face demo:
https://huggingface.co/spaces/sczhou/CodeFormer
"""

import sys
sys.path.append('CodeFormer')
import os
import cv2
import numpy as np
import torch
import torch.nn.functional as F
import uuid, threading, time, glob
import gradio as gr

from torchvision.transforms.functional import normalize

from basicsr.utils import imwrite, img2tensor, tensor2img
from basicsr.utils.download_util import load_file_from_url
from facelib.utils.face_restoration_helper import FaceRestoreHelper
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.realesrgan_utils import RealESRGANer
from facelib.utils.misc import is_gray

from basicsr.utils.registry import ARCH_REGISTRY


os.system("pip freeze")

pretrain_model_url = {
    'codeformer': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',
    'detection': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
    'parsing': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth',
    'realesrgan': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth'
}
# download weights
if not os.path.exists('CodeFormer/weights/CodeFormer/codeformer.pth'):
    load_file_from_url(url=pretrain_model_url['codeformer'], model_dir='CodeFormer/weights/CodeFormer', progress=True, file_name=None)
if not os.path.exists('CodeFormer/weights/facelib/detection_Resnet50_Final.pth'):
    load_file_from_url(url=pretrain_model_url['detection'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)
if not os.path.exists('CodeFormer/weights/facelib/parsing_parsenet.pth'):
    load_file_from_url(url=pretrain_model_url['parsing'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)
if not os.path.exists('CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth'):
    load_file_from_url(url=pretrain_model_url['realesrgan'], model_dir='CodeFormer/weights/realesrgan', progress=True, file_name=None)

# download images
torch.hub.download_url_to_file(
    'https://replicate.com/api/models/sczhou/codeformer/files/fa3fe3d1-76b0-4ca8-ac0d-0a925cb0ff54/06.png',
    '01.png')
torch.hub.download_url_to_file(
    'https://replicate.com/api/models/sczhou/codeformer/files/a1daba8e-af14-4b00-86a4-69cec9619b53/04.jpg',
    '02.jpg')
torch.hub.download_url_to_file(
    'https://replicate.com/api/models/sczhou/codeformer/files/542d64f9-1712-4de7-85f7-3863009a7c3d/03.jpg',
    '03.jpg')
torch.hub.download_url_to_file(
    'https://replicate.com/api/models/sczhou/codeformer/files/a11098b0-a18a-4c02-a19a-9a7045d68426/010.jpg',
    '04.jpg')
torch.hub.download_url_to_file(
    'https://replicate.com/api/models/sczhou/codeformer/files/7cf19c2c-e0cf-4712-9af8-cf5bdbb8d0ee/012.jpg',
    '05.jpg')
torch.hub.download_url_to_file(
    'https://raw.githubusercontent.com/sczhou/CodeFormer/master/inputs/cropped_faces/0729.png',
    '06.png')


def imread_unicode_safe(path):
    with open(path, "rb") as f:
        data = np.frombuffer(f.read(), dtype=np.uint8)
    return cv2.imdecode(data, cv2.IMREAD_COLOR)

def delayed_remove(path, delay=60):
    time.sleep(delay)
    try:
        if os.path.exists(path):
            os.remove(path)
            print(f"[CLEANUP] removed: {path}")
        else:
            print(f"[CLEANUP] already gone: {path}")
    except Exception as e:
        print(f"[CLEANUP] failed: {path} | {e}")

# set enhancer with RealESRGAN
def set_realesrgan():
    half = True if torch.cuda.is_available() else False
    model = RRDBNet(
        num_in_ch=3,
        num_out_ch=3,
        num_feat=64,
        num_block=23,
        num_grow_ch=32,
        scale=2,
    )
    upsampler = RealESRGANer(
        scale=2,
        model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth",
        model=model,
        tile=400,
        tile_pad=40,
        pre_pad=0,
        half=half,
    )
    return upsampler

upsampler = set_realesrgan()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
codeformer_net = ARCH_REGISTRY.get("CodeFormer")(
    dim_embd=512,
    codebook_size=1024,
    n_head=8,
    n_layers=9,
    connect_list=["32", "64", "128", "256"],
).to(device)
ckpt_path = "CodeFormer/weights/CodeFormer/codeformer.pth"
checkpoint = torch.load(ckpt_path)["params_ema"]
codeformer_net.load_state_dict(checkpoint)
codeformer_net.eval()

os.makedirs('output', exist_ok=True)

def inference(image, face_align, background_enhance, face_upsample, upscale, codeformer_fidelity):
    """Run a single prediction on the model"""
    try: # global try
        # take the default setting for the demo
        only_center_face = False
        draw_box = False
        detection_model = "retinaface_resnet50"

        face_align = face_align if face_align is not None else True
        background_enhance = background_enhance if background_enhance is not None else True
        face_upsample = face_upsample if face_upsample is not None else True
        upscale = upscale if (upscale is not None and upscale > 0) else 2

        has_aligned = not face_align
        upscale = 1 if has_aligned else upscale
        
        if isinstance(image, dict):
            image_path = image.get("name")
        elif isinstance(image, str):
            image_path = image
        else:
            image_path = None
            raise gr.Error("Invalid input image.")
        
        if not os.path.exists(image_path):
            raise gr.Error("Invalid input image.")

        print('Inp:', image_path, background_enhance, face_upsample, upscale, codeformer_fidelity)
        
        img = imread_unicode_safe(image_path)

        if img is None:
            raise gr.Error("Failed to read input image.")
        
        print('\timage size:', img.shape)

        upscale = int(upscale) # convert type to int
        if upscale > 4: # avoid memory exceeded due to too large upscale
            upscale = 4 
        if upscale > 2 and max(img.shape[:2])>1000: # avoid memory exceeded due to too large img resolution
            upscale = 2 
        if min(img.shape[:2]) > 1100 or max(img.shape[:2])>1500: # avoid memory exceeded due to too large img resolution
            upscale = 1
            background_enhance = False
            face_upsample = False

        h, w = img.shape[:2]
        if h * w > 4_000_000: # avoid memory exceeded due to too large img resolution
            raise gr.Error(
                "Image resolution is too large (>4 megapixels). "
                "To keep the demo responsive and avoid long queue times, this case is skipped. "
                "For such inputs, please deploy this demo locally and remove this limit."
            )
            
        face_helper = FaceRestoreHelper(
            upscale,
            face_size=512,
            crop_ratio=(1, 1),
            det_model=detection_model,
            save_ext="png",
            use_parse=True,
            device=device,
        )
        bg_upsampler = upsampler if background_enhance else None
        face_upsampler = upsampler if face_upsample else None

        if has_aligned:
            # the input faces are already cropped and aligned
            img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
            face_helper.is_gray = is_gray(img, threshold=5)
            if face_helper.is_gray:
                print('\tgrayscale input: True')
            face_helper.cropped_faces = [img]
        else:
            face_helper.read_image(img)
            # get face landmarks for each face
            num_det_faces = face_helper.get_face_landmarks_5(
            only_center_face=only_center_face, resize=640, eye_dist_threshold=5
            )
            print(f'\tdetect {num_det_faces} faces')
            # align and warp each face
            face_helper.align_warp_face()

            if min(img.shape[:2]) > 1000 and num_det_faces > 15:
                raise gr.Error(
                    "Too many faces detected (>15) in a high-resolution image. "
                    "To keep the demo responsive and avoid long queue times, this case is skipped. "
                    "For such inputs, please deploy this demo locally and remove this limit."
                )

            
        # face restoration for each cropped face
        for idx, cropped_face in enumerate(face_helper.cropped_faces):
            # prepare data
            cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
            cropped_face_t = cropped_face_t.unsqueeze(0).to(device)

            try:
                with torch.no_grad():
                    output = codeformer_net(
                        cropped_face_t, w=codeformer_fidelity, adain=True
                    )[0]
                    restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
                del output
                torch.cuda.empty_cache()
            except RuntimeError as error:
                print(f"Failed inference for CodeFormer: {error}")
                restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))

            restored_face = restored_face.astype("uint8")
            face_helper.add_restored_face(restored_face, cropped_face)

        # paste_back
        if not has_aligned:
            # upsample the background
            if bg_upsampler is not None:
                # Now only support RealESRGAN for upsampling background
                bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
            else:
                bg_img = None
            face_helper.get_inverse_affine(None)
            # paste each restored face to the input image
            if face_upsample and face_upsampler is not None:
                restored_img = face_helper.paste_faces_to_input_image(
                    upsample_img=bg_img,
                    draw_box=draw_box,
                    face_upsampler=face_upsampler,
                )
            else:
                restored_img = face_helper.paste_faces_to_input_image(
                    upsample_img=bg_img, draw_box=draw_box
                )
        else:
            restored_img = restored_face

        # save restored img
        # save_path = f'output/out.png'
        # imwrite(restored_img, str(save_path))

        # restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
        # return restored_img


        #save restored img
        save_path = f"output/{uuid.uuid4().hex}.png"
        imwrite(restored_img, save_path)
        print(f"[SAVE] path={save_path} outputs={len(glob.glob('output/*.png'))}")
        
        threading.Thread(
            target=delayed_remove,
            args=(save_path,30),
            daemon=True
        ).start()
    
        return save_path, None

    except gr.Error:
        raise
        
    except Exception as error:
        print('[UNEXPECTED ERROR]', error)
        raise gr.Error("Unexpected error. Please try another image.")
        

title = "CodeFormer: Robust Face Restoration and Enhancement Network"

description = r"""<center><img src='https://user-images.githubusercontent.com/14334509/189166076-94bb2cac-4f4e-40fb-a69f-66709e3d98f5.png' alt='CodeFormer logo'></center>
<br>
<b>Official Gradio demo</b> for <a href='https://github.com/sczhou/CodeFormer' target='_blank'><b>Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022)</b></a><br>
🔥 CodeFormer is a robust face restoration algorithm for old photos or AI-generated faces.<br>
🤗 Try CodeFormer for improved stable-diffusion generation!<br>
"""

article = r"""
If CodeFormer is helpful, please help to ⭐ the <a href='https://github.com/sczhou/CodeFormer' target='_blank'>Github Repo</a>. Thanks! 
[![GitHub Stars](https://img.shields.io/github/stars/sczhou/CodeFormer?style=social)](https://github.com/sczhou/CodeFormer)

---

📝 **Citation**

If our work is useful for your research, please consider citing:
```bibtex
@inproceedings{zhou2022codeformer,
    author = {Zhou, Shangchen and Chan, Kelvin C.K. and Li, Chongyi and Loy, Chen Change},
    title = {Towards Robust Blind Face Restoration with Codebook Lookup TransFormer},
    booktitle = {NeurIPS},
    year = {2022}
}
```

📋 **License**

This project is licensed under <a rel="license" href="https://github.com/sczhou/CodeFormer/blob/master/LICENSE">S-Lab License 1.0</a>. 
Redistribution and use for non-commercial purposes should follow this license.

📧 **Contact**

If you have any questions, please feel free to reach me out at <b>shangchenzhou@gmail.com</b>.

🤗 **Find Me:**
<style type="text/css">
td {
    padding-right: 0px !important;
}

.gradio-container-4-37-2 .prose table, .gradio-container-4-37-2 .prose tr, .gradio-container-4-37-2 .prose td, .gradio-container-4-37-2 .prose th {
    border: 0px solid #ffffff;
    border-bottom: 0px solid #ffffff;
}

</style>

<table>
<tr>
    <td><a href="https://github.com/sczhou"><img style="margin:-0.8em 0 2em 0" src="https://img.shields.io/github/followers/sczhou?style=social" alt="Github Follow"></a></td>
    <td><a href="https://twitter.com/ShangchenZhou"><img style="margin:-0.8em 0 2em 0" src="https://img.shields.io/twitter/follow/ShangchenZhou?label=%40ShangchenZhou&style=social" alt="Twitter Follow"></a></td>
</tr>
</table>

<center><img src='https://api.infinitescript.com/badgen/count?name=sczhou/CodeFormer&ltext=Visitors&color=6dc9aa' alt='visitors'></center>
"""

with gr.Blocks() as demo:
    gr.Markdown(title)
    gr.Markdown(description)
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="filepath", label="Input")
            face_align = gr.Checkbox(value=True, label="Pre_Face_Align")
            background_enhance = gr.Checkbox(value=True, label="Background_Enhance")
            face_enhance = gr.Checkbox(value=True, label="Face_Upsample")
            upscale_factor = gr.Number(value=2, label="Rescaling_Factor (up to 4)")
            codeformer_fidelity = gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)')
            submit = gr.Button('Enhance Image')
        with gr.Column():
            output_img = gr.Image(type="filepath", label="Output")
            note = gr.Markdown("**Please download the output within 30 seconds.**")
            
    inps = [input_img, face_align, background_enhance, face_enhance, upscale_factor, codeformer_fidelity]
    outs = [output_img, note]
    submit.click(fn=inference, inputs=inps, outputs=outs)
            
    ex = gr.Examples([
        ['01.png', True, True, True, 2, 0.7],
        ['02.jpg', True, True, True, 2, 0.7],
        ['03.jpg', True, True, True, 2, 0.7],
        ['04.jpg', True, True, True, 2, 0.1],
        ['05.jpg', True, True, True, 2, 0.1],
        ['06.png', False, True, True, 1, 0.5]
      ],
        inputs=inps,
        cache_examples=False)
    
    gr.Markdown(article)
    
DEBUG = os.getenv('DEBUG') == '1'
demo.queue(api_open=False, max_size=10, default_concurrency_limit=2)
demo.launch(debug=DEBUG)