Update app.py

app.py

@@ -87,23 +87,26 @@ model1 = Generator(3, 1, 3)
 model1.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
 model1.eval()
 
-# model3 = Generator(3, 1, 3)
-# model3.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
-# model3.eval()
+model3 = Generator(3, 1, 3)
+model3.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
+model3.eval()
 
 # model2 = Generator(3, 1, 3)
 # model2.load_state_dict(torch.load('model2.pth', map_location=torch.device('cpu')))
 # model2.eval()
 
 def predict(input_img):
-    input_img = Image.open(input_img)
+    input_img = Image.open(input_img, ver)
     transform = transforms.Compose([transforms.Resize(256, Image.BICUBIC), transforms.ToTensor()])
     input_img = transform(input_img)
     input_img = torch.unsqueeze(input_img, 0)
 
     drawing = 0
     with torch.no_grad():
-        drawing = model1(input_img)[0].detach()
+        if ver == 'Simple Lines':
+            drawing = model3(input_img)[0].detach()
+        else:
+            drawing = model1(input_img)[0].detach()
     
     drawing = transforms.ToPILImage()(drawing)
     return drawing