Update app.py

app.py

@@ -114,7 +114,8 @@ def inference(image, upscale, large_input_flag, color_fix):
 
 	# img2tensor
 	y = np.array(image).astype(np.float32) / 255.
-	y = torch.from_numpy(np.transpose(y[:, :, [2, 1, 0]], (2, 0, 1))).float()
+	# y = torch.from_numpy(np.transpose(y[:, :, [2, 1, 0]], (2, 0, 1))).float()
+	y = torch.from_numpy(np.transpose(y, (2, 0, 1))).float()
 	y = y.unsqueeze(0).to(device)
 
 	# inference
@@ -147,11 +148,15 @@ def inference(image, upscale, large_input_flag, color_fix):
 		output = wavelet_reconstruction(output, y)
 	# tensor2img
 	output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
-	if output.ndim == 3:
-		output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
+	# if output.ndim == 3:
+	# 	output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
 	output = (output * 255.0).round().astype(np.uint8)
 	
-	return (image, Image.fromarray(output)), output
+    # save results
+	save_path = './out.png'
+	cv2.imwrite(save_path, output)
+	
+	return (image, Image.fromarray(output)), save_path
     
 
 
@@ -233,8 +238,11 @@ demo = gr.Interface(
 		ImageSlider(label="Super-Resolved Image", 
                         type="pil",
                         show_download_button=True,
-                        ),
-		gr.Image()
+                    ),
+		gr.Image(
+			label="Download Output",
+			type='filepath'
+		),
     ],
     title=title,
     description=description,