from PIL import Image from PIL import ImageDraw from copy import copy, deepcopy from pathlib import Path from modules import shared from modules import devices from modules.processing import process_images from sd_bmab import util from sd_bmab.base.common import StopGeneration from sd_bmab.base.context import Context from sd_bmab.sd_override import StableDiffusionProcessingTxt2ImgOv, StableDiffusionProcessingImg2ImgOv def apply_extensions(p, cn_enabled=False): script_runner = copy(p.scripts) script_args = deepcopy(p.script_args) active_script = ['dynamic_thresholding', 'wildcards'] if cn_enabled: active_script.append('controlnet') for idx, obj in enumerate(script_args): if 'controlnet' in obj.__class__.__name__.lower(): if hasattr(obj, 'enabled'): obj.enabled = False if hasattr(obj, 'input_mode'): obj.input_mode = getattr(obj.input_mode, 'SIMPLE', 'simple') elif isinstance(obj, dict) and 'module' in obj: obj['enabled'] = False filtered_alwayson = [] for script_object in script_runner.alwayson_scripts: filepath = script_object.filename filename = Path(filepath).stem if filename in active_script: filtered_alwayson.append(script_object) script_runner.alwayson_scripts = filtered_alwayson return script_runner, script_args def build_img2img(context: Context, img, options): p = context.sdprocessing img = img.convert('RGB') if 'inpaint_full_res' in options: res = options['inpaint_full_res'] if res == 'Whole picture': options['inpaint_full_res'] = 0 if res == 'Only masked': options['inpaint_full_res'] = 1 i2i_param = dict( init_images=[img], resize_mode=0, denoising_strength=0.4, mask=None, mask_blur=4, inpainting_fill=1, inpaint_full_res=True, inpaint_full_res_padding=32, inpainting_mask_invert=0, initial_noise_multiplier=1.0, outpath_samples=p.outpath_samples, outpath_grids=p.outpath_grids, prompt=p.prompt, negative_prompt=p.negative_prompt, styles=p.styles, seed=p.seed, subseed=p.subseed, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w, sampler_name=p.sampler_name, batch_size=1, n_iter=1, steps=p.steps, cfg_scale=p.cfg_scale, width=img.width, height=img.height, restore_faces=False, tiling=p.tiling, extra_generation_params=p.extra_generation_params, do_not_save_samples=True, do_not_save_grid=True, override_settings={ 'sd_model_checkpoint': shared.sd_model.sd_checkpoint_info.name_for_extra }, ) if hasattr(p, 'scheduler'): i2i_param['scheduler'] = p.scheduler else: if 'scheduler' in options: del options['scheduler'] context.apply_checkpoint(i2i_param) if options is not None: i2i_param.update(options) return i2i_param def process_img2img(context: Context, img, options=None): if shared.state.skipped or shared.state.interrupted: return img i2i_param = build_img2img(context, img, options) img2img = StableDiffusionProcessingImg2ImgOv(**i2i_param) img2img.cached_c = [None, None] img2img.cached_uc = [None, None] img2img.scripts, img2img.script_args = apply_extensions(context.sdprocessing) with StopGeneration(): processed = process_images(img2img) img = processed.images[0] img2img.close() devices.torch_gc() return img def process_img2img_with_controlnet(context: Context, image, options, controlnet): i2i_param = build_img2img(context, image, options) img2img = StableDiffusionProcessingImg2ImgOv(**i2i_param) img2img.cached_c = [None, None] img2img.cached_uc = [None, None] img2img.scripts, img2img.script_args = apply_extensions(context.sdprocessing, cn_enabled=True) cn_args_begin, cn_args_end = util.get_cn_args(img2img) cn_args = range(cn_args_begin, cn_args_end) sc_args = list(img2img.script_args) for ix in range(0, len(controlnet), 1): idx = cn_args[ix] sc_args[idx] = controlnet[ix] img2img.script_args = sc_args processed = process_images(img2img) image = processed.images[0] img2img.close() devices.torch_gc() return image def process_txt2img(context, options=None, controlnet=None, processor=None): p = context.sdprocessing t2i_param = dict( denoising_strength=0.4, outpath_samples=p.outpath_samples, outpath_grids=p.outpath_grids, prompt=p.prompt, negative_prompt=p.negative_prompt, styles=p.styles, seed=p.seed, subseed=p.subseed, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w, sampler_name=p.sampler_name, batch_size=1, n_iter=1, steps=p.steps, cfg_scale=p.cfg_scale, width=p.width, height=p.height, restore_faces=False, tiling=p.tiling, extra_generation_params=p.extra_generation_params, do_not_save_samples=True, do_not_save_grid=True, override_settings={ 'sd_model_checkpoint': shared.sd_model.sd_checkpoint_info.name_for_extra }, ) if hasattr(p, 'scheduler'): t2i_param['scheduler'] = p.scheduler else: if 'scheduler' in options: del options['scheduler'] context.apply_checkpoint(t2i_param) if options is not None: t2i_param.update(options) if processor: txt2img = processor(**t2i_param) else: txt2img = StableDiffusionProcessingTxt2ImgOv(**t2i_param) txt2img.context = context txt2img.cached_c = [None, None] txt2img.cached_uc = [None, None] if controlnet is None: txt2img.scripts, txt2img.script_args = apply_extensions(p, False) else: txt2img.scripts, txt2img.script_args = apply_extensions(p, True) cn_args_begin, cn_args_end = util.get_cn_args(txt2img) cn_args = range(cn_args_begin, cn_args_end) sc_args = list(txt2img.script_args) for ix in range(0, len(controlnet), 1): idx = cn_args[ix] sc_args[idx] = controlnet[ix] txt2img.script_args = sc_args with StopGeneration(): processed = process_images(txt2img) img = processed.images[0] devices.torch_gc() return img def masked_image(img, xyxy): x1, y1, x2, y2 = xyxy check = img.convert('RGBA') dd = Image.new('RGBA', img.size, (0, 0, 0, 0)) dr = ImageDraw.Draw(dd, 'RGBA') dr.rectangle((x1, y1, x2, y2), fill=(255, 0, 0, 255)) check = Image.blend(check, dd, alpha=0.5) check.convert('RGB').save('check.png')