support Apple mps

README.md

@@ -46,8 +46,8 @@ from hit_sng_arch import HiT_SNG
 from hit_srf_arch import HiT_SRF
 import cv2
 
-# use GPU (True) or CPU (False)
-cuda_flag = True
+# detect device
+device = torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
 
 # initialize model (change model and upscale according to your setting)
 model = HiT_SRF(upscale=4) 
@@ -55,11 +55,10 @@ model = HiT_SRF(upscale=4)
 # load model (change repo_name according to your setting)
 repo_name = "XiangZ/hit-srf-4x"
 model = model.from_pretrained(repo_name)
-if cuda_flag:
-    model.cuda()
+model.to(device)
 
 # test and save results
-sr_results = model.infer_image("path-to-input-image", cuda=cuda_flag)
+sr_results = model.infer_image("path-to-input-image", device=device)
 cv2.imwrite("path-to-output-location", sr_results)
 ```