Spaces:
Runtime error
Runtime error
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) | |