Update model.py

model.py

@@ -19,6 +19,23 @@ import shlex
 import subprocess
 import sys
 
+import mmcv
+from mmdet.apis import inference_detector, init_detector
+from mmpose.apis import (inference_top_down_pose_model, init_pose_model, process_mmdet_results, vis_pose_result)
+
+skeleton = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9], [8, 10],
+            [1, 2], [0, 1], [0, 2], [1, 3], [2, 4], [3, 5], [4, 6]]
+
+pose_kpt_color = [[51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [0, 255, 0],
+                  [255, 128, 0], [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0],
+                  [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0]]
+
+pose_link_color = [[0, 255, 0], [0, 255, 0], [255, 128, 0], [255, 128, 0],
+                   [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [0, 255, 0], [255, 128, 0],
+                   [0, 255, 0], [255, 128, 0], [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255],
+                   [51, 153, 255], [51, 153, 255], [51, 153, 255]]
+
+
 sys.path.append('T2I-Adapter')
 
 config_path =  'https://github.com/TencentARC/T2I-Adapter/raw/main/configs/stable-diffusion/'
@@ -101,6 +118,10 @@ class Model:
     def process_sketch(self, input_img, type_in, color_back, prompt, neg_prompt, fix_sample, scale, con_strength, base_model):    
         global current_base
         device = 'cuda' 
+        if base_model == 'sd-v1-4.ckpt':
+            model = self.model
+        else:
+            model = self.model_anything
         # if current_base != base_model:
         #     ckpt = os.path.join("models", base_model)
         #     pl_sd = torch.load(ckpt, map_location="cpu")
@@ -132,8 +153,8 @@ class Model:
             im = im.float()
             im_edge = tensor2img(im)
     
-        c = self.model.get_learned_conditioning([prompt])
-        nc = self.model.get_learned_conditioning([neg_prompt])
+        c = model.get_learned_conditioning([prompt])
+        nc = model.get_learned_conditioning([neg_prompt])
         
         with torch.no_grad():
             # extract condition features
@@ -155,7 +176,7 @@ class Model:
                                         mode = 'sketch',
                                         con_strength = con_strength)
     
-        x_samples_ddim = self.model.decode_first_stage(samples_ddim)
+        x_samples_ddim = model.decode_first_stage(samples_ddim)
         x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
         x_samples_ddim = x_samples_ddim.permute(0, 2, 3, 1).cpu().numpy()[0]
         x_samples_ddim = 255.*x_samples_ddim
@@ -166,7 +187,11 @@ class Model:
     @torch.inference_mode()
     def process_pose(self, input_img, prompt, neg_prompt, fix_sample, scale, con_strength, base_model):    
         global current_base
-        device = 'cuda' 
+        device = 'cuda'
+        if base_model == 'sd-v1-4.ckpt':
+            model = self.model
+        else:
+            model = self.model_anything
         # if current_base != base_model:
         #     ckpt = os.path.join("models", base_model)
         #     pl_sd = torch.load(ckpt, map_location="cpu")
@@ -186,8 +211,8 @@ class Model:
 
         im_pose = tensor2img(pose)
     
-        c = self.model.get_learned_conditioning([prompt])
-        nc = self.model.get_learned_conditioning([neg_prompt])
+        c = model.get_learned_conditioning([prompt])
+        nc = model.get_learned_conditioning([neg_prompt])
         
         with torch.no_grad():
             # extract condition features
@@ -209,11 +234,10 @@ class Model:
                                         mode = 'sketch',
                                         con_strength = con_strength)
     
-        x_samples_ddim = self.model.decode_first_stage(samples_ddim)
+        x_samples_ddim = model.decode_first_stage(samples_ddim)
         x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
         x_samples_ddim = x_samples_ddim.permute(0, 2, 3, 1).cpu().numpy()[0]
         x_samples_ddim = 255.*x_samples_ddim
         x_samples_ddim = x_samples_ddim.astype(np.uint8)
     
-        return [im_pose, x_samples_ddim]
-
+        return [im_pose, x_samples_ddim]