first commit

annotator/midas/__init__.py

@@ -22,9 +22,9 @@ from .api import MiDaSInference
 
 class MidasDetector:
     def __init__(self):
-        self.model = MiDaSInference(model_type="dpt_large").cuda()
+        self.model = MiDaSInference(model_type="dpt_hybrid").cuda()
 
-    def __call__(self, input_image, a=np.pi * 0.2, bg_th=0.02):
+    def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1):
         assert input_image.ndim == 3
         image_depth = input_image
         with torch.no_grad():

annotator/midas/api.py

@@ -30,8 +30,8 @@ ISL_PATHS = {
     "midas_v21_small": "",
 }
 
-# remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
-remote_model_path = "https://storage.googleapis.com/sfr-unicontrol-data-research/annotator/ckpts/dpt_large_384.pt"  #"https://huggingface.co/Salesforce/UniControl/blob/main/annotator/ckpts/dpt_large_384.pt"
+remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
+# remote_model_path = "https://storage.googleapis.com/sfr-unicontrol-data-research/annotator/ckpts/dpt_large_384.pt"  #"https://huggingface.co/Salesforce/UniControl/blob/main/annotator/ckpts/dpt_large_384.pt"
 
 def disabled_train(self, mode=True):
     """Overwrite model.train with this function to make sure train/eval mode

annotator/openpose/__init__.py

@@ -31,15 +31,15 @@ hand_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/ann
 class OpenposeDetector:
     def __init__(self):
         body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")
-        hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")
+        # hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")
 
         if not os.path.exists(hand_modelpath):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(body_model_path, model_dir=annotator_ckpts_path)
-            load_file_from_url(hand_model_path, model_dir=annotator_ckpts_path)
+            # load_file_from_url(hand_model_path, model_dir=annotator_ckpts_path)
 
         self.body_estimation = Body(body_modelpath)
-        self.hand_estimation = Hand(hand_modelpath)
+        # self.hand_estimation = Hand(hand_modelpath)
 
     def __call__(self, oriImg, hand=False):
         oriImg = oriImg[:, :, ::-1].copy()
@@ -47,13 +47,13 @@ class OpenposeDetector:
             candidate, subset = self.body_estimation(oriImg)
             canvas = np.zeros_like(oriImg)
             canvas = util.draw_bodypose(canvas, candidate, subset)
-            if hand:
-                hands_list = util.handDetect(candidate, subset, oriImg)
-                all_hand_peaks = []
-                for x, y, w, is_left in hands_list:
-                    peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :])
-                    peaks[:, 0] = np.where(peaks[:, 0] == 0, peaks[:, 0], peaks[:, 0] + x)
-                    peaks[:, 1] = np.where(peaks[:, 1] == 0, peaks[:, 1], peaks[:, 1] + y)
-                    all_hand_peaks.append(peaks)
-                canvas = util.draw_handpose(canvas, all_hand_peaks)
+            # if hand:
+            #     hands_list = util.handDetect(candidate, subset, oriImg)
+            #     all_hand_peaks = []
+            #     for x, y, w, is_left in hands_list:
+            #         peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :])
+            #         peaks[:, 0] = np.where(peaks[:, 0] == 0, peaks[:, 0], peaks[:, 0] + x)
+            #         peaks[:, 1] = np.where(peaks[:, 1] == 0, peaks[:, 1], peaks[:, 1] + y)
+            #         all_hand_peaks.append(peaks)
+            #     canvas = util.draw_handpose(canvas, all_hand_peaks)
             return canvas, dict(candidate=candidate.tolist(), subset=subset.tolist())

model.py

@@ -315,7 +315,6 @@ class ControlNet(nn.Module):
             num_heads = ch // num_head_channels
             dim_head = num_head_channels
         if legacy:
-            # num_heads = 1
             dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
         self.middle_block = TimestepEmbedSequential(
             ResBlock(
@@ -360,7 +359,6 @@ class ControlNet(nn.Module):
         hint -> 4, 3, 512, 512
         context - > 4, 77, 768
         '''
-        BS = 1 # x.shape[0], one batch one task
         BS_Real = x.shape[0]
         if kwargs is not None:
             task_name = kwargs['task']['name']
@@ -407,7 +405,6 @@ class ControlLDM(LatentDiffusion):
         super().__init__(*args, **kwargs)
         self.mapping_task = {"control_hed": "hed edge to image", "control_canny": "canny edge to image", "control_seg": "segmentation map to image", "control_depth": "depth map to image", "control_normal": "normal surface map to image", "control_img": "image editing", "control_openpose": "human pose skeleton to image", "control_hedsketch": "sketch to image", "control_bbox": "bounding box to image", "control_outpainting": "image outpainting", "control_grayscale": "gray image to color image", "control_blur": "deblur image to clean image", "control_inpainting": "image inpainting"}
         self.all_tasks_num = len(self.mapping_task)
-#         self.task_weight_all = nn.Parameter(torch.zeros(self.all_tasks_num,), requires_grad=True)
         self.task_loss_ema = torch.zeros(self.all_tasks_num,)
 
         self.control_model = instantiate_from_config(control_stage_config) # -> ControlNet
@@ -552,11 +549,6 @@ class ControlLDM(LatentDiffusion):
         c_cat, c = c["c_concat"][0][:N], c["c_crossattn"][0][:N]
         N = min(z.shape[0], N)
         n_row = min(z.shape[0], n_row)
-#         log["reconstruction"] = self.decode_first_stage(z)
-#         log["control"] = c_cat * 2.0 - 1.0
-#         log["conditioning"] = log_txt_as_img((512, 512), batch[self.cond_stage_key], size=16)
-
-
 
         uc_cross = self.get_unconditional_conditioning(N)
         uc_cat = c_cat  # torch.zeros_like(c_cat)