Update annotator/hed/__init__.py

annotator/hed/__init__.py

@@ -100,13 +100,20 @@ class HEDdetector:
         if not os.path.exists(modelpath):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)
-        self.netNetwork = Network(modelpath).cuda().eval()
+        if torch.cuda.is_available():
+            self.netNetwork = Network(modelpath).cuda().eval()
+        else:
+            self.netNetwork = Network(modelpath).eval()
+            
 
     def __call__(self, input_image):
         assert input_image.ndim == 3
         input_image = input_image[:, :, ::-1].copy()
         with torch.no_grad():
-            image_hed = torch.from_numpy(input_image).float().cuda()
+            if torch.cuda.is_available():
+                image_hed = torch.from_numpy(input_image).float().cuda()
+            else:
+                image_hed = torch.from_numpy(input_image).float()
             image_hed = image_hed / 255.0
             image_hed = rearrange(image_hed, 'h w c -> 1 c h w')
             edge = self.netNetwork(image_hed)[0]