1
0
mirror of https://github.com/vladmandic/sdnext.git synced 2026-01-29 05:02:09 +03:00
Files
sdnext/modules/control/run.py
Vladimir Mandic 7aa09957dc fix lllite
Signed-off-by: Vladimir Mandic <mandic00@live.com>
2025-05-04 13:55:45 -04:00

841 lines
51 KiB
Python

import os
import time
from typing import List, Union
import cv2
import numpy as np
from PIL import Image
from modules.control import util # helper functions
from modules.control import unit # control units
from modules.control import processors # image preprocessors
from modules.control import tile # tiling module
from modules.control.units import controlnet # lllyasviel ControlNet
from modules.control.units import xs # VisLearn ControlNet-XS
from modules.control.units import lite # Kohya ControlLLLite
from modules.control.units import t2iadapter # TencentARC T2I-Adapter
from modules.control.units import reference # ControlNet-Reference
from modules import devices, shared, errors, processing, images, sd_models, scripts, masking
from modules.processing_class import StableDiffusionProcessingControl
from modules.ui_common import infotext_to_html
from modules.api import script
from modules.timer import process as process_timer
debug = os.environ.get('SD_CONTROL_DEBUG', None) is not None
debug_log = shared.log.trace if debug else lambda *args, **kwargs: None
pipe = None
instance = None
original_pipeline = None
p_extra_args = {}
unified_models = ['Flex2Pipeline'] # models that have controlnet builtin
def restore_pipeline():
global pipe, instance # pylint: disable=global-statement
if instance is not None and hasattr(instance, 'restore'):
instance.restore()
if original_pipeline is not None and (original_pipeline.__class__.__name__ != shared.sd_model.__class__.__name__):
debug_log(f'Control restored pipeline: class={shared.sd_model.__class__.__name__} to={original_pipeline.__class__.__name__}')
shared.sd_model = original_pipeline
pipe = None
instance = None
devices.torch_gc()
def terminate(msg):
restore_pipeline()
shared.log.error(f'Control terminated: {msg}')
return msg
def is_unified_model():
return shared.sd_model.__class__.__name__ in unified_models
def set_pipe(p, has_models, unit_type, selected_models, active_model, active_strength, control_conditioning, control_guidance_start, control_guidance_end, inits):
global pipe, instance # pylint: disable=global-statement
pipe = None
if has_models:
p.ops.append('control')
p.extra_generation_params["Control type"] = unit_type # overriden later with pretty-print
p.extra_generation_params["Control model"] = ';'.join([(m.model_id or '') for m in active_model if m.model is not None])
p.extra_generation_params["Control conditioning"] = control_conditioning if isinstance(control_conditioning, list) else [control_conditioning]
p.extra_generation_params['Control start'] = control_guidance_start if isinstance(control_guidance_start, list) else [control_guidance_start]
p.extra_generation_params['Control end'] = control_guidance_end if isinstance(control_guidance_end, list) else [control_guidance_end]
p.extra_generation_params["Control conditioning"] = ';'.join([str(c) for c in p.extra_generation_params["Control conditioning"]])
p.extra_generation_params['Control start'] = ';'.join([str(c) for c in p.extra_generation_params['Control start']])
p.extra_generation_params['Control end'] = ';'.join([str(c) for c in p.extra_generation_params['Control end']])
if unit_type == 't2i adapter' and has_models:
p.extra_generation_params["Control type"] = 'T2I-Adapter'
p.task_args['adapter_conditioning_scale'] = control_conditioning
instance = t2iadapter.AdapterPipeline(selected_models, shared.sd_model)
pipe = instance.pipeline
if inits is not None:
shared.log.warning('Control: T2I-Adapter does not support separate init image')
elif unit_type == 'controlnet' and has_models:
p.extra_generation_params["Control type"] = 'ControlNet'
if shared.sd_model_type == 'f1':
p.task_args['controlnet_conditioning_scale'] = control_conditioning if isinstance(control_conditioning, list) else [control_conditioning]
else:
p.task_args['controlnet_conditioning_scale'] = control_conditioning
p.task_args['control_guidance_start'] = control_guidance_start
p.task_args['control_guidance_end'] = control_guidance_end
p.task_args['guess_mode'] = p.guess_mode
if not is_unified_model():
instance = controlnet.ControlNetPipeline(selected_models, shared.sd_model, p=p)
pipe = instance.pipeline
else:
pipe = shared.sd_model
elif unit_type == 'xs' and has_models:
p.extra_generation_params["Control type"] = 'ControlNet-XS'
p.controlnet_conditioning_scale = control_conditioning
p.control_guidance_start = control_guidance_start
p.control_guidance_end = control_guidance_end
instance = xs.ControlNetXSPipeline(selected_models, shared.sd_model)
pipe = instance.pipeline
if inits is not None:
shared.log.warning('Control: ControlNet-XS does not support separate init image')
elif unit_type == 'lite' and has_models:
p.extra_generation_params["Control type"] = 'ControlLLLite'
p.controlnet_conditioning_scale = control_conditioning
instance = lite.ControlLLitePipeline(shared.sd_model)
pipe = instance.pipeline
if inits is not None:
shared.log.warning('Control: ControlLLLite does not support separate init image')
elif unit_type == 'reference' and has_models:
p.extra_generation_params["Control type"] = 'Reference'
p.extra_generation_params["Control attention"] = p.attention
p.task_args['reference_attn'] = 'Attention' in p.attention
p.task_args['reference_adain'] = 'Adain' in p.attention
p.task_args['attention_auto_machine_weight'] = p.query_weight
p.task_args['gn_auto_machine_weight'] = p.adain_weight
p.task_args['style_fidelity'] = p.fidelity
instance = reference.ReferencePipeline(shared.sd_model)
pipe = instance.pipeline
if inits is not None:
shared.log.warning('Control: ControlNet-XS does not support separate init image')
else: # run in txt2img/img2img mode
if len(active_strength) > 0:
p.strength = active_strength[0]
pipe = shared.sd_model
instance = None
debug_log(f'Control: run type={unit_type} models={has_models} pipe={pipe.__class__.__name__ if pipe is not None else None}')
return pipe
def check_active(p, unit_type, units):
active_process: List[processors.Processor] = [] # all active preprocessors
active_model: List[Union[controlnet.ControlNet, xs.ControlNetXS, t2iadapter.Adapter]] = [] # all active models
active_strength: List[float] = [] # strength factors for all active models
active_start: List[float] = [] # start step for all active models
active_end: List[float] = [] # end step for all active models
num_units = 0
for u in units:
if u.type != unit_type:
continue
num_units += 1
debug_log(f'Control unit: i={num_units} type={u.type} enabled={u.enabled} cn={u.controlnet} proc={u.process}')
if not u.enabled:
if u.controlnet is not None and u.controlnet.model is not None:
debug_log(f'Control unit offload: model="{u.controlnet.model_id}" device={devices.cpu}')
sd_models.move_model(u.controlnet.model, devices.cpu)
continue
if u.controlnet is not None and u.controlnet.model is not None:
debug_log(f'Control unit offload: model="{u.controlnet.model_id}" device={devices.device}')
sd_models.move_model(u.controlnet.model, devices.device)
if unit_type == 't2i adapter' and u.adapter.model is not None:
active_process.append(u.process)
active_model.append(u.adapter)
active_strength.append(float(u.strength))
p.adapter_conditioning_factor = u.factor
shared.log.debug(f'Control T2I-Adapter unit: i={num_units} process="{u.process.processor_id}" model="{u.adapter.model_id}" strength={u.strength} factor={u.factor}')
elif unit_type == 'controlnet' and (u.controlnet.model is not None or is_unified_model()):
active_process.append(u.process)
active_model.append(u.controlnet)
active_strength.append(float(u.strength))
active_start.append(float(u.start))
active_end.append(float(u.end))
p.guess_mode = u.guess
if isinstance(u.mode, str):
if not hasattr(p, 'control_mode'):
p.control_mode = []
p.control_mode.append(u.choices.index(u.mode) if u.mode in u.choices else 0)
p.is_tile = p.is_tile or 'tile' in u.mode.lower()
p.control_tile = u.tile
p.extra_generation_params["Control mode"] = u.mode
shared.log.debug(f'Control ControlNet unit: i={num_units} process="{u.process.processor_id}" model="{u.controlnet.model_id}" strength={u.strength} guess={u.guess} start={u.start} end={u.end} mode={u.mode}')
elif unit_type == 'xs' and u.controlnet.model is not None:
active_process.append(u.process)
active_model.append(u.controlnet)
active_strength.append(float(u.strength))
active_start.append(float(u.start))
active_end.append(float(u.end))
shared.log.debug(f'Control ControlNet-XS unit: i={num_units} process={u.process.processor_id} model={u.controlnet.model_id} strength={u.strength} guess={u.guess} start={u.start} end={u.end}')
elif unit_type == 'lite' and u.controlnet.model is not None:
active_process.append(u.process)
active_model.append(u.controlnet)
active_strength.append(float(u.strength))
shared.log.debug(f'Control ControlLLite unit: i={num_units} process={u.process.processor_id} model={u.controlnet.model_id} strength={u.strength} guess={u.guess} start={u.start} end={u.end}')
elif unit_type == 'reference':
p.override = u.override
p.attention = u.attention
p.query_weight = float(u.query_weight)
p.adain_weight = float(u.adain_weight)
p.fidelity = u.fidelity
shared.log.debug('Control Reference unit')
else:
if u.process.processor_id is not None:
active_process.append(u.process)
shared.log.debug(f'Control process unit: i={num_units} process={u.process.processor_id}')
active_strength.append(float(u.strength))
debug_log(f'Control active: process={len(active_process)} model={len(active_model)}')
return active_process, active_model, active_strength, active_start, active_end
def check_enabled(p, unit_type, units, active_model, active_strength, active_start, active_end):
has_models = False
selected_models: List[Union[controlnet.ControlNetModel, xs.ControlNetXSModel, t2iadapter.AdapterModel]] = None
control_conditioning = None
control_guidance_start = None
control_guidance_end = None
if unit_type == 't2i adapter' or unit_type == 'controlnet' or unit_type == 'xs' or unit_type == 'lite':
if len(active_model) == 0:
selected_models = None
elif len(active_model) == 1:
selected_models = active_model[0].model if active_model[0].model is not None else None
p.is_tile = p.is_tile or 'tile' in (active_model[0].model_id or '').lower()
has_models = (selected_models is not None) or is_unified_model()
control_conditioning = active_strength[0] if len(active_strength) > 0 else 1 # strength or list[strength]
control_guidance_start = active_start[0] if len(active_start) > 0 else 0
control_guidance_end = active_end[0] if len(active_end) > 0 else 1
else:
selected_models = [m.model for m in active_model if m.model is not None]
has_models = len(selected_models) > 0
control_conditioning = active_strength[0] if len(active_strength) == 1 else list(active_strength) # strength or list[strength]
control_guidance_start = active_start[0] if len(active_start) == 1 else list(active_start)
control_guidance_end = active_end[0] if len(active_end) == 1 else list(active_end)
elif unit_type == 'reference':
has_models = any(u.enabled for u in units if u.type == 'reference')
else:
pass
return has_models, selected_models, control_conditioning, control_guidance_start, control_guidance_end
def control_set(kwargs):
if kwargs:
global p_extra_args # pylint: disable=global-statement
p_extra_args = {}
debug_log(f'Control extra args: {kwargs}')
for k, v in kwargs.items():
p_extra_args[k] = v
def control_run(state: str = '',
units: List[unit.Unit] = [], inputs: List[Image.Image] = [], inits: List[Image.Image] = [], mask: Image.Image = None, unit_type: str = None, is_generator: bool = True,
input_type: int = 0,
prompt: str = '', negative_prompt: str = '', styles: List[str] = [],
steps: int = 20, sampler_index: int = None,
seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1,
cfg_scale: float = 6.0, clip_skip: float = 1.0, image_cfg_scale: float = 6.0, diffusers_guidance_rescale: float = 0.7, pag_scale: float = 0.0, pag_adaptive: float = 0.5, cfg_end: float = 1.0,
vae_type: str = 'Full', tiling: bool = False, hidiffusion: bool = False,
detailer_enabled: bool = True, detailer_prompt: str = '', detailer_negative: str = '', detailer_steps: int = 10, detailer_strength: float = 0.3,
hdr_mode: int = 0, hdr_brightness: float = 0, hdr_color: float = 0, hdr_sharpen: float = 0, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 0.95,
hdr_maximize: bool = False, hdr_max_center: float = 0.6, hdr_max_boundry: float = 1.0, hdr_color_picker: str = None, hdr_tint_ratio: float = 0,
resize_mode_before: int = 0, resize_name_before: str = 'None', resize_context_before: str = 'None', width_before: int = 512, height_before: int = 512, scale_by_before: float = 1.0, selected_scale_tab_before: int = 0,
resize_mode_after: int = 0, resize_name_after: str = 'None', resize_context_after: str = 'None', width_after: int = 0, height_after: int = 0, scale_by_after: float = 1.0, selected_scale_tab_after: int = 0,
resize_mode_mask: int = 0, resize_name_mask: str = 'None', resize_context_mask: str = 'None', width_mask: int = 0, height_mask: int = 0, scale_by_mask: float = 1.0, selected_scale_tab_mask: int = 0,
denoising_strength: float = 0.3, batch_count: int = 1, batch_size: int = 1,
enable_hr: bool = False, hr_sampler_index: int = None, hr_denoising_strength: float = 0.0, hr_resize_mode: int = 0, hr_resize_context: str = 'None', hr_upscaler: str = None, hr_force: bool = False, hr_second_pass_steps: int = 20,
hr_scale: float = 1.0, hr_resize_x: int = 0, hr_resize_y: int = 0, refiner_steps: int = 5, refiner_start: float = 0.0, refiner_prompt: str = '', refiner_negative: str = '',
video_skip_frames: int = 0, video_type: str = 'None', video_duration: float = 2.0, video_loop: bool = False, video_pad: int = 0, video_interpolate: int = 0,
*input_script_args,
):
# handle optional initialization via ui
for u in units:
if not u.enabled:
continue
if u.process_name is not None and u.process_name != '' and u.process_name != 'None':
u.process.load(u.process_name, force=False)
if u.model_name is not None and u.model_name != '' and u.model_name != 'None':
if u.type == 't2i adapter':
u.adapter.load(u.model_name, force=False)
else:
u.controlnet.load(u.model_name, force=False)
u.update_choices(u.model_name)
if u.process is not None and u.process.override is None and u.override is not None:
u.process.override = u.override
global pipe, original_pipeline # pylint: disable=global-statement
debug_log(f'Control: type={unit_type} input={inputs} init={inits} type={input_type}')
if inputs is None or (type(inputs) is list and len(inputs) == 0):
inputs = [None]
output_images: List[Image.Image] = [] # output images
processed_image: Image.Image = None # last processed image
if mask is not None and input_type == 0:
input_type = 1 # inpaint always requires control_image
if sampler_index is None:
shared.log.warning('Sampler: invalid')
sampler_index = 0
if hr_sampler_index is None:
hr_sampler_index = sampler_index
p = StableDiffusionProcessingControl(
prompt = prompt,
negative_prompt = negative_prompt,
styles = styles,
steps = steps,
n_iter = batch_count,
batch_size = batch_size,
sampler_name = processing.get_sampler_name(sampler_index),
seed = seed,
subseed = subseed,
subseed_strength = subseed_strength,
seed_resize_from_h = seed_resize_from_h,
seed_resize_from_w = seed_resize_from_w,
# advanced
cfg_scale = cfg_scale,
cfg_end = cfg_end,
clip_skip = clip_skip,
image_cfg_scale = image_cfg_scale,
diffusers_guidance_rescale = diffusers_guidance_rescale,
pag_scale = pag_scale,
pag_adaptive = pag_adaptive,
vae_type = vae_type,
tiling = tiling,
hidiffusion = hidiffusion,
# detailer
detailer_enabled = detailer_enabled,
detailer_prompt = detailer_prompt,
detailer_negative = detailer_negative,
detailer_steps = detailer_steps,
detailer_strength = detailer_strength,
# resize
resize_mode = resize_mode_before if resize_name_before != 'None' else 0,
resize_name = resize_name_before,
scale_by = scale_by_before,
selected_scale_tab = selected_scale_tab_before,
denoising_strength = denoising_strength,
# inpaint
inpaint_full_res = masking.opts.mask_only,
inpainting_mask_invert = 1 if masking.opts.invert else 0,
inpainting_fill = 1,
# hdr
hdr_mode=hdr_mode, hdr_brightness=hdr_brightness, hdr_color=hdr_color, hdr_sharpen=hdr_sharpen, hdr_clamp=hdr_clamp,
hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold, hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry, hdr_color_picker=hdr_color_picker, hdr_tint_ratio=hdr_tint_ratio,
# path
outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_control_samples,
outpath_grids=shared.opts.outdir_grids or shared.opts.outdir_control_grids,
)
p.state = state
p.is_tile = False
# processing.process_init(p)
resize_mode_before = resize_mode_before if resize_name_before != 'None' and inputs is not None and len(inputs) > 0 else 0
# TODO modernui: monkey-patch for missing tabs.select event
if selected_scale_tab_before == 0 and resize_name_before != 'None' and scale_by_before != 1 and inputs is not None and len(inputs) > 0:
shared.log.debug('Control: override resize mode=before')
selected_scale_tab_before = 1
if selected_scale_tab_after == 0 and resize_name_after != 'None' and scale_by_after != 1:
shared.log.debug('Control: override resize mode=after')
selected_scale_tab_after = 1
if selected_scale_tab_mask == 0 and resize_name_mask != 'None' and scale_by_mask != 1:
shared.log.debug('Control: override resize mode=mask')
selected_scale_tab_mask = 1
# set control sizing
if resize_mode_before != 0 or inputs is None or inputs == [None]:
p.width, p.height = width_before, height_before # pylint: disable=attribute-defined-outside-init
p.width_before = width_before
p.height_before = height_before
if resize_name_before != 'None':
p.resize_mode_before = resize_mode_before
p.resize_name_before = resize_name_before
p.scale_by_before = scale_by_before
p.selected_scale_tab_before = selected_scale_tab_before
else:
del p.width
del p.height
if resize_name_after != 'None':
p.resize_mode_after = resize_mode_after
p.resize_name_after = resize_name_after
p.width_after = width_after
p.height_after = height_after
p.scale_by_after = scale_by_after
p.selected_scale_tab_after = selected_scale_tab_after
if resize_name_mask != 'None':
p.resize_mode_mask = resize_mode_mask
p.resize_name_mask = resize_name_mask
p.width_mask = width_mask
p.height_mask = height_mask
p.scale_by_mask = scale_by_mask
p.selected_scale_tab_mask = selected_scale_tab_mask
# hires/refine defined outside of main init
p.enable_hr = enable_hr
p.hr_sampler_name = processing.get_sampler_name(hr_sampler_index)
p.hr_denoising_strength = hr_denoising_strength
p.hr_resize_mode = hr_resize_mode
p.hr_resize_context = hr_resize_context
p.hr_upscaler = hr_upscaler
p.hr_force = hr_force
p.hr_second_pass_steps = hr_second_pass_steps
p.hr_scale = hr_scale
p.hr_resize_x = hr_resize_x
p.hr_resize_y = hr_resize_y
p.refiner_steps = refiner_steps
p.refiner_start = refiner_start
p.refiner_prompt = refiner_prompt
p.refiner_negative = refiner_negative
if p.enable_hr and (p.hr_resize_x == 0 or p.hr_resize_y == 0):
p.hr_upscale_to_x, p.hr_upscale_to_y = 8 * int(width_before * p.hr_scale / 8), 8 * int(height_before * p.hr_scale / 8)
elif p.enable_hr and (p.hr_upscale_to_x == 0 or p.hr_upscale_to_y == 0):
p.hr_upscale_to_x, p.hr_upscale_to_y = 8 * int(p.hr_resize_x / 8), 8 * int(hr_resize_y / 8)
if is_unified_model():
p.init_images = inputs
global p_extra_args # pylint: disable=global-statement
for k, v in p_extra_args.items():
setattr(p, k, v)
p_extra_args = {}
if shared.sd_model is None:
shared.log.warning('Aborted: op=control model not loaded')
return [], '', '', 'Error: model not loaded'
unit_type = unit_type.strip().lower() if unit_type is not None else ''
t0 = time.time()
active_process, active_model, active_strength, active_start, active_end = check_active(p, unit_type, units)
has_models, selected_models, control_conditioning, control_guidance_start, control_guidance_end = check_enabled(p, unit_type, units, active_model, active_strength, active_start, active_end)
processed: processing.Processed = None
image_txt = ''
info_txt = []
p.is_tile = p.is_tile and has_models
pipe = set_pipe(p, has_models, unit_type, selected_models, active_model, active_strength, control_conditioning, control_guidance_start, control_guidance_end, inits)
debug_log(f'Control pipeline: class={pipe.__class__.__name__} args={vars(p)}')
t1, t2, t3 = time.time(), 0, 0
status = True
frame = None
video = None
output_filename = None
index = 0
frames = 0
blended_image = None
# set pipeline
if pipe is None:
return [], '', '', 'Pipeline not set'
elif pipe.__class__.__name__ != shared.sd_model.__class__.__name__:
original_pipeline = shared.sd_model
shared.sd_model = pipe
sd_models.move_model(shared.sd_model, shared.device)
shared.sd_model.to(dtype=devices.dtype)
debug_log(f'Control device={devices.device} dtype={devices.dtype}')
sd_models.copy_diffuser_options(shared.sd_model, original_pipeline) # copy options from original pipeline
sd_models.set_diffuser_options(shared.sd_model)
else:
original_pipeline = None
try:
with devices.inference_context():
if isinstance(inputs, str): # only video, the rest is a list
if input_type == 2: # separate init image
if isinstance(inits, str) and inits != inputs:
shared.log.warning('Control: separate init video not support for video input')
input_type = 1
try:
video = cv2.VideoCapture(inputs)
if not video.isOpened():
if is_generator:
yield terminate(f'Video open failed: path={inputs}')
return [], '', '', 'Error: video open failed'
frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
fps = int(video.get(cv2.CAP_PROP_FPS))
w, h = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
codec = util.decode_fourcc(video.get(cv2.CAP_PROP_FOURCC))
status, frame = video.read()
if status:
shared.state.frame_count = 1 + frames // (video_skip_frames + 1)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
shared.log.debug(f'Control: input video: path={inputs} frames={frames} fps={fps} size={w}x{h} codec={codec}')
except Exception as e:
if is_generator:
yield terminate(f'Video open failed: path={inputs} {e}')
return [], '', '', 'Error: video open failed'
while status:
if pipe is None: # pipe may have been reset externally
pipe = set_pipe(p, has_models, unit_type, selected_models, active_model, active_strength, control_conditioning, control_guidance_start, control_guidance_end, inits)
debug_log(f'Control pipeline reinit: class={pipe.__class__.__name__}')
possible = sd_models.get_call(pipe).keys()
processed_image = None
if frame is not None:
inputs = [Image.fromarray(frame)] # cv2 to pil
for i, input_image in enumerate(inputs):
debug_log(f'Control Control image: {i + 1} of {len(inputs)}')
if shared.state.skipped:
shared.state.skipped = False
continue
if shared.state.interrupted:
shared.state.interrupted = False
if is_generator:
yield terminate('Interrupted')
return [], '', '', 'Interrupted'
# get input
if isinstance(input_image, str):
try:
input_image = Image.open(inputs[i])
except Exception as e:
shared.log.error(f'Control: image open failed: path={inputs[i]} type=control error={e}')
continue
# match init input
if input_type == 1:
debug_log('Control Init image: same as control')
init_image = input_image
elif inits is None:
debug_log('Control Init image: none')
init_image = None
elif isinstance(inits[i], str):
debug_log(f'Control: init image: {inits[i]}')
try:
init_image = Image.open(inits[i])
except Exception as e:
shared.log.error(f'Control: image open failed: path={inits[i]} type=init error={e}')
continue
else:
debug_log(f'Control Init image: {i % len(inits) + 1} of {len(inits)}')
init_image = inits[i % len(inits)]
if video is not None and index % (video_skip_frames + 1) != 0:
index += 1
continue
index += 1
# resize before
if resize_mode_before != 0 and resize_name_before != 'None':
if selected_scale_tab_before == 1 and input_image is not None:
width_before, height_before = int(input_image.width * scale_by_before), int(input_image.height * scale_by_before)
if input_image is not None:
p.extra_generation_params["Control resize"] = f'{resize_name_before}'
debug_log(f'Control resize: op=before image={input_image} width={width_before} height={height_before} mode={resize_mode_before} name={resize_name_before} context="{resize_context_before}"')
input_image = images.resize_image(resize_mode_before, input_image, width_before, height_before, resize_name_before, context=resize_context_before)
if input_image is not None and init_image is not None and init_image.size != input_image.size:
debug_log(f'Control resize init: image={init_image} target={input_image}')
init_image = images.resize_image(resize_mode=1, im=init_image, width=input_image.width, height=input_image.height)
if input_image is not None and p.override is not None and p.override.size != input_image.size:
debug_log(f'Control resize override: image={p.override} target={input_image}')
p.override = images.resize_image(resize_mode=1, im=p.override, width=input_image.width, height=input_image.height)
if input_image is not None:
p.width = input_image.width
p.height = input_image.height
debug_log(f'Control: input image={input_image}')
processed_images = []
if mask is not None:
p.extra_generation_params["Mask only"] = masking.opts.mask_only if masking.opts.mask_only else None
p.extra_generation_params["Mask auto"] = masking.opts.auto_mask if masking.opts.auto_mask != 'None' else None
p.extra_generation_params["Mask invert"] = masking.opts.invert if masking.opts.invert else None
p.extra_generation_params["Mask blur"] = masking.opts.mask_blur if masking.opts.mask_blur > 0 else None
p.extra_generation_params["Mask erode"] = masking.opts.mask_erode if masking.opts.mask_erode > 0 else None
p.extra_generation_params["Mask dilate"] = masking.opts.mask_dilate if masking.opts.mask_dilate > 0 else None
p.extra_generation_params["Mask model"] = masking.opts.model if masking.opts.model is not None else None
masked_image = masking.run_mask(input_image=input_image, input_mask=mask, return_type='Masked', invert=p.inpainting_mask_invert==1) if mask is not None else input_image
else:
masked_image = input_image
for i, process in enumerate(active_process): # list[image]
debug_log(f'Control: i={i+1} process="{process.processor_id}" input={masked_image} override={process.override}')
processed_image = process(
image_input=masked_image,
mode='RGB',
resize_mode=resize_mode_before,
resize_name=resize_name_before,
scale_tab=selected_scale_tab_before,
scale_by=scale_by_before,
)
if processed_image is not None:
processed_images.append(processed_image)
if shared.opts.control_unload_processor and process.processor_id is not None:
processors.config[process.processor_id]['dirty'] = True # to force reload
process.model = None
debug_log(f'Control processed: {len(processed_images)}')
if len(processed_images) > 0:
try:
if len(p.extra_generation_params["Control process"]) == 0:
p.extra_generation_params["Control process"] = None
else:
p.extra_generation_params["Control process"] = ';'.join([p.processor_id for p in active_process if p.processor_id is not None])
except Exception:
pass
if any(img is None for img in processed_images):
if is_generator:
yield terminate('Attempting process but output is none')
return [], '', '', 'Error: output is none'
if len(processed_images) > 1 and len(active_process) != len(active_model):
processed_image = [np.array(i) for i in processed_images]
processed_image = util.blend(processed_image) # blend all processed images into one
processed_image = Image.fromarray(processed_image)
blended_image = processed_image
elif len(processed_images) == 1:
processed_image = processed_images
blended_image = processed_image[0]
else:
blended_image = [np.array(i) for i in processed_images]
blended_image = util.blend(blended_image) # blend all processed images into one
blended_image = Image.fromarray(blended_image)
if isinstance(selected_models, list) and len(processed_images) == len(selected_models):
debug_log(f'Control: inputs match: input={len(processed_images)} models={len(selected_models)}')
p.init_images = processed_images
elif isinstance(selected_models, list) and len(processed_images) != len(selected_models):
if is_generator:
yield terminate(f'Number of inputs does not match: input={len(processed_images)} models={len(selected_models)}')
return [], '', '', 'Error: number of inputs does not match'
elif selected_models is not None:
p.init_images = processed_image
else:
debug_log('Control processed: using input direct')
processed_image = input_image
if unit_type == 'reference' and has_models:
p.ref_image = p.override or input_image
p.task_args.pop('image', None)
p.task_args['ref_image'] = p.ref_image
debug_log(f'Control: process=None image={p.ref_image}')
if p.ref_image is None:
if is_generator:
yield terminate('Attempting reference mode but image is none')
return [], '', '', 'Reference mode without image'
elif unit_type == 'controlnet' and has_models:
if input_type == 0: # Control only
if 'control_image' in possible:
p.task_args['control_image'] = [p.init_images] if isinstance(p.init_images, Image.Image) else p.init_images
elif 'image' in possible:
p.task_args['image'] = [p.init_images] if isinstance(p.init_images, Image.Image) else p.init_images
if 'control_mode' in possible:
p.task_args['control_mode'] = getattr(p, 'control_mode', None)
if 'strength' in possible:
p.task_args['strength'] = p.denoising_strength
p.init_images = None
elif input_type == 1: # Init image same as control
p.init_images = [p.override or input_image] * max(1, len(active_model))
if 'inpaint_image' in possible: # flex
p.task_args['inpaint_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images
p.task_args['inpaint_mask'] = Image.new('L', p.task_args['inpaint_image'].size, int(p.denoising_strength * 255))
p.task_args['control_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images
p.task_args['width'] = p.width
p.task_args['height'] = p.height
elif 'control_image' in possible:
p.task_args['control_image'] = p.init_images # switch image and control_image
if 'control_mode' in possible:
p.task_args['control_mode'] = getattr(p, 'control_mode', None)
if 'strength' in possible:
p.task_args['strength'] = p.denoising_strength
elif input_type == 2: # Separate init image
if init_image is None:
shared.log.warning('Control: separate init image not provided')
init_image = input_image
if 'inpaint_image' in possible: # flex
p.task_args['inpaint_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images
p.task_args['inpaint_mask'] = Image.new('L', p.task_args['inpaint_image'].size, int(p.denoising_strength * 255))
p.task_args['control_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images
p.task_args['width'] = p.width
p.task_args['height'] = p.height
elif 'control_image' in possible:
p.task_args['control_image'] = p.init_images # switch image and control_image
if 'control_mode' in possible:
p.task_args['control_mode'] = getattr(p, 'control_mode', None)
if 'strength' in possible:
p.task_args['strength'] = p.denoising_strength
p.init_images = [init_image] * len(active_model)
if hasattr(shared.sd_model, 'controlnet') and hasattr(p.task_args, 'control_image') and len(p.task_args['control_image']) > 1 and (shared.sd_model.__class__.__name__ == 'StableDiffusionXLControlNetUnionPipeline'): # special case for controlnet-union
p.task_args['control_image'] = [[x] for x in p.task_args['control_image']]
p.task_args['control_mode'] = [[x] for x in p.task_args['control_mode']]
if is_generator:
image_txt = f'{blended_image.width}x{blended_image.height}' if blended_image is not None else 'None'
msg = f'process | {index} of {frames if video is not None else len(inputs)} | {"Image" if video is None else "Frame"} {image_txt}'
debug_log(f'Control yield: {msg}')
if is_generator:
yield (None, blended_image, f'Control {msg}')
t2 += time.time() - t2
# determine txt2img, img2img, inpaint pipeline
if unit_type == 'reference' and has_models: # special case
p.is_control = True
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
elif not has_models: # run in txt2img/img2img/inpaint mode
if mask is not None:
p.task_args['strength'] = p.denoising_strength
p.image_mask = mask
p.init_images = input_image if isinstance(input_image, list) else [input_image]
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.INPAINTING)
elif processed_image is not None:
p.init_images = processed_image if isinstance(processed_image, list) else [processed_image]
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
else:
p.init_hr(p.scale_by, p.resize_name, force=True)
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
elif has_models: # actual control
p.is_control = True
if mask is not None:
p.task_args['strength'] = denoising_strength
p.image_mask = mask
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.INPAINTING) # only controlnet supports inpaint
if hasattr(p, 'init_images') and p.init_images is not None:
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) # only controlnet supports img2img
else:
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
if hasattr(p, 'init_images') and p.init_images is not None and 'image' in possible:
p.task_args['image'] = p.init_images # need to set explicitly for txt2img
del p.init_images
if unit_type == 'lite':
if input_type == 0:
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
shared.sd_model.no_task_switch = True
elif input_type == 1:
p.init_images = [input_image]
elif input_type == 2:
if init_image is None:
shared.log.warning('Control: separate init image not provided')
init_image = input_image
p.init_images = [init_image]
instance.apply(selected_models, processed_image, control_conditioning)
if hasattr(p, 'init_images') and p.init_images is None: # delete empty
del p.init_images
# final check
if has_models and shared.sd_model.__class__.__name__ not in unified_models:
if unit_type in ['controlnet', 't2i adapter', 'lite', 'xs'] \
and p.task_args.get('image', None) is None \
and p.task_args.get('control_image', None) is None \
and getattr(p, 'init_images', None) is None \
and getattr(p, 'image', None) is None:
if is_generator:
shared.log.debug(f'Control args: {p.task_args}')
yield terminate(f'Mode={p.extra_generation_params.get("Control type", None)} input image is none')
return [], '', '', 'Error: Input image is none'
# resize mask
if mask is not None and resize_mode_mask != 0 and resize_name_mask != 'None':
if selected_scale_tab_mask == 1:
width_mask, height_mask = int(input_image.width * scale_by_mask), int(input_image.height * scale_by_mask)
p.width, p.height = width_mask, height_mask
debug_log(f'Control resize: op=mask image={mask} width={width_mask} height={height_mask} mode={resize_mode_mask} name={resize_name_mask} context="{resize_context_mask}"')
# pipeline
output = None
script_run = False
if pipe is not None: # run new pipeline
if not hasattr(pipe, 'restore_pipeline') and video is None:
pipe.restore_pipeline = restore_pipeline
shared.sd_model.restore_pipeline = restore_pipeline
debug_log(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)} class={pipe.__class__}')
# debug_log(f'Control exec pipeline: p={vars(p)}')
# debug_log(f'Control exec pipeline: args={p.task_args} image={p.task_args.get("image", None)} control={p.task_args.get("control_image", None)} mask={p.task_args.get("mask_image", None) or p.image_mask} ref={p.task_args.get("ref_image", None)}')
if sd_models.get_diffusers_task(pipe) != sd_models.DiffusersTaskType.TEXT_2_IMAGE: # force vae back to gpu if not in txt2img mode
sd_models.move_model(pipe.vae, devices.device)
p.scripts = scripts.scripts_control
p.script_args = input_script_args or []
if len(p.script_args) == 0:
script_runner = scripts.scripts_control
if not script_runner.scripts:
script_runner.initialize_scripts(False)
p.script_args = script.init_default_script_args(script_runner)
# actual processing
if p.is_tile:
processed: processing.Processed = tile.run_tiling(p, input_image)
if processed is None and p.scripts is not None:
processed = p.scripts.run(p, *p.script_args)
if processed is None:
processed: processing.Processed = processing.process_images(p) # run actual pipeline
else:
script_run = True
# postprocessing
if p.scripts is not None:
processed = p.scripts.after(p, processed, *p.script_args)
output = None
if processed is not None:
output = processed.images
info_txt = [processed.infotext(p, i) for i in range(len(output))]
# output = pipe(**vars(p)).images # alternative direct pipe exec call
else: # blend all processed images and return
output = [processed_image]
t3 += time.time() - t3
# outputs
output = output or []
for i, output_image in enumerate(output):
if output_image is not None:
# resize after
is_grid = len(output) == p.batch_size * p.n_iter + 1 and i == 0
if selected_scale_tab_after == 1:
width_after = int(output_image.width * scale_by_after)
height_after = int(output_image.height * scale_by_after)
if resize_mode_after != 0 and resize_name_after != 'None' and not is_grid:
debug_log(f'Control resize: op=after image={output_image} width={width_after} height={height_after} mode={resize_mode_after} name={resize_name_after} context="{resize_context_after}"')
output_image = images.resize_image(resize_mode_after, output_image, width_after, height_after, resize_name_after, context=resize_context_after)
output_images.append(output_image)
if shared.opts.include_mask and not script_run:
if processed_image is not None and isinstance(processed_image, Image.Image):
output_images.append(processed_image)
if is_generator:
image_txt = f'{output_image.width}x{output_image.height}' if output_image is not None else 'None'
if video is not None:
msg = f'Control output | {index} of {frames} skip {video_skip_frames} | Frame {image_txt}'
else:
msg = f'Control output | {index} of {len(inputs)} | Image {image_txt}'
yield (output_image, blended_image, msg) # result is control_output, proces_output
if video is not None and frame is not None:
status, frame = video.read()
if status:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
debug_log(f'Control: video frame={index} frames={frames} status={status} skip={index % (video_skip_frames + 1)} progress={index/frames:.2f}')
else:
status = False
if video is not None:
video.release()
debug_log(f'Control: pipeline units={len(active_model)} process={len(active_process)} time={t3-t0:.2f} init={t1-t0:.2f} proc={t2-t1:.2f} ctrl={t3-t2:.2f} outputs={len(output_images)}')
process_timer.add('init', t1-t0)
process_timer.add('proc', t2-t1)
except Exception as e:
shared.log.error(f'Control pipeline failed: type={unit_type} units={len(active_model)} error={e}')
errors.display(e, 'Control')
if len(output_images) == 0:
output_images = None
image_txt = '| Images None'
else:
image_txt = ''
p.init_images = output_images # may be used for hires
if video_type != 'None' and isinstance(output_images, list):
p.do_not_save_grid = True # pylint: disable=attribute-defined-outside-init
output_filename = images.save_video(p, filename=None, images=output_images, video_type=video_type, duration=video_duration, loop=video_loop, pad=video_pad, interpolate=video_interpolate, sync=True)
if shared.opts.gradio_skip_video:
output_filename = ''
image_txt = f'| Frames {len(output_images)} | Size {output_images[0].width}x{output_images[0].height}'
p.close()
restore_pipeline()
debug_log(f'Ready: {image_txt}')
html_txt = f'<p>Ready {image_txt}</p>' if image_txt != '' else ''
if len(info_txt) > 0:
html_txt = html_txt + infotext_to_html(info_txt[0])
if is_generator:
yield (output_images, blended_image, html_txt, output_filename)
return (output_images, blended_image, html_txt, output_filename)