Update demo/model.py

demo/model.py

@@ -104,8 +104,7 @@ class Model_all:
         self.config = OmegaConf.load("configs/stable-diffusion/app.yaml")
         self.config.model.params.cond_stage_config.params.device = device
         self.base_model = load_model_from_config(self.config, "models/sd-v1-4.ckpt").to(device)
-        self.current_base_pose = 'sd-v1-4.ckpt'
-        self.current_base_sketch = 'sd-v1-4.ckpt'
+        self.current_base = 'sd-v1-4.ckpt'
         self.sampler = PLMSSampler(self.base_model)
 
         # sketch part
@@ -144,7 +143,7 @@ class Model_all:
 
     @torch.no_grad()
     def process_sketch(self, input_img, type_in, color_back, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model):
-        if self.current_base_sketch != base_model:
+        if self.current_base != base_model:
             ckpt = os.path.join("models", base_model)
             pl_sd = torch.load(ckpt, map_location="cpu")
             if "state_dict" in pl_sd:
@@ -152,7 +151,7 @@ class Model_all:
             else:
                 sd = pl_sd
             self.base_model.load_state_dict(sd, strict=False)
-            self.current_base_sketch = base_model
+            self.current_base = base_model
             # del sd
             # del pl_sd
         con_strength = int((1-con_strength)*50)
@@ -218,7 +217,7 @@ class Model_all:
 
     @torch.no_grad()
     def process_draw(self, input_img, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model):
-        if self.current_base_sketch != base_model:
+        if self.current_base != base_model:
             ckpt = os.path.join("models", base_model)
             pl_sd = torch.load(ckpt, map_location="cpu")
             if "state_dict" in pl_sd:
@@ -226,7 +225,7 @@ class Model_all:
             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.current_base_sketch = base_model
+            self.current_base = base_model
         con_strength = int((1-con_strength)*50)
         if fix_sample == 'True':
             seed_everything(42)
@@ -288,7 +287,7 @@ class Model_all:
 
     @torch.no_grad()
     def process_keypose(self, input_img, type_in, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model):
-        if self.current_base_pose != base_model:
+        if self.current_base != base_model:
             ckpt = os.path.join("models", base_model)
             pl_sd = torch.load(ckpt, map_location="cpu")
             if "state_dict" in pl_sd:
@@ -296,7 +295,7 @@ class Model_all:
             else:
                 sd = pl_sd
             self.base_model.load_state_dict(sd, strict=False)
-            self.current_base_pose = base_model
+            self.current_base = base_model
         con_strength = int((1-con_strength)*50)
         if fix_sample == 'True':
             seed_everything(42)