import os
from huggingface_hub import hf_hub_download
os.system("pip -qq install facenet_pytorch")
from facenet_pytorch import MTCNN
from torchvision import transforms
import torch, PIL
from tqdm.notebook import tqdm
import gradio as gr
import torch

device = "cuda:0" if torch.cuda.is_available() else "cpu"

image_size = 512

means = [0.5, 0.5, 0.5]
stds = [0.5, 0.5, 0.5]

model_path = hf_hub_download(repo_id="jjeamin/ArcaneStyleTransfer", filename="pytorch_model.bin")

if 'cuda' in device:
    style_transfer = torch.jit.load(model_path).eval().cuda().half()
    t_stds = torch.tensor(stds).cuda().half()[:,None,None]
    t_means = torch.tensor(means).cuda().half()[:,None,None]
else:
    style_transfer = torch.jit.load(model_path).eval().cpu()
    t_stds = torch.tensor(stds).cpu()[:,None,None]
    t_means = torch.tensor(means).cpu()[:,None,None]
    
mtcnn = MTCNN(image_size=image_size, margin=80)

def detect(img):
 
        # Detect faces
        batch_boxes, batch_probs, batch_points = mtcnn.detect(img, landmarks=True)
        # Select faces
        if not mtcnn.keep_all:
            batch_boxes, batch_probs, batch_points = mtcnn.select_boxes(
                batch_boxes, batch_probs, batch_points, img, method=mtcnn.selection_method
            )
 
        return batch_boxes, batch_points

def makeEven(_x):
  return _x if (_x % 2 == 0) else _x+1

def scale(boxes, _img, max_res=1_500_000, target_face=256, fixed_ratio=0, max_upscale=2, VERBOSE=False):
 
    x, y = _img.size
 
    ratio = 2 #initial ratio
 
    #scale to desired face size
    if (boxes is not None):
      if len(boxes)>0:
        ratio = target_face/max(boxes[0][2:]-boxes[0][:2]); 
        ratio = min(ratio, max_upscale)
        if VERBOSE: print('up by', ratio)

    if fixed_ratio>0:
      if VERBOSE: print('fixed ratio')
      ratio = fixed_ratio
 
    x*=ratio
    y*=ratio
 
    #downscale to fit into max res 
    res = x*y
    if res > max_res:
      ratio = pow(res/max_res,1/2); 
      if VERBOSE: print(ratio)
      x=int(x/ratio)
      y=int(y/ratio)
 
    #make dimensions even, because usually NNs fail on uneven dimensions due skip connection size mismatch
    x = makeEven(int(x))
    y = makeEven(int(y))
    
    size = (x, y)

    return _img.resize(size)

def scale_by_face_size(_img, max_res=1_500_000, target_face=256, fix_ratio=0, max_upscale=2, VERBOSE=False):
    boxes = None
    boxes, _ = detect(_img)
    if VERBOSE: print('boxes',boxes)
    img_resized = scale(boxes, _img, max_res, target_face, fix_ratio, max_upscale, VERBOSE)
    return img_resized


img_transforms = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(means, stds)])
 
def tensor2im(var):
     return var.mul(t_stds).add(t_means).mul(255.).clamp(0,255).permute(1,2,0)

def proc_pil_img(input_image):
    if 'cuda' in device: 
        transformed_image = img_transforms(input_image)[None,...].cuda().half()
    else:
        transformed_image = img_transforms(input_image)[None,...].cpu()
            
    with torch.no_grad():
        result_image = style_transfer(transformed_image)[0]
        output_image = tensor2im(result_image)
        output_image = output_image.detach().cpu().numpy().astype('uint8')
        output_image = PIL.Image.fromarray(output_image)
    return output_image

def process(im):
    im = scale_by_face_size(im, target_face=image_size, max_res=1_500_000, max_upscale=1)
    res = proc_pil_img(im)
    return res
        
gr.Interface(
    process, 
    inputs=gr.inputs.Image(type="pil", label="Input", shape=(image_size, image_size)),
    outputs=gr.outputs.Image(type="pil", label="Output"),
    title="Arcane Style Transfer",
    description="Gradio demo for Arcane Style Transfer",
    article = "<p style='text-align: center'><a href='https://github.com/jjeamin/anime_style_transfer_pytorch' target='_blank'>Github Repo Pytorch by jjeamin</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=jjeamin_arcane_st' alt='visitor badge'></center></p>",
    examples=[['billie.png'], ['gongyoo.jpeg'], ['IU.png']],
    enable_queue=True,
    allow_flagging=False,
    allow_screenshot=False
    ).launch(enable_queue=True,cache_examples=True)