commit from ssq

app.py

@@ -6,8 +6,12 @@ from skimage import img_as_ubyte
 
 from Allweather.util import load_img, save_img
 from basicsr.models.archs.histoformer_arch import Histoformer
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+# True
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+# Tesla T4
 
-model_restoration = Histoformer.from_pretrained("sunsean/Histoformer-real").to("cuda")
+model_restoration = Histoformer.from_pretrained("sunsean/Histoformer-real").to("cpu")
 
 model_restoration.eval()
 
@@ -15,7 +19,7 @@ factor = 8
 def predict(input_img):
     img = np.float32(load_img(input_img))/255.
     img = torch.from_numpy(img).permute(2,0,1)
-    input_ = img.unsqueeze(0).cuda()
+    input_ = img.unsqueeze(0)
 
     # Padding in case images are not multiples of 8
     h,w = input_.shape[2], input_.shape[3]
@@ -27,7 +31,7 @@ def predict(input_img):
     restored = model_restoration(input_)
     output_path = "restored.png"
     restored = restored[:,:,:h,:w]
-    restored = torch.clamp(restored,0,1).cpu().detach().permute(0, 2, 3, 1).squeeze(0).numpy()
+    restored = torch.clamp(restored,0,1).detach().permute(0, 2, 3, 1).squeeze(0).numpy()
 
     save_img(output_path, img_as_ubyte(restored))