Update app.py

app.py

@@ -104,7 +104,7 @@ def patch2img(outs, idxes, sr_size, scale=4, crop_size=512):
     return (preds / count_mt).to(outs.device)
 
 
-def load_img (filename, norm=True):
+def load_img(filename, norm=True):
     img = np.array(Image.open(filename).convert("RGB"))
     h, w = img.shape[:2]
     
@@ -114,20 +114,21 @@ def load_img (filename, norm=True):
     return img
 
 
-def inference(image, upscale, large_input_flag, color_fix):
-	model = set_safmn(upscale)
-	
-    img = np.array(image).astype(np.float32) / 255.
-    
-    y = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
-    y = y.unsqueeze(0).to(device)
 
+def inference(image, upscale, large_input_flag, color_fix):
 	upscale = int(upscale) # convert type to int
 	if upscale > 4: 
 		upscale = 4 
 	if 0 < upscale < 3:
 		upscale = 2
 
+	model = set_safmn(upscale)
+
+	img = np.array(image)
+	img = img.astype(np.float32) / 255.
+	y = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
+	y = y.unsqueeze(0).to(device)
+
 	# inference
 	if large_input_flag:
 		patches, idx, size = img2patch(y, scale=upscale)
@@ -154,7 +155,7 @@ def inference(image, upscale, large_input_flag, color_fix):
 
 	# color fix
 	if color_fix:
-		y = F.interpolate(y, scale_factor=upscale, mode='bilinear')
+		y = F.interpolate(img, scale_factor=upscale, mode='bilinear')
 		output = wavelet_reconstruction(output, y)
 	# tensor2img
 	output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
@@ -162,7 +163,8 @@ def inference(image, upscale, large_input_flag, color_fix):
 		output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
 	output = (output * 255.0).round().astype(np.uint8)
 
-	return (image, Image.fromarray(output))
+	return (image, output)
+