modules/control/processors.py

import os
import time
import numpy as np
from PIL import Image
from modules.shared import log
from modules.errors import display
from modules import devices, images

from modules.control.proc.hed import HEDdetector
from modules.control.proc.canny import CannyDetector
from modules.control.proc.edge import EdgeDetector
from modules.control.proc.lineart import LineartDetector
from modules.control.proc.lineart_anime import LineartAnimeDetector
from modules.control.proc.pidi import PidiNetDetector
from modules.control.proc.mediapipe_face import MediapipeFaceDetector
from modules.control.proc.shuffle import ContentShuffleDetector
from modules.control.proc.leres import LeresDetector
from modules.control.proc.midas import MidasDetector
from modules.control.proc.mlsd import MLSDdetector
from modules.control.proc.normalbae import NormalBaeDetector
from modules.control.proc.openpose import OpenposeDetector
from modules.control.proc.dwpose import DWposeDetector
from modules.control.proc.segment_anything import SamDetector
from modules.control.proc.zoe import ZoeDetector
from modules.control.proc.marigold import MarigoldDetector
from modules.control.proc.dpt import DPTDetector
from modules.control.proc.glpn import GLPNDetector
from modules.control.proc.depth_anything import DepthAnythingDetector


models = {}
cache_dir = 'models/control/processors'
debug = log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
debug('Trace: CONTROL')
config = {
    # placeholder
    'None': {},
    # pose models
    'OpenPose': {'class': OpenposeDetector, 'checkpoint': True, 'params': {'include_body': True, 'include_hand': False, 'include_face': False}},
    'DWPose': {'class': DWposeDetector, 'checkpoint': False, 'model': 'Tiny', 'params': {'min_confidence': 0.3}},
    'MediaPipe Face': {'class': MediapipeFaceDetector, 'checkpoint': False, 'params': {'max_faces': 1, 'min_confidence': 0.5}},
    # outline models
    'Canny': {'class': CannyDetector, 'checkpoint': False, 'params': {'low_threshold': 100, 'high_threshold': 200}},
    'Edge': {'class': EdgeDetector, 'checkpoint': False, 'params': {'pf': True, 'mode': 'edge'}},
    'LineArt Realistic': {'class': LineartDetector, 'checkpoint': True, 'params': {'coarse': False}},
    'LineArt Anime': {'class': LineartAnimeDetector, 'checkpoint': True, 'params': {}},
    'HED': {'class': HEDdetector, 'checkpoint': True, 'params': {'scribble': False, 'safe': False}},
    'PidiNet': {'class': PidiNetDetector, 'checkpoint': True, 'params': {'scribble': False, 'safe': False, 'apply_filter': False}},
    # depth models
    'Midas Depth Hybrid': {'class': MidasDetector, 'checkpoint': True, 'params': {'bg_th': 0.1, 'depth_and_normal': False}},
    'Leres Depth': {'class': LeresDetector, 'checkpoint': True, 'params': {'boost': False, 'thr_a':0, 'thr_b':0}},
    'Zoe Depth': {'class': ZoeDetector, 'checkpoint': True, 'params': {'gamma_corrected': False}, 'load_config': {'pretrained_model_or_path': 'halffried/gyre_zoedepth', 'filename': 'ZoeD_M12_N.safetensors', 'model_type': "zoedepth"}},
    'Marigold Depth': {'class': MarigoldDetector, 'checkpoint': True, 'params': {'denoising_steps': 10, 'ensemble_size': 10, 'processing_res': 512, 'match_input_res': True, 'color_map': 'None'}, 'load_config': {'pretrained_model_or_path': 'Bingxin/Marigold'}},
    'Normal Bae': {'class': NormalBaeDetector, 'checkpoint': True, 'params': {}},
    # segmentation models
    'SegmentAnything': {'class': SamDetector, 'checkpoint': True, 'model': 'Base', 'params': {}},
    # other models
    'MLSD': {'class': MLSDdetector, 'checkpoint': True, 'params': {'thr_v': 0.1, 'thr_d': 0.1}},
    'Shuffle': {'class': ContentShuffleDetector, 'checkpoint': False, 'params': {}},
    'DPT Depth Hybrid': {'class': DPTDetector, 'checkpoint': False, 'params': {}},
    'GLPN Depth': {'class': GLPNDetector, 'checkpoint': False, 'params': {}},
    'Depth Anything': {'class': DepthAnythingDetector, 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'LiheYoung/depth_anything_vitl14' }, 'params': { 'color_map': 'inferno' }},
    # 'Midas Depth Large': {'class': MidasDetector, 'checkpoint': True, 'params': {'bg_th': 0.1, 'depth_and_normal': False}, 'load_config': {'pretrained_model_or_path': 'Intel/dpt-large', 'model_type': "dpt_large", 'filename': ''}},
    # 'Zoe Depth Zoe': {'class': ZoeDetector, 'checkpoint': True, 'params': {}},
    # 'Zoe Depth NK': {'class': ZoeDetector, 'checkpoint': True, 'params': {}, 'load_config': {'pretrained_model_or_path': 'halffried/gyre_zoedepth', 'filename': 'ZoeD_M12_NK.safetensors', 'model_type': "zoedepth_nk"}},
}


def list_models(refresh=False):
    global models # pylint: disable=global-statement
    if not refresh and len(models) > 0:
        return models
    models = list(config)
    debug(f'Control list processors: path={cache_dir} models={models}')
    return models


def update_settings(*settings):
    debug(f'Control settings: {settings}')
    def update(what, val):
        processor_id = what[0]
        if len(what) == 2 and config[processor_id][what[1]] != val:
            config[processor_id][what[1]] = val
            config[processor_id]['dirty'] = True
            log.debug(f'Control settings: id="{processor_id}" {what[-1]}={val}')
        elif len(what) == 3 and config[processor_id][what[1]][what[2]] != val:
            config[processor_id][what[1]][what[2]] = val
            config[processor_id]['dirty'] = True
            log.debug(f'Control settings: id="{processor_id}" {what[-1]}={val}')
        elif len(what) == 4 and config[processor_id][what[1]][what[2]][what[3]] != val:
            config[processor_id][what[1]][what[2]][what[3]] = val
            config[processor_id]['dirty'] = True
            log.debug(f'Control settings: id="{processor_id}" {what[-1]}={val}')

    update(['HED', 'params', 'scribble'], settings[0])
    update(['Midas Depth Hybrid', 'params', 'bg_th'], settings[1])
    update(['Midas Depth Hybrid', 'params', 'depth_and_normal'], settings[2])
    update(['MLSD', 'params', 'thr_v'], settings[3])
    update(['MLSD', 'params', 'thr_d'], settings[4])
    update(['OpenPose', 'params', 'include_body'], settings[5])
    update(['OpenPose', 'params', 'include_hand'], settings[6])
    update(['OpenPose', 'params', 'include_face'], settings[7])
    update(['PidiNet', 'params', 'scribble'], settings[8])
    update(['PidiNet', 'params', 'apply_filter'], settings[9])
    update(['LineArt Realistic', 'params', 'coarse'], settings[10])
    update(['Leres Depth', 'params', 'boost'], settings[11])
    update(['Leres Depth', 'params', 'thr_a'], settings[12])
    update(['Leres Depth', 'params', 'thr_b'], settings[13])
    update(['MediaPipe Face', 'params', 'max_faces'], settings[14])
    update(['MediaPipe Face', 'params', 'min_confidence'], settings[15])
    update(['Canny', 'params', 'low_threshold'], settings[16])
    update(['Canny', 'params', 'high_threshold'], settings[17])
    update(['DWPose', 'model'], settings[18])
    update(['DWPose', 'params', 'min_confidence'], settings[19])
    update(['SegmentAnything', 'model'], settings[20])
    update(['Edge', 'params', 'pf'], settings[21])
    update(['Edge', 'params', 'mode'], settings[22])
    update(['Zoe Depth', 'params', 'gamma_corrected'], settings[23])
    update(['Marigold Depth', 'params', 'color_map'], settings[24])
    update(['Marigold Depth', 'params', 'denoising_steps'], settings[25])
    update(['Marigold Depth', 'params', 'ensemble_size'], settings[26])
    update(['Depth Anything', 'params', 'color_map'], settings[27])


class Processor():
    def __init__(self, processor_id: str = None, resize = True):
        self.model = None
        self.processor_id = None
        self.override = None
        self.resize = resize
        self.reset()
        self.config(processor_id)
        if processor_id is not None:
            self.load()

    def reset(self, processor_id: str = None):
        if self.model is not None:
            debug(f'Control Processor unloaded: id="{self.processor_id}"')
        self.model = None
        self.processor_id = processor_id
        # self.override = None
        devices.torch_gc()
        self.load_config = { 'cache_dir': cache_dir }

    def config(self, processor_id = None):
        if processor_id is not None:
            self.processor_id = processor_id
        from_config = config.get(self.processor_id, {}).get('load_config', None)
        """
        if load_config is not None:
            for k, v in load_config.items():
                self.load_config[k] = v
        """
        if from_config is not None:
            for k, v in from_config.items():
                self.load_config[k] = v

    def load(self, processor_id: str = None) -> str:
        try:
            t0 = time.time()
            processor_id = processor_id or self.processor_id
            if processor_id is None or processor_id == 'None':
                self.reset()
                return ''
            if self.processor_id != processor_id:
                self.reset()
                self.config(processor_id)
            cls = config[processor_id]['class']
            log.debug(f'Control Processor loading: id="{processor_id}" class={cls.__name__}')
            debug(f'Control Processor config={self.load_config}')
            if 'DWPose' in processor_id:
                det_ckpt = 'https://download.openmmlab.com/mmdetection/v2.0/yolox/yolox_l_8x8_300e_coco/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth'
                if 'Tiny' == config['DWPose']['model']:
                    pose_config = 'config/rtmpose-t_8xb64-270e_coco-ubody-wholebody-256x192.py'
                    pose_ckpt = 'https://huggingface.co/yzd-v/DWPose/resolve/main/dw-tt_ucoco.pth'
                elif 'Medium' == config['DWPose']['model']:
                    pose_config = 'config/rtmpose-m_8xb64-270e_coco-ubody-wholebody-256x192.py'
                    pose_ckpt = 'https://huggingface.co/yzd-v/DWPose/resolve/main/dw-mm_ucoco.pth'
                elif 'Large' == config['DWPose']['model']:
                    pose_config = 'config/rtmpose-l_8xb32-270e_coco-ubody-wholebody-384x288.py'
                    pose_ckpt = 'https://huggingface.co/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.pth'
                else:
                    log.error(f'Control Processor load failed: id="{processor_id}" error=unknown model type')
                    return f'Processor failed to load: {processor_id}'
                self.model = cls(det_ckpt=det_ckpt, pose_config=pose_config, pose_ckpt=pose_ckpt, device="cpu")
            elif 'SegmentAnything' in processor_id:
                if 'Base' == config['SegmentAnything']['model']:
                    self.model = cls.from_pretrained(model_path = 'segments-arnaud/sam_vit_b', filename='sam_vit_b_01ec64.pth', model_type='vit_b', **self.load_config)
                elif 'Large' == config['SegmentAnything']['model']:
                    self.model = cls.from_pretrained(model_path = 'segments-arnaud/sam_vit_l', filename='sam_vit_l_0b3195.pth', model_type='vit_l', **self.load_config)
                else:
                    log.error(f'Control Processor load failed: id="{processor_id}" error=unknown model type')
                    return f'Processor failed to load: {processor_id}'
            elif config[processor_id].get('load_config', None) is not None:
                self.model = cls.from_pretrained(**self.load_config)
            elif config[processor_id]['checkpoint']:
                self.model = cls.from_pretrained("lllyasviel/Annotators", **self.load_config)
            else:
                self.model = cls() # class instance only
            t1 = time.time()
            self.processor_id = processor_id
            log.debug(f'Control Processor loaded: id="{processor_id}" class={self.model.__class__.__name__} time={t1-t0:.2f}')
            return f'Processor loaded: {processor_id}'
        except Exception as e:
            log.error(f'Control Processor load failed: id="{processor_id}" error={e}')
            display(e, 'Control Processor load')
            return f'Processor load filed: {processor_id}'

    def __call__(self, image_input: Image, mode: str = 'RGB', resize_mode: int = 0, resize_name: str = 'None', scale_tab: int = 1, scale_by: float = 1.0, local_config: dict = {}):
        if self.processor_id is None or self.processor_id == 'None':
            return self.override if self.override is not None else image_input
        if self.override is not None:
            debug(f'Control Processor: id="{self.processor_id}" override={self.override}')
            image_input = self.override
            if resize_mode != 0 and resize_name != 'None':
                if scale_tab == 1:
                    width_before, height_before = int(image_input.width * scale_by), int(image_input.height * scale_by)
                    debug(f'Control resize: op=before image={image_input} width={width_before} height={height_before} mode={resize_mode} name={resize_name}')
                    image_input = images.resize_image(resize_mode, image_input, width_before, height_before, resize_name)
        image_process = image_input
        if image_input is None:
            # log.error('Control Processor: no input')
            return image_process
        if config[self.processor_id].get('dirty', False):
            processor_id = self.processor_id
            config[processor_id].pop('dirty')
            self.reset()
            self.load(processor_id)
        if self.model is None:
            # log.error('Control Processor: model not loaded')
            return image_process
        try:
            t0 = time.time()
            kwargs = config.get(self.processor_id, {}).get('params', None)
            if kwargs:
                kwargs.update(local_config)
            if self.resize:
                image_resized = image_input.resize((512, 512), Image.Resampling.LANCZOS)
            else:
                image_resized = image_input
            with devices.inference_context():
                image_process = self.model(image_resized, **kwargs)
            if isinstance(image_process, np.ndarray):
                if np.max(image_process) < 2:
                    image_process = (255.0 * image_process).astype(np.uint8)
                image_process = Image.fromarray(image_process, 'L')
            if self.resize and image_process.size != image_input.size:
                image_process = image_process.resize(image_input.size, Image.Resampling.LANCZOS)
            t1 = time.time()
            log.debug(f'Control Processor: id="{self.processor_id}" mode={mode} args={kwargs} time={t1-t0:.2f}')
        except Exception as e:
            log.error(f'Control Processor failed: id="{self.processor_id}" error={e}')
            display(e, 'Control Processor')
        if mode != 'RGB':
            image_process = image_process.convert(mode)
        return image_process

    def preview(self):
        import modules.ui_control_helpers as helpers
        input_image = helpers.input_source
        if isinstance(input_image, list):
            input_image = input_image[0]
        debug('Control process preview')
        return self.__call__(input_image)