try fix the preprocess OOM issue

.gitignore

@@ -0,0 +1,17 @@
+__pycache__
+.vscode/
+data/
+data
+log/
+log
+pretrained/
+pretrained
+assets/nosync/
+assets/demo/temp/temp_*
+*.out
+gradio_cached_examples/
+src/*/build
+src/*/dist
+src/*/*.egg-info/
+extensions/
+extensions

app.py

@@ -29,13 +29,13 @@ n_sample_image = 1
 #     ['canny_v11p'        , ('canny'   , 'pretrained/controlnet/control_v11p_sd15_canny_slimmed.safetensors')],
 #     ['depth'             , ('depth'   , 'pretrained/controlnet/control_sd15_depth_slimmed.safetensors')],
 #     ['hed'               , ('hed'     , 'pretrained/controlnet/control_sd15_hed_slimmed.safetensors')],
+#     ['softedge_v11p'     , ('hed'     , 'pretrained/controlnet/control_v11p_sd15_softedge_slimmed.safetensors')],
 #     ['mlsd'              , ('mlsd'    , 'pretrained/controlnet/control_sd15_mlsd_slimmed.safetensors')],
-#     ['mlsd_v11p'         , ('mlsd'    , 'pretrained/controlnet/control_v11p_sd15_mlsd_slimmed.safetensors')],
-#     ['normal'            , ('normal'  , 'pretrained/controlnet/control_sd15_normal_slimmed.safetensors')],
+#     # ['mlsd_v11p'         , ('mlsd'    , 'pretrained/controlnet/control_v11p_sd15_mlsd_slimmed.safetensors')],
+#     # ['normal'            , ('normal'  , 'pretrained/controlnet/control_sd15_normal_slimmed.safetensors')],
 #     ['openpose'          , ('openpose', 'pretrained/controlnet/control_sd15_openpose_slimmed.safetensors')],
 #     ['openpose_v11p'     , ('openpose', 'pretrained/controlnet/control_v11p_sd15_openpose_slimmed.safetensors')],
 #     ['scribble'          , ('scribble', 'pretrained/controlnet/control_sd15_scribble_slimmed.safetensors')],
-#     ['softedge_v11p'     , ('scribble', 'pretrained/controlnet/control_v11p_sd15_softedge_slimmed.safetensors')],
 #     ['seg'               , ('none'    , 'pretrained/controlnet/control_sd15_seg_slimmed.safetensors')],
 #     ['lineart_v11p'      , ('none'    , 'pretrained/controlnet/control_v11p_sd15_lineart_slimmed.safetensors')],
 #     ['lineart_anime_v11p', ('none'    , 'pretrained/controlnet/control_v11p_sd15s2_lineart_anime_slimmed.safetensors')],
@@ -45,7 +45,7 @@ controlnet_path = OrderedDict([
     ['canny'             , ('canny'   , hf_hub_download('shi-labs/prompt-free-diffusion', 'pretrained/controlnet/control_sd15_canny_slimmed.safetensors'))],
     # ['canny_v11p'        , ('canny'   , hf_hub_download('shi-labs/prompt-free-diffusion', 'pretrained/controlnet/control_v11p_sd15_canny_slimmed.safetensors'))],
     ['depth'             , ('depth'   , hf_hub_download('shi-labs/prompt-free-diffusion', 'pretrained/controlnet/control_sd15_depth_slimmed.safetensors'))],
-    # ['hed'               , ('hed'     , hf_hub_download('shi-labs/prompt-free-diffusion', 'pretrained/controlnet/control_sd15_hed_slimmed.safetensors'))],
+    ['hed'               , ('hed'     , hf_hub_download('shi-labs/prompt-free-diffusion', 'pretrained/controlnet/control_sd15_hed_slimmed.safetensors'))],
     ['mlsd'              , ('mlsd'    , hf_hub_download('shi-labs/prompt-free-diffusion', 'pretrained/controlnet/control_sd15_mlsd_slimmed.safetensors'))],
     # ['mlsd_v11p'         , ('mlsd'    , hf_hub_download('shi-labs/prompt-free-diffusion', 'pretrained/controlnet/control_v11p_sd15_mlsd_slimmed.safetensors'))],
     # ['normal'            , ('normal'  , hf_hub_download('shi-labs/prompt-free-diffusion', 'pretrained/controlnet/control_sd15_normal_slimmed.safetensors'))],
@@ -61,7 +61,7 @@ controlnet_path = OrderedDict([
 preprocess_method = [
     'canny'                ,
     'depth'                ,
-    # 'hed'                  ,
+    'hed'                  ,
     'mlsd'                 ,
     # 'normal'               ,
     'openpose'             ,

lib/model_zoo/controlnet.py

@@ -296,12 +296,12 @@ class ControlNet(nn.Module):
         if type == 'none' or type is None:
             return None
 
-        elif type in ['input', 'shuffle_v11e']:
+        elif type in ['input']:
             y_torch = torch.stack([tvtrans.ToTensor()(xi) for xi in x_list])
             y_torch = y_torch.to(device).to(torch.float32)
             return y_torch
 
-        elif type in ['canny', 'canny_v11p']:
+        elif type in ['canny']:
             low_threshold = kwargs.pop('low_threshold', 100)
             high_threshold = kwargs.pop('high_threshold', 200)
             from .controlnet_annotator.canny import apply_canny
@@ -320,7 +320,7 @@ class ControlNet(nn.Module):
             unload_midas_model()
             return y_torch
 
-        elif type in ['hed', 'softedge_v11p']:
+        elif type in ['hed']:
             from .controlnet_annotator.hed import apply_hed, unload_hed_model
             y_list = [apply_hed(np.array(xi), device=device) for xi in x_list]
             y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
@@ -349,7 +349,7 @@ class ControlNet(nn.Module):
             unload_midas_model()
             return y_torch
 
-        elif type in ['openpose', 'openpose_v11p']:
+        elif type in ['openpose']:
             from .controlnet_annotator.openpose import OpenposeModel
             from functools import partial
             wrapper = OpenposeModel()
@@ -359,10 +359,10 @@ class ControlNet(nn.Module):
             y_list = [apply_openpose(np.array(xi)) for xi in x_list]
             y_torch = torch.stack([tvtrans.ToTensor()(yi.copy()) for yi in y_list])
             y_torch = y_torch.to(device).to(torch.float32)
-            OpenposeModel.unload()
+            wrapper.unload()
             return y_torch
 
-        elif type in ['openpose_withface', 'openpose_withface_v11p']:
+        elif type in ['openpose_withface']:
             from .controlnet_annotator.openpose import OpenposeModel
             from functools import partial
             wrapper = OpenposeModel()
@@ -372,10 +372,10 @@ class ControlNet(nn.Module):
             y_list = [apply_openpose(np.array(xi)) for xi in x_list]
             y_torch = torch.stack([tvtrans.ToTensor()(yi.copy()) for yi in y_list])
             y_torch = y_torch.to(device).to(torch.float32)
-            OpenposeModel.unload()
+            wrapper.unload()
             return y_torch
 
-        elif type in ['openpose_withfacehand', 'openpose_withfacehand_v11p']:
+        elif type in ['openpose_withfacehand']:
             from .controlnet_annotator.openpose import OpenposeModel
             from functools import partial
             wrapper = OpenposeModel()
@@ -385,7 +385,7 @@ class ControlNet(nn.Module):
             y_list = [apply_openpose(np.array(xi)) for xi in x_list]
             y_torch = torch.stack([tvtrans.ToTensor()(yi.copy()) for yi in y_list])
             y_torch = y_torch.to(device).to(torch.float32)
-            OpenposeModel.unload()
+            wrapper.unload()
             return y_torch
 
         elif type == 'scribble':

lib/model_zoo/controlnet_annotator/midas/__init__.py

@@ -16,8 +16,7 @@ def apply_midas(input_image, a=np.pi * 2.0, bg_th=0.1, device='cpu'):
     global model
     if model is None:
         model = MiDaSInference(model_type="dpt_hybrid")
-        model = model.to(device)
-    
+    model = model.to(device)    
     assert input_image.ndim == 3
     image_depth = input_image
     with torch.no_grad():

lib/model_zoo/controlnet_annotator/openpose/__init__.py

@@ -18,6 +18,8 @@ from .body import Body, BodyResult, Keypoint
 from .hand import Hand
 from .face import Face
 
+openposemodel = None
+
 models_path = "pretrained/controlnet/preprocess"
 
 from typing import NamedTuple, Tuple, List, Callable, Union
@@ -170,11 +172,21 @@ class OpenposeDetector:
         """
         Unload the Openpose models by moving them to the CPU.
         """
+        self.device = "cpu"
         if self.body_estimation is not None:
             self.body_estimation.model.to("cpu")
             self.hand_estimation.model.to("cpu")
             self.face_estimation.model.to("cpu")
 
+    def set_device(self, device):
+        self.device = device
+        if self.body_estimation is not None:
+            self.body_estimation.model.to(device)
+        if self.hand_estimation is not None:
+            self.hand_estimation.model.to(device)
+        if self.face_estimation is not None:
+            self.face_estimation.model.to(device)
+
     def detect_hands(self, body: BodyResult, oriImg) -> Tuple[Union[HandResult, None], Union[HandResult, None]]:
         left_hand = None
         right_hand = None
@@ -291,7 +303,7 @@ class OpenposeDetector:
 
 class OpenposeModel(object):
     def __init__(self) -> None:
-        self.model_openpose = None
+        pass
 
     def run_model(
             self,
@@ -302,13 +314,17 @@ class OpenposeModel(object):
             json_pose_callback: Callable[[str], None] = None,
             device = 'cpu', ):
 
+        global openposemodel
+
         if json_pose_callback is None:
             json_pose_callback = lambda x: None
 
-        if self.model_openpose is None:
-            self.model_openpose = OpenposeDetector(device=device)
+        if openposemodel is None:
+            openposemodel = OpenposeDetector(device=device)
+        else:
+            openposemodel.set_device(device)
 
-        return self.model_openpose(
+        return openposemodel(
             img,
             include_body=include_body,
             include_hand=include_hand,
@@ -316,5 +332,6 @@ class OpenposeModel(object):
             json_pose_callback=json_pose_callback)
 
     def unload(self):
-        if self.model_openpose is not None:
-            self.model_openpose.unload_model()
+        global openposemodel
+        if openposemodel is not None:
+            openposemodel.unload_model()