Update demo/model.py

demo/model.py

@@ -150,7 +150,9 @@ class Model_all:
                 sd = pl_sd["state_dict"]
             else:
                 sd = pl_sd
+            self.base_model = self.base_model.cpu()
             self.base_model.load_state_dict(sd, strict=False)
+            self.base_model = self.base_model.cuda()
             self.current_base = base_model
             # del sd
             # del pl_sd
@@ -224,7 +226,9 @@ class Model_all:
                 sd = pl_sd["state_dict"]
             else:
                 sd = pl_sd
-            self.base_model.load_state_dict(sd, strict=False) #load_model_from_config(config, os.path.join("models", base_model)).to(device)
+            self.base_model = self.base_model.cpu()
+            self.base_model.load_state_dict(sd, strict=False)
+            self.base_model = self.base_model.cuda()
             self.current_base = base_model
         con_strength = int((1-con_strength)*50)
         if fix_sample == 'True':
@@ -294,7 +298,9 @@ class Model_all:
                 sd = pl_sd["state_dict"]
             else:
                 sd = pl_sd
+            self.base_model = self.base_model.cpu()
             self.base_model.load_state_dict(sd, strict=False)
+            self.base_model = self.base_model.cuda()
             self.current_base = base_model
         con_strength = int((1-con_strength)*50)
         if fix_sample == 'True':