Create app.py

app.py

@@ -0,0 +1,110 @@
+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 network.Transformer import Transformer
+from huggingface_hub import hf_hub_download
+
+from PIL import Image
+
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Constants
+
+MODEL_PATH = "models"
+COLOUR_MODEL = "RGB"
+
+MODEL_REPO = "NDugar/horse_to_zebra_cycle_GAN"
+MODEL_FILE = "h2z-85epoch.pth"
+
+# Model Initalisation
+#shinkai_model_hfhub = hf_hub_download(repo_id=MODEL_REPO_SHINKAI, filename=MODEL_FILE_SHINKAI)
+#hosoda_model_hfhub = hf_hub_download(repo_id=MODEL_REPO_HOSODA, filename=MODEL_FILE_HOSODA)
+#miyazaki_model_hfhub = hf_hub_download(repo_id=MODEL_REPO_MIYAZAKI, filename=MODEL_FILE_MIYAZAKI)
+model_hfhub = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+
+#shinkai_model = Transformer()
+#hosoda_model = Transformer()
+#miyazaki_model = Transformer()
+model = Transformer()
+
+enable_gpu = torch.cuda.is_available()
+map_location = torch.device("cuda") if enable_gpu else "cpu"
+
+model.load_state_dict(torch.load(model_hfhub, map_location=map_location))
+
+shinkai_model.eval()
+hosoda_model.eval()
+miyazaki_model.eval()
+kon_model.eval()
+
+
+# Functions
+
+def get_model():
+    return model
+
+
+def adjust_image_for_model(img):
+    logger.info(f"Image Height: {img.height}, Image Width: {img.width}")
+    return img
+
+
+def inference(img, style):
+    img = adjust_image_for_model(img)
+    input_image = img.convert(COLOUR_MODEL)
+    input_image = np.asarray(input_image)
+    input_image = input_image[:, :, [2, 1, 0]]
+    input_image = transforms.ToTensor()(input_image).unsqueeze(0)
+    input_image = -1 + 2 * input_image
+
+    if enable_gpu:
+        logger.info(f"CUDA found. Using GPU.")
+        input_image = Variable(input_image).cuda()
+    else:
+        logger.info(f"CUDA not found. Using CPU.")
+        input_image = Variable(input_image).float()
+    model = get_model()
+    output_image = model(input_image)
+    output_image = output_image[0]
+    # BGR -> RGB
+    output_image = output_image[[2, 1, 0], :, :]
+    output_image = output_image.data.cpu().float() * 0.5 + 0.5
+
+    return transforms.ToPILImage()(output_image)
+
+
+# Gradio setup
+
+title = "Horse 2 Zebra GAN"
+description = "Gradio Demo for CycleGAN"
+
+gr.Interface(
+    fn=inference,
+    inputs=[
+        gr.inputs.Image(
+            type="pil",
+            label="Input Photo",
+        ),
+    ],
+    outputs=gr.outputs.Image(
+        type="pil",
+        label="Output Image",
+    ),
+    title=title,
+    description=description,
+    article=article,
+    examples=examples,
+    allow_flagging="never",
+    allow_screenshot=False,
+).launch(enable_queue=True)