Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-08-26 06:30:44 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
e17efc19f76c30d9e02ab1ee83af388e50a49249
automatic/modules/control/processors.py
T
QualiaRain c0f5f41042 fix(control): remove inert DWPose (Legacy) pose-model selector 
The DWPose (Legacy) "Pose Model" (Tiny/Medium/Large) radio wrote to
config['DWPose (Legacy)']['model'], but the rtmlib backend
(RtmlibPoseDetector.from_pretrained) only takes the mode string and
always loads RTMPose 'lightweight' - it never reads that key (the only
config[...]['model'] readers are the SegmentAnything branch). The
selector was a no-op left over from the removed mmpose DWPose processor
(see #4956).

Remove the radio, its dead 'model' config key (both the module-level and
delay_load_config dicts), and its update_settings slot; reindex the
remaining settings 18..27. The settings list stays 1:1 with the UI
controls (re-derived the append-order -> update_settings map to confirm).

No runtime behavior change. The alternative - keep the capability by
mapping Tiny/Medium/Large onto rtmlib's mode - is intentionally not taken
here since it would change which model loads.

Co-Authored-By: Claude <noreply@anthropic.com>
2026-06-22 10:28:59 -04:00

381 lines
26 KiB
Python
Raw Blame History

import os
import time
import numpy as np
from PIL import Image
from modules.logger import log
from modules.errors import display
from modules import devices, images
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': None, 'group': 'Pose', 'checkpoint': True, 'params': {'include_body': True, 'include_hand': False, 'include_face': False}},
'RTMW': {'class': None, 'group': 'Pose', 'checkpoint': False, 'params': {'min_confidence': 0.3, 'draw_body_pose': True, 'draw_hand_pose': True, 'draw_face_pose': True}},
'RTMO': {'class': None, 'group': 'Pose', 'checkpoint': False, 'params': {'min_confidence': 0.3}},
'ViTPose': {'class': None, 'group': 'Pose', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'usyd-community/vitpose-plus-base'}, 'params': {'min_confidence': 0.3}},
# edge models
'Canny': {'class': None, 'group': 'Edge', 'checkpoint': False, 'params': {'low_threshold': 100, 'high_threshold': 200}},
'Edge': {'class': None, 'group': 'Edge', 'checkpoint': False, 'params': {'pf': True, 'mode': 'edge'}},
'LineArt Realistic': {'class': None, 'group': 'Edge', 'checkpoint': True, 'params': {'coarse': False}},
'LineArt Anime': {'class': None, 'group': 'Edge', 'checkpoint': True, 'params': {}},
'HED': {'class': None, 'group': 'Edge', 'checkpoint': True, 'params': {'scribble': False, 'safe': False}},
'PidiNet': {'class': None, 'group': 'Edge', 'checkpoint': True, 'params': {'scribble': False, 'safe': False, 'apply_filter': False}},
'MLSD': {'class': None, 'group': 'Edge', 'checkpoint': True, 'params': {'thr_v': 0.1, 'thr_d': 0.1}},
# depth models
'Midas Depth Hybrid': {'class': None, 'group': 'Depth', 'checkpoint': True, 'params': {'bg_th': 0.1, 'depth_and_normal': False}},
'Leres Depth': {'class': None, 'group': 'Depth', 'checkpoint': True, 'params': {'boost': False, 'thr_a': 0, 'thr_b': 0}},
'Zoe Depth': {'class': None, 'group': 'Depth', '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': None, 'group': 'Depth', 'checkpoint': True, 'params': {'denoising_steps': 4, 'ensemble_size': 4, 'processing_res': 768, 'match_input_res': True, 'color_map': 'None'}, 'load_config': {'pretrained_model_or_path': 'prs-eth/marigold-depth-v1-1'}},
'DPT Depth Hybrid': {'class': None, 'group': 'Depth', 'checkpoint': False, 'params': {}},
'GLPN Depth': {'class': None, 'group': 'Depth', 'checkpoint': False, 'params': {}},
'Depth Anything': {'class': None, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'LiheYoung/depth_anything_vitl14'}, 'params': {'color_map': 'inferno'}},
'Depth Pro': {'class': None, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'apple/DepthPro-hf'}, 'params': {'color_map': 'inferno'}},
'Depth Anything V2 Small': {'class': None, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'depth-anything/Depth-Anything-V2-Small-hf'}, 'params': {'color_map': 'inferno'}},
'Depth Anything V2 Large': {'class': None, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'depth-anything/Depth-Anything-V2-Large-hf'}, 'params': {'color_map': 'inferno'}},
'Marigold Depth LCM': {'class': None, 'group': 'Depth', 'checkpoint': True, 'params': {'denoising_steps': 1, 'ensemble_size': 1, 'processing_res': 768, 'match_input_res': True, 'color_map': 'None'}, 'load_config': {'pretrained_model_or_path': 'prs-eth/marigold-depth-lcm-v1-0'}},
'Lotus Depth': {'class': None, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'jingheya/lotus-depth-g-v2-1-disparity'}, 'params': {'color_map': 'inferno'}},
# normal models
'DSINE': {'class': None, 'group': 'Normal', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'hugoycj/DSINE-hub'}, 'params': {}},
'StableNormal': {'class': None, 'group': 'Normal', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'Stable-X/StableNormal'}, 'params': {}},
'Marigold Normals': {'class': None, 'group': 'Normal', 'checkpoint': True, 'params': {'denoising_steps': 4, 'ensemble_size': 4, 'processing_res': 768, 'match_input_res': True}, 'load_config': {'pretrained_model_or_path': 'prs-eth/marigold-normals-v1-1'}},
# segmentation models
'SegmentAnything 1.0': {'class': None, 'group': 'Segmentation', 'checkpoint': True, 'model': 'Base', 'params': {}},
'SegmentAnything 2.1': {'class': None, 'group': 'Segmentation', 'checkpoint': True, 'model': 'Large', 'load_config': {'pretrained_model_or_path': 'facebook/sam2.1-hiera-large'}, 'params': {}},
'OneFormer': {'class': None, 'group': 'Segmentation', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'shi-labs/oneformer_ade20k_swin_large'}, 'params': {}},
# other models
'Shuffle': {'class': None, 'group': 'Other', 'checkpoint': False, 'params': {}},
# legacy models
'MediaPipe Face (Legacy)': {'class': None, 'group': 'Pose', 'checkpoint': False, 'params': {'max_faces': 1, 'min_confidence': 0.5}},
'DWPose (Legacy)': {'class': None, 'group': 'Pose', 'checkpoint': False, 'params': {'min_confidence': 0.3}},
'TEED (Legacy)': {'class': None, 'group': 'Edge', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'fal/teed'}, 'params': {}},
'Anyline (Legacy)': {'class': None, 'group': 'Edge', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'TheMistoAI/MistoLine'}, 'params': {}},
'Normal Bae (Legacy)': {'class': None, 'group': 'Normal', 'checkpoint': True, 'params': {}},
}
def delay_load_config():
global config # pylint: disable=global-statement
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.openpose import OpenposeDetector
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
from modules.control.proc.depth_pro import DepthProDetector
from modules.control.proc.depth_anything_v2 import DepthAnythingV2Detector
from modules.control.proc.teed import TEEDDetector
from modules.control.proc.normalbae import NormalBaeDetector
from modules.control.proc.anyline import AnylineDetector
from modules.control.proc.rtmlib_pose import RtmlibPoseDetector
from modules.control.proc.vitpose import ViTPoseDetector
from modules.control.proc.sam2 import Sam2Detector
from modules.control.proc.oneformer import OneFormerDetector
from modules.control.proc.dsine import DSINEDetector
from modules.control.proc.stablenormal import StableNormalDetector
from modules.control.proc.marigold_normals import MarigoldNormalsDetector
from modules.control.proc.lotus import LotusDetector
config = {
# placeholder
'None': {},
# pose models
'OpenPose': {'class': OpenposeDetector, 'group': 'Pose', 'checkpoint': True, 'params': {'include_body': True, 'include_hand': False, 'include_face': False}},
'MediaPipe Face (Legacy)': {'class': MediapipeFaceDetector, 'group': 'Pose', 'checkpoint': False, 'params': {'max_faces': 1, 'min_confidence': 0.5}},
'DWPose (Legacy)': {'class': RtmlibPoseDetector, 'group': 'Pose', 'checkpoint': False, 'params': {'min_confidence': 0.3}},
'RTMW': {'class': RtmlibPoseDetector, 'group': 'Pose', 'checkpoint': False, 'params': {'min_confidence': 0.3, 'draw_body_pose': True, 'draw_hand_pose': True, 'draw_face_pose': True}},
'RTMO': {'class': RtmlibPoseDetector, 'group': 'Pose', 'checkpoint': False, 'params': {'min_confidence': 0.3}},
'ViTPose': {'class': ViTPoseDetector, 'group': 'Pose', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'usyd-community/vitpose-plus-base'}, 'params': {'min_confidence': 0.3}},
# edge models
'Canny': {'class': CannyDetector, 'group': 'Edge', 'checkpoint': False, 'params': {'low_threshold': 100, 'high_threshold': 200}},
'Edge': {'class': EdgeDetector, 'group': 'Edge', 'checkpoint': False, 'params': {'pf': True, 'mode': 'edge'}},
'LineArt Realistic': {'class': LineartDetector, 'group': 'Edge', 'checkpoint': True, 'params': {'coarse': False}},
'LineArt Anime': {'class': LineartAnimeDetector, 'group': 'Edge', 'checkpoint': True, 'params': {}},
'HED': {'class': HEDdetector, 'group': 'Edge', 'checkpoint': True, 'params': {'scribble': False, 'safe': False}},
'PidiNet': {'class': PidiNetDetector, 'group': 'Edge', 'checkpoint': True, 'params': {'scribble': False, 'safe': False, 'apply_filter': False}},
'MLSD': {'class': MLSDdetector, 'group': 'Edge', 'checkpoint': True, 'params': {'thr_v': 0.1, 'thr_d': 0.1}},
'TEED (Legacy)': {'class': TEEDDetector, 'group': 'Edge', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'fal/teed'}, 'params': {}},
'Anyline (Legacy)': {'class': AnylineDetector, 'group': 'Edge', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'TheMistoAI/MistoLine'}, 'params': {}},
# depth models
'Midas Depth Hybrid': {'class': MidasDetector, 'group': 'Depth', 'checkpoint': True, 'params': {'bg_th': 0.1, 'depth_and_normal': False}},
'Leres Depth': {'class': LeresDetector, 'group': 'Depth', 'checkpoint': True, 'params': {'boost': False, 'thr_a': 0, 'thr_b': 0}},
'Zoe Depth': {'class': ZoeDetector, 'group': 'Depth', '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, 'group': 'Depth', 'checkpoint': True, 'params': {'denoising_steps': 4, 'ensemble_size': 4, 'processing_res': 768, 'match_input_res': True, 'color_map': 'None'}, 'load_config': {'pretrained_model_or_path': 'prs-eth/marigold-depth-v1-1'}},
'DPT Depth Hybrid': {'class': DPTDetector, 'group': 'Depth', 'checkpoint': False, 'params': {}},
'GLPN Depth': {'class': GLPNDetector, 'group': 'Depth', 'checkpoint': False, 'params': {}},
'Depth Anything': {'class': DepthAnythingDetector, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'LiheYoung/depth_anything_vitl14'}, 'params': {'color_map': 'inferno'}},
'Depth Pro': {'class': DepthProDetector, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'apple/DepthPro-hf'}, 'params': {'color_map': 'inferno'}},
'Depth Anything V2 Small': {'class': DepthAnythingV2Detector, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'depth-anything/Depth-Anything-V2-Small-hf'}, 'params': {'color_map': 'inferno'}},
'Depth Anything V2 Large': {'class': DepthAnythingV2Detector, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'depth-anything/Depth-Anything-V2-Large-hf'}, 'params': {'color_map': 'inferno'}},
'Marigold Depth LCM': {'class': MarigoldDetector, 'group': 'Depth', 'checkpoint': True, 'params': {'denoising_steps': 1, 'ensemble_size': 1, 'processing_res': 768, 'match_input_res': True, 'color_map': 'None'}, 'load_config': {'pretrained_model_or_path': 'prs-eth/marigold-depth-lcm-v1-0'}},
'Lotus Depth': {'class': LotusDetector, 'group': 'Depth', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'jingheya/lotus-depth-g-v2-1-disparity'}, 'params': {'color_map': 'inferno'}},
# normal models
'Normal Bae (Legacy)': {'class': NormalBaeDetector, 'group': 'Normal', 'checkpoint': True, 'params': {}},
'DSINE': {'class': DSINEDetector, 'group': 'Normal', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'hugoycj/DSINE-hub'}, 'params': {}},
'StableNormal': {'class': StableNormalDetector, 'group': 'Normal', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'Stable-X/StableNormal'}, 'params': {}},
'Marigold Normals': {'class': MarigoldNormalsDetector, 'group': 'Normal', 'checkpoint': True, 'params': {'denoising_steps': 4, 'ensemble_size': 4, 'processing_res': 768, 'match_input_res': True}, 'load_config': {'pretrained_model_or_path': 'prs-eth/marigold-normals-v1-1'}},
# segmentation models
'SegmentAnything 1.0': {'class': SamDetector, 'group': 'Segmentation', 'checkpoint': True, 'model': 'Base', 'params': {}},
'SegmentAnything 2.1': {'class': Sam2Detector, 'group': 'Segmentation', 'checkpoint': True, 'model': 'Large', 'load_config': {'pretrained_model_or_path': 'facebook/sam2.1-hiera-large'}, 'params': {}},
'OneFormer': {'class': OneFormerDetector, 'group': 'Segmentation', 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'shi-labs/oneformer_ade20k_swin_large'}, 'params': {}},
# other models
'Shuffle': {'class': ContentShuffleDetector, 'group': 'Other', 'checkpoint': False, 'params': {}},
}
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 (Legacy)', 'params', 'max_faces'], settings[14])
update(['MediaPipe Face (Legacy)', 'params', 'min_confidence'], settings[15])
update(['Canny', 'params', 'low_threshold'], settings[16])
update(['Canny', 'params', 'high_threshold'], settings[17])
update(['DWPose (Legacy)', 'params', 'min_confidence'], settings[18])
update(['SegmentAnything 1.0', 'model'], settings[19])
update(['Edge', 'params', 'pf'], settings[20])
update(['Edge', 'params', 'mode'], settings[21])
update(['Zoe Depth', 'params', 'gamma_corrected'], settings[22])
update(['Marigold Depth', 'params', 'color_map'], settings[23])
update(['Marigold Depth', 'params', 'denoising_steps'], settings[24])
update(['Marigold Depth', 'params', 'ensemble_size'], settings[25])
update(['Depth Anything', 'params', 'color_map'], settings[26])
update(['Depth Pro', 'params', 'color_map'], settings[27])
class Processor:
def __init__(self, processor_id: str | None = None, resize = True):
self.model = None
self.processor_id: str | None = None
self.override: Image.Image | None = None
self.resize = resize
self.reset()
self.config(processor_id)
if processor_id is not None:
self.load()
def __str__(self):
return f' Processor(id={self.processor_id} model={self.model.__class__.__name__})' if self.processor_id and self.model else ''
def reset(self, processor_id: str | None = None):
if self.model is not None:
debug(f'Control Processor unloaded: id="{self.processor_id}"')
self.model = None
self.processor_id = processor_id
devices.torch_gc(force=True, reason='processor')
self.load_config = { 'cache_dir': cache_dir }
from modules.shared import opts
if opts.offline_mode:
self.load_config["local_files_only"] = True
os.environ['HF_HUB_OFFLINE'] = '1'
else:
os.environ.pop('HF_HUB_OFFLINE', None)
os.unsetenv('HF_HUB_OFFLINE')
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 = None, force: bool = True) -> str:
from modules.shared import state
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)
else:
if not force and self.model is not None:
# log.debug(f'Control Processor: id={processor_id} already loaded')
return ''
if processor_id not in config:
log.error(f'Control Processor unknown: id="{processor_id}" available={list(config)}')
return f'Processor failed to load: {processor_id}'
cls = config[processor_id]['class']
if cls is None:
delay_load_config()
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}')
jobid = state.begin('Load processor')
if processor_id in ('DWPose (Legacy)', 'RTMW', 'RTMO'):
model_type = {'DWPose (Legacy)': 'DWPose', 'RTMW': 'RTMW-l', 'RTMO': 'RTMO-l'}[processor_id]
self.model = cls.from_pretrained(model_type, **self.load_config)
elif processor_id == 'SegmentAnything 1.0':
if 'Base' == config[processor_id]['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[processor_id]['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()
state.end(jobid)
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.Image, mode: str = 'RGB', width: int = 0, height: int = 0, resize_mode: int = 0, resize_name: str = 'None', scale_tab: int = 1, scale_by: float = 1.0, local_config: dict | None = None):
"""Run the preprocessor on an input image and return the processed control map.
Args:
image_input: Source image to preprocess.
mode: Output color mode ('RGB', 'L', etc.).
width, height: Target dimensions for resize when an override image is provided.
resize_mode: Resize strategy index (0 = no resize).
resize_name: Resize algorithm name ('None' to skip).
scale_tab: Scale mode selector (1 = scale by multiplier).
scale_by: Scale multiplier when scale_tab is 1.
local_config: Per-call parameter overrides merged on top of the processor's global
config[processor_id]['params']. Keys must match the processor's accepted kwargs
(e.g. {'low_threshold': 50, 'high_threshold': 150} for Canny). Passed from
the API via Unit.process_params.
"""
if local_config is None:
local_config = {}
if self.override is not None:
debug(f'Control Processor: id="{self.processor_id}" override={self.override}')
width = image_input.width if image_input is not None else width
height = image_input.height if image_input is not None else height
if (width != self.override.width) or (height != self.override.height):
debug(f'Control resize: op=override image={self.override} width={width} height={height} mode={resize_mode} name={resize_name}')
image_input = images.resize_image(resize_mode, self.override, width, height, resize_name)
else:
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)
if self.processor_id is None or self.processor_id == 'None':
return image_input
image_process = image_input
if image_input is None:
# log.error('Control Processor: no input')
return image_process
if isinstance(image_input, list):
image_input = image_input[0]
if self.processor_id not in config:
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 = dict(config.get(self.processor_id, {}).get('params', {}))
if local_config:
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 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)
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 input_image is None:
return []
if isinstance(input_image, list):
input_image = input_image[0]
debug('Control process preview')
return [self.__call__(input_image)]
Reference in New Issue View Git Blame Copy Permalink