Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import matplotlib.pyplot as plt 
-
+from PIL import Image
 import torch.nn as nn
 import torch.nn.functional as F
 import torch
@@ -235,14 +235,14 @@ class GeneratorUNet(nn.Module, HugGANModelHubMixin):
         return self.final(u7)
 
 def load_image_infer(image_file):
-  imageA = Image.fromarray(np.array(imageA)[:, ::-1, :], "RGB")
-  input_image = normalize_test(input_image)
+  image_file = Image.fromarray(np.array(image_file)[:, ::-1, :], "RGB")
+  image_file = normalize_test(image_file)
 
-  return input_image
+  return image_file
   
 def generate_images(test_input):
   test_input = load_image_infer(test_input)
-  prediction = generator(np.expand_dims(test_input, axis=0), training=True)
+  prediction = generator(test_input).data
   fig = plt.figure(figsize=(128, 128))
   title = ['Predicted Image']