Update demo/model.py

demo/model.py

@@ -17,6 +17,7 @@ import cv2
 import numpy as np
 import torch.nn.functional as F
 
+
 def preprocessing(image, device):
     # Resize
     scale = 640 / max(image.shape[:2])
@@ -39,6 +40,7 @@ def preprocessing(image, device):
 
     return image, raw_image
 
+
 def imshow_keypoints(img,
                      pose_result,
                      skeleton=None,
@@ -138,18 +140,22 @@ class Model_all:
                                     use_conv=False).to(device)
         self.model_sketch.load_state_dict(torch.load("models/t2iadapter_sketch_sd14v1.pth", map_location=device))
         self.model_edge = pidinet().to(device)
-        self.model_edge.load_state_dict({k.replace('module.', ''): v for k, v in torch.load('models/table5_pidinet.pth', map_location=device)['state_dict'].items()})
+        self.model_edge.load_state_dict({k.replace('module.', ''): v for k, v in
+                                         torch.load('models/table5_pidinet.pth', map_location=device)[
+                                             'state_dict'].items()})
 
         # segmentation part
         self.model_seger = seger().to(device)
         self.model_seger.eval()
         self.coler = Colorize(n=182)
-        self.model_seg = Adapter(cin=int(3*64), channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False).to(device)
+        self.model_seg = Adapter(cin=int(3 * 64), channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True,
+                                 use_conv=False).to(device)
         self.model_seg.load_state_dict(torch.load("models/t2iadapter_seg_sd14v1.pth", map_location=device))
         self.depth_model = MiDaSInference(model_type='dpt_hybrid').to(device)
 
         # depth part
-        self.model_depth = Adapter(cin=3*64, channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False).to(device)
+        self.model_depth = Adapter(cin=3 * 64, channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True,
+                                   use_conv=False).to(device)
         self.model_depth.load_state_dict(torch.load("models/t2iadapter_depth_sd14v1.pth", map_location=device))
 
         # keypose part
@@ -183,7 +189,7 @@ class Model_all:
                                 [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]]
-    
+
     def load_vae(self):
         vae_sd = torch.load(os.path.join('models', 'anything-v4.0.vae.pt'), map_location="cuda")
         sd = vae_sd["state_dict"]
@@ -254,7 +260,7 @@ class Model_all:
 
     @torch.no_grad()
     def process_depth(self, input_img, type_in, prompt, neg_prompt, pos_prompt, fix_sample, scale,
-                       con_strength, base_model):
+                      con_strength, base_model):
         if self.current_base != base_model:
             ckpt = os.path.join("models", base_model)
             pl_sd = torch.load(ckpt, map_location="cuda")
@@ -312,7 +318,8 @@ class Model_all:
         return [im_depth, x_samples_ddim]
 
     @torch.no_grad()
-    def process_depth_keypose(self, input_img_depth, input_img_keypose, type_in_depth, type_in_keypose, w_depth, w_keypose, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model):
+    def process_depth_keypose(self, input_img_depth, input_img_keypose, type_in_depth, type_in_keypose, w_depth,
+                              w_keypose, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model):
         if self.current_base != base_model:
             ckpt = os.path.join("models", base_model)
             pl_sd = torch.load(ckpt, map_location="cuda")
@@ -343,8 +350,7 @@ class Model_all:
 
         # get keypose
         if type_in_keypose == 'Keypose':
-            im_keypose_out = im_keypose.copy()
-            pose = img2tensor(im_keypose).unsqueeze(0) / 255.
+            im_keypose_out = im_keypose.copy()[:,:,::-1]
         elif type_in_keypose == 'Image':
             image = im_keypose.copy()
             im_keypose = img2tensor(im_keypose).unsqueeze(0) / 255.
@@ -378,7 +384,7 @@ class Model_all:
                 pose_link_color=self.pose_link_color,
                 radius=2,
                 thickness=2)
-            im_keypose_out = im_keypose_out.astype(np.uint8)[:,:,::-1]
+            im_keypose_out = im_keypose_out.astype(np.uint8)[:, :, ::-1]
 
         # extract condition features
         c = self.base_model.get_learned_conditioning([prompt + ', ' + pos_prompt])
@@ -387,7 +393,8 @@ class Model_all:
         pose = img2tensor(im_keypose_out, bgr2rgb=True, float32=True) / 255.
         pose = pose.unsqueeze(0)
         features_adapter_keypose = self.model_pose(pose.to(self.device))
-        features_adapter = [f_d*w_depth + f_k*w_keypose for f_d, f_k in zip(features_adapter_depth, features_adapter_keypose)]
+        features_adapter = [f_d * w_depth + f_k * w_keypose for f_d, f_k in
+                            zip(features_adapter_depth, features_adapter_keypose)]
         shape = [4, 64, 64]
 
         # sampling
@@ -416,7 +423,7 @@ class Model_all:
 
     @torch.no_grad()
     def process_seg(self, input_img, type_in, prompt, neg_prompt, pos_prompt, fix_sample, scale,
-                       con_strength, base_model):
+                    con_strength, base_model):
         if self.current_base != base_model:
             ckpt = os.path.join("models", base_model)
             pl_sd = torch.load(ckpt, map_location="cuda")
@@ -450,10 +457,10 @@ class Model_all:
             labelmap = np.argmax(probs, axis=0)
 
             labelmap = self.coler(labelmap)
-            labelmap = np.transpose(labelmap, (1,2,0))
+            labelmap = np.transpose(labelmap, (1, 2, 0))
             labelmap = cv2.resize(labelmap, (512, 512))
-            labelmap = img2tensor(labelmap, bgr2rgb=False, float32=True)/255.
-            im_seg = tensor2img(labelmap)[:,:,::-1]
+            labelmap = img2tensor(labelmap, bgr2rgb=False, float32=True) / 255.
+            im_seg = tensor2img(labelmap)[:, :, ::-1]
             labelmap = labelmap.unsqueeze(0)
 
         # extract condition features
@@ -564,8 +571,7 @@ class Model_all:
         im = cv2.resize(input_img, (512, 512))
 
         if type_in == 'Keypose':
-            im_pose = im.copy()
-            im = img2tensor(im).unsqueeze(0) / 255.
+            im_pose = im.copy()[:,:,::-1]
         elif type_in == 'Image':
             image = im.copy()
             im = img2tensor(im).unsqueeze(0) / 255.
@@ -599,7 +605,7 @@ class Model_all:
                 pose_link_color=self.pose_link_color,
                 radius=2,
                 thickness=2)
-        im_pose = cv2.resize(im_pose, (512, 512))
+        # im_pose = cv2.resize(im_pose, (512, 512))
 
         # extract condition features
         c = self.base_model.get_learned_conditioning([prompt + ', ' + pos_prompt])