NDugar's picture
Update app.py
5ee2f62
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):
_,_,preds = model.get_preds(dl=[b], with_decoded=True)
return preds
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)