import gradio as gr
import utils
from PIL import Image
import torch
import math
from torchvision import transforms


device = "cpu"
years = [str(y) for y in range(1880, 2020, 10)]


orig_models = {}

for year in years:
    G, w_avg = utils.load_stylegan2(f"pretrained_models/{year}.pkl", device)
    orig_models[year] = { "G": G.eval()}
    

def run_alignment(image_path,idx=None):
    import dlib
    from align_all_parallel import align_face
    predictor = dlib.shape_predictor("pretrained_models/shape_predictor_68_face_landmarks.dat")
    aligned_image = align_face(filepath=image_path, predictor=predictor, idx=idx) 
    print("Aligned image has shape: {}".format(aligned_image.size))
    return aligned_image 

def predict(inp):
  #with torch.no_grad():     
  return inp


gr.Interface(fn=predict, 
             inputs=gr.Image(type="pil"),
             outputs=gr.Image(type="pil"),
             #examples=["lion.jpg", "cheetah.jpg"]
             ).launch()