Update app.py

app.py

@@ -55,9 +55,9 @@ class Generator(nn.Module):
             model2 += [ResidualBlock(in_features)]
         self.model2 = nn.Sequential(*model2)
 
-        # Upsampling
+        # More downsampling
         model3 = []
-        out_features = in_features//2
+        out_features = in_features*3
         for _ in range(2):
             model3 += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
                         norm_layer(out_features),
@@ -87,9 +87,13 @@ model1 = Generator(3, 1, 3)
 model1.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
 model1.eval()
 
-model2 = Generator(3, 1, 3)
-model2.load_state_dict(torch.load('model2.pth', map_location=torch.device('cpu')))
-model2.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)
@@ -99,7 +103,7 @@ def predict(input_img):
 
     drawing = 0
     with torch.no_grad():
-        drawing = model1(input_img)[0].detach()
+        drawing = model3(input_img)[0].detach()
     
     drawing = transforms.ToPILImage()(drawing)
     return drawing