from cgitb import enable
from ctypes.wintypes import HFONT
import os
import sys
import torch
import gradio as gr
import numpy as np
import torchvision.transforms as transforms
from torch.autograd import Variable
from huggingface_hub import hf_hub_download
from PIL import Image
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
MODEL_PATH = "models"
COLOUR_MODEL = "RGB"
MODEL_REPO = "NDugar/horse_to_zebra_cycle_GAN"
MODEL_FILE = "h2z-85epoch.pth"
model_hfhub = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
enable_gpu = torch.cuda.is_available()
map_location = torch.device("cuda") if enable_gpu else "cpu"
from huggingface_hub import hf_hub_download
from fastai.learner import load_learner

model = torch.load("h2z-85epoch.pth",map_location=torch.device('cpu'))
def generate_img(img):
    prediction = model(img, training=False)
    return prediction


image = gr.inputs.Image(type="numpy")
op = gr.outputs.Image(type="numpy")

iface = gr.Interface(
    generate_img,
    image,
    op,
    title="CycleGAN-using UPIT",
    description='CycleGAN model using Horse to Zebra using UPIT - https://github.com/tmabraham/UPIT'
)

iface.launch(cache_examples=False)