cpu option

app.py

@@ -267,8 +267,6 @@ if __name__ == "__main__":
     parser.add_argument('--share', action="store_true")
     args = parser.parse_args()
     cfg = OmegaConf.load(args.config)
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    print(f"Using device: {device}")
-    matcher = zero_shot.Matcher(cfg, device)
+    matcher = zero_shot.Matcher(cfg)
     datadicts = None
     demo.launch(share=args.share)

zero_shot.py

@@ -101,9 +101,10 @@ class Matcher(object):
     def __init__(self, cfg):
         self.cfg = cfg
         self.device = torch.device(f'cuda:{cfg["gpu"]}' if torch.cuda.is_available() else 'cpu')
+        print(f"Using device: {self.device}")
         self.diffusion_model = None
         if self.cfg.get("sds", False):
-            self.diffusion_model = DiffModel(cfg["sds_conf"])
+            self.diffusion_model = DiffModel(cfg["sds_conf"], self.device)
         self.n_fmap = self.cfg["deepfeat_conf"]["fmap"]["n_fmap"]
         self.n_loop = 0
         if self.cfg.get("optimize", False):