Update app.py

app.py

@@ -25,55 +25,24 @@ from fastai.learner import load_learner
 model = load_learner(
     hf_hub_download("NDugar/horse_to_zebra_cycle_GAN", "h2z-85epoch.pth")
 )
-model.eval()
-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
+def generate_img(img_path):
+    img = tf.io.read_file(img_path)
+    img = tf.image.decode_png(img)
+    img = tf.expand_dims(img, axis=0)
+    img = preprocess_test_image(img)
+    prediction = model(img, training=False)[0].numpy()
+    return prediction
 
-    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)
-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)
+image = gr.inputs.Image(type="filepath")
+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)