diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py index 51e955887..fe1edfafe 100644 --- a/modules/postprocess/yolo.py +++ b/modules/postprocess/yolo.py @@ -6,11 +6,12 @@ from copy import copy import numpy as np import gradio as gr from PIL import Image, ImageDraw +from modules.logger import log from modules import shared, processing, devices, processing_class, ui_common, ui_components, ui_symbols, images, extra_networks, sd_models from modules.detailer import Detailer -def _p_or_opt(p, key): +def get_opt(p, key): val = getattr(p, key, None) if p is not None else None if val is not None: return val @@ -75,12 +76,14 @@ class YoloRestorer(Detailer): name = os.path.splitext(os.path.basename(f))[0] if name not in files: self.list[name] = os.path.join(shared.opts.yolo_dir, f) - shared.log.info(f'Available Detailer: path="{shared.opts.yolo_dir}" items={len(list(self.list))} downloaded={downloaded}') + log.info(f'Available Detailer: path="{shared.opts.yolo_dir}" items={len(list(self.list))} downloaded={downloaded}') return list(self.list) def dependencies(self): - import installer - installer.install('ultralytics==8.3.40', ignore=True, quiet=True) + from installer import install + install("ultralytics==8.3.40", ignore=True, quiet=True) + install("omegaconf") + install("antlr4-python3-runtime") def predict( self, @@ -100,7 +103,7 @@ class YoloRestorer(Detailer): ) -> list[YoloResult]: if augment is None: - augment = _p_or_opt(p, 'detailer_augment') if _p_or_opt(p, 'detailer_augment') is not None else shared.opts.detailer_augment + augment = get_opt(p, 'detailer_augment') if get_opt(p, 'detailer_augment') is not None else shared.opts.detailer_augment if offload is None: offload = shared.opts.detailer_unload if model is None or (isinstance(model, str) and len(model) == 0): @@ -115,8 +118,8 @@ class YoloRestorer(Detailer): if model is None: return result args = { - 'conf': _p_or_opt(p, 'detailer_conf'), - 'iou': _p_or_opt(p, 'detailer_iou'), + 'conf': get_opt(p, 'detailer_conf'), + 'iou': get_opt(p, 'detailer_iou'), } try: if TYPE_CHECKING: @@ -137,7 +140,7 @@ class YoloRestorer(Detailer): if offload: model.to('cpu') except Exception as e: - shared.log.error(f'Detailer predict: {e}') + log.error(f'Detailer predict: {e}') return result classes_str = classes if classes is not None else shared.opts.detailer_classes @@ -166,9 +169,9 @@ class YoloRestorer(Detailer): box = box.tolist() w, h = box[2] - box[0], box[3] - box[1] x_size, y_size = w/image.width, h/image.height - _min = _p_or_opt(p, 'detailer_min_size') + _min = get_opt(p, 'detailer_min_size') min_size = _min if _min is not None and 0 <= _min <= 1 else 0 - _max = _p_or_opt(p, 'detailer_max_size') + _max = get_opt(p, 'detailer_max_size') max_size = _max if _max is not None and 0 < _max <= 1 else 1 if x_size >= min_size and y_size >=min_size and x_size <= max_size and y_size <= max_size: use_seg = segmentation if segmentation is not None else shared.opts.detailer_seg @@ -192,7 +195,7 @@ class YoloRestorer(Detailer): args=args, ) result.append(res) - if len(result) >= _p_or_opt(p, 'detailer_max'): + if len(result) >= get_opt(p, 'detailer_max'): break return result @@ -207,14 +210,14 @@ class YoloRestorer(Detailer): else: model_url = self.list.get(model_name, None) if model_url is None: - shared.log.error(f'Load: type=Detailer name="{model_name}" error="model not found"') + log.error(f'Load: type=Detailer name="{model_name}" error="model not found"') return None, None file_name = os.path.basename(model_url) model_file = None try: model_file = modelloader.load_file_from_url(url=model_url, model_dir=shared.opts.yolo_dir, file_name=file_name) if model_file is None: - shared.log.error(f'Load: type=Detailer name="{model_name}" url="{model_url}" error="failed to fetch model"') + log.error(f'Load: type=Detailer name="{model_name}" url="{model_url}" error="failed to fetch model"') elif model_file.endswith('.onnx'): import onnxruntime as ort options = ort.SessionOptions() @@ -227,11 +230,11 @@ class YoloRestorer(Detailer): import ultralytics model = ultralytics.YOLO(model_file) classes = list(model.names.values()) - shared.log.info(f'Load: type=Detailer name="{model_name}" model="{model_file}" ultralytics={ultralytics.__version__} classes={classes}') + log.info(f'Load: type=Detailer name="{model_name}" model="{model_file}" ultralytics={ultralytics.__version__} classes={classes}') self.models[model_name] = model return model_name, model except Exception as e: - shared.log.error(f'Load: type=Detailer name="{model_name}" error="{e}"') + log.error(f'Load: type=Detailer name="{model_name}" error="{e}"') return None, None def merge(self, items: list[YoloResult]) -> list[YoloResult]: @@ -260,7 +263,7 @@ class YoloRestorer(Detailer): size = min(image.width, image.height) // 32 font = images.get_font(size) color = (0, 190, 190) - shared.log.debug(f'Detailer: draw={items}') + log.debug(f'Detailer: draw={items}') for i, item in enumerate(items): use_seg = segmentation if segmentation is not None else shared.opts.detailer_seg if use_seg and item.mask is not None: @@ -292,7 +295,7 @@ class YoloRestorer(Detailer): shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.INPAINTING) if (sd_models.get_diffusers_task(shared.sd_model) != sd_models.DiffusersTaskType.INPAINTING) and (shared.sd_model.__class__.__name__ not in sd_models.pipe_switch_task_exclude): - shared.log.error(f'Detailer: model="{shared.sd_model.__class__.__name__}" not compatible') + log.error(f'Detailer: model="{shared.sd_model.__class__.__name__}" not compatible') return np_image models = [] @@ -305,9 +308,9 @@ class YoloRestorer(Detailer): if len(models) == 0: models = shared.opts.detailer_models if len(models) == 0: - shared.log.warning('Detailer: model=None') + log.warning('Detailer: model=None') return np_image - shared.log.debug(f'Detailer: models={models}') + log.debug(f'Detailer: models={models}') # resolve per-request detailer settings with fallback to global opts use_seg = getattr(p, 'detailer_segmentation', None) @@ -335,7 +338,7 @@ class YoloRestorer(Detailer): name, model = self.load(model_name) if model is None: - shared.log.warning(f'Detailer: model="{name}" not loaded') + log.warning(f'Detailer: model="{name}" not loaded') continue if image is None: @@ -343,12 +346,12 @@ class YoloRestorer(Detailer): items = self.predict(model, image, classes=use_classes, segmentation=use_seg, p=p) if len(items) == 0: - shared.log.info(f'Detailer: model="{name}" no items detected') + log.info(f'Detailer: model="{name}" no items detected') continue do_merge = use_merge if use_merge is not None else shared.opts.detailer_merge if do_merge and len(items) > 1: - shared.log.debug(f'Detailer: model="{name}" items={len(items)} merge') + log.debug(f'Detailer: model="{name}" items={len(items)} merge') items = self.merge(items) orig_prompt: str = orig_p.get('all_prompts', [''])[0] @@ -380,15 +383,15 @@ class YoloRestorer(Detailer): 'styles': [], 'inpaint_full_res': True, 'inpainting_mask_invert': 0, - 'mask_blur': _p_or_opt(p, 'detailer_blur'), - 'inpaint_full_res_padding': _p_or_opt(p, 'detailer_padding'), + 'mask_blur': get_opt(p, 'detailer_blur'), + 'inpaint_full_res_padding': get_opt(p, 'detailer_padding'), 'width': p.detailer_resolution, 'height': p.detailer_resolution, 'vae_type': orig_p.get('vae_type', 'Full'), } args.update(model_args) if args['denoising_strength'] == 0: - shared.log.debug(f'Detailer: model="{name}" strength=0 skip') + log.debug(f'Detailer: model="{name}" strength=0 skip') return np_image control_pipeline = None orig_class = shared.sd_model.__class__ @@ -406,7 +409,7 @@ class YoloRestorer(Detailer): p.steps = orig_p.get('steps', 0) # report = [{'label': i.label, 'score': i.score, 'size': f'{i.width}x{i.height}' } for i in items] - # shared.log.info(f'Detailer: model="{name}" items={report} args={args}') + # log.info(f'Detailer: model="{name}" items={report} args={args}') models_used.append(name) mask_all = [] @@ -414,8 +417,8 @@ class YoloRestorer(Detailer): pc = copy(p) pc.ops.append('detailer') - pc.schedulers_sigma_adjust = _p_or_opt(p, 'detailer_sigma_adjust') - pc.schedulers_sigma_adjust_max = _p_or_opt(p, 'detailer_sigma_adjust_max') + pc.schedulers_sigma_adjust = get_opt(p, 'detailer_sigma_adjust') + pc.schedulers_sigma_adjust_max = get_opt(p, 'detailer_sigma_adjust_max') pc.mask_apply_overlay = True do_sort = use_sort if use_sort is not None else shared.opts.detailer_sort @@ -435,7 +438,7 @@ class YoloRestorer(Detailer): pc.negative_prompts = [pc.negative_prompt] pc.prompts, pc.network_data = extra_networks.parse_prompts(pc.prompts) extra_networks.activate(pc, pc.network_data) - shared.log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label}" score={item.score:.2f} seg={use_seg if use_seg is not None else shared.opts.detailer_seg} prompt="{pc.prompt}"') + log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label}" score={item.score:.2f} seg={use_seg if use_seg is not None else shared.opts.detailer_seg} prompt="{pc.prompt}"') pc.init_images = [image] pc.image_mask = [item.mask] pc.overlay_images = [] @@ -470,7 +473,7 @@ class YoloRestorer(Detailer): p.state = orig_p.get('state', None) p.ops = orig_p.get('ops', []) - if len(mask_all) > 0 and _p_or_opt(p, 'include_mask'): + if len(mask_all) > 0 and get_opt(p, 'include_mask'): from modules.control.util import blend p.image_mask = blend([np.array(m) for m in mask_all]) p.image_mask = Image.fromarray(p.image_mask) @@ -510,9 +513,9 @@ class YoloRestorer(Detailer): shared.opts.detailer_seg = seg # shared.opts.detailer_resolution = resolution shared.opts.save(silent=True) - shared.log.debug(f'Detailer settings: models={detailers} classes={classes} strength={strength} conf={min_confidence} max={max_detected} iou={iou} size={min_size}-{max_size} padding={padding} steps={steps} resolution={resolution} save={save} sort={sort} seg={seg}') + log.debug(f'Detailer settings: models={detailers} classes={classes} strength={strength} conf={min_confidence} max={max_detected} iou={iou} size={min_size}-{max_size} padding={padding} steps={steps} resolution={resolution} save={save} sort={sort} seg={seg}') if not self.ui_mode: - shared.log.debug(f'Detailer expert: {text}') + log.debug(f'Detailer expert: {text}') with gr.Accordion(open=False, label="Detailer", elem_id=f"{tab}_detailer_accordion", elem_classes=["small-accordion"]): with gr.Row():