From df60fe202deed9dd5b94d733a51918530f41c18e Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 10 Jan 2025 08:16:08 -0500 Subject: [PATCH] control processor handle null output Signed-off-by: Vladimir Mandic --- modules/control/proc/dwpose/__init__.py | 14 +++++++++++--- modules/control/processors.py | 3 +++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/modules/control/proc/dwpose/__init__.py b/modules/control/proc/dwpose/__init__.py index 4469e10c8..d8fdb9618 100644 --- 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 index 2f15c2d3f..aa1275d62 100644 --- 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)