mirror of https://github.com/vladmandic/automatic
vladmandic
parent
2fb9e2115e
commit
9013ba7bf1
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Reference in New Issue