app.py

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)