Update app.py

app.py

@@ -134,16 +134,22 @@ model = Generator(3, 32, 3, 4).cpu() # input_dim, num_filter, output_dim, num_re
 model.load_state_dict(torch.load('G_A_HW4_SAVE.pt',map_location=torch.device('cpu')))
 print(model)
 model.eval()
+model_2 = Generator(3, 32, 3, 4)
+model_2.load_state_dict(torch.load('G_B_HW4_SAVE.pt',map_location=torch.device('cpu')))
 
 totensor = torchvision.transforms.ToTensor()
 normalize_fn = torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 topilimage = torchvision.transforms.ToPILImage()
 
-def predict(input):
-    im = normalize_fn(totensor(input))
-    print(im.shape)
-    preds = model(im.unsqueeze(0))/2 + 0.5
-    print(preds.shape)
-    return topilimage(preds.squeeze(0).detach())
+def predict(input_1, input_2):
+    im1 = normalize_fn(totensor(input))
+    print(im1.shape)
+    preds1 = model(im1.unsqueeze(0))/2 + 0.5
+    print(preds1.shape)
+    im2 = normalize_fn(totensor(input))
+    print(im2.shape)
+    preds2 = model_2(im1.unsqueeze(0))/2 + 0.5
+    print(preds2.shape)
+    return topilimage(preds1.squeeze(0).detach()), topilimage(preds2.squeeze(0).detach())
 
 gr_interface = gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(256,256)), outputs="image", title='Emoji_CycleGAN').launch()