control processor handle null output

Signed-off-by: Vladimir Mandic <[email protected]>
vladmandic · Jan 10, 2025 · df60fe2 · df60fe2
1 parent 465d7ca
commit df60fe2
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/modules/control/proc/dwpose/__init__.py b/modules/control/proc/dwpose/__init__.py
@@ -13,12 +13,14 @@
 from modules.control.util import HWC3, resize_image
 from .draw import draw_bodypose, draw_handpose, draw_facepose
 checked_ok = False
+busy = False
 
 
 def check_dependencies():
-    global checked_ok # pylint: disable=global-statement
+    global checked_ok, busy # pylint: disable=global-statement
     debug = log.trace if os.environ.get('SD_DWPOSE_DEBUG', None) is not None else lambda *args, **kwargs: None
     packages = [
+        'termcolor',
         'openmim==0.3.9',
         'mmengine==0.10.4',
         'mmcv==2.1.0',
@@ -68,8 +70,13 @@ def __init__(self, det_config=None, det_ckpt=None, pose_config=None, pose_ckpt=N
         if not checked_ok:
             if not check_dependencies():
                 return
-        from .wholebody import Wholebody
-        self.pose_estimation = Wholebody(det_config, det_ckpt, pose_config, pose_ckpt, device)
+        Wholebody = None
+        try:
+            from .wholebody import Wholebody
+        except Exception as e:
+            log.error(f'DWPose: {e}')
+        if Wholebody is not None:
+            self.pose_estimation = Wholebody(det_config, det_ckpt, pose_config, pose_ckpt, device)
 
     def to(self, device):
         self.pose_estimation.to(device)
@@ -78,6 +85,7 @@ def to(self, device):
     def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", min_confidence=0.3, **kwargs):
         if self.pose_estimation is None:
             log.error("DWPose: not loaded")
+            return None
         input_image = cv2.cvtColor(np.array(input_image, dtype=np.uint8), cv2.COLOR_RGB2BGR)
 
         input_image = HWC3(input_image)

diff --git a/modules/control/processors.py b/modules/control/processors.py
@@ -253,6 +253,9 @@ def __call__(self, image_input: Image, mode: str = 'RGB', resize_mode: int = 0,
                 image_resized = image_input
             with devices.inference_context():
                 image_process = self.model(image_resized, **kwargs)
+            if image_process is None:
+                log.error(f'Control Processor: id="{self.processor_id}" no image')
+                return image_input
             if isinstance(image_process, np.ndarray):
                 if np.max(image_process) < 2:
                     image_process = (255.0 * image_process).astype(np.uint8)