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