mirror of https://github.com/vladmandic/automatic
553 lines
33 KiB
Python
553 lines
33 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.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
|
|
|
|
|
|
debug = shared.log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
|
|
debug('Trace: CONTROL')
|
|
pipe = None
|
|
original_pipeline = None
|
|
|
|
|
|
def restore_pipeline():
|
|
global pipe # pylint: disable=global-statement
|
|
pipe = None
|
|
if original_pipeline is not None:
|
|
shared.sd_model = original_pipeline
|
|
debug(f'Control restored pipeline: class={shared.sd_model.__class__.__name__}')
|
|
devices.torch_gc()
|
|
|
|
|
|
def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_generator: bool, input_type: int,
|
|
prompt, negative, styles, steps, sampler_index,
|
|
seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
|
|
cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, cfg_end, full_quality, restore_faces, tiling,
|
|
hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio,
|
|
resize_mode_before, resize_name_before, width_before, height_before, scale_by_before, selected_scale_tab_before,
|
|
resize_mode_after, resize_name_after, width_after, height_after, scale_by_after, selected_scale_tab_after,
|
|
denoising_strength, batch_count, batch_size,
|
|
video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate,
|
|
ip_adapter, ip_scale, ip_image,
|
|
*input_script_args # pylint: disable=unused-argument
|
|
):
|
|
global pipe, original_pipeline # pylint: disable=global-statement
|
|
debug(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
|
|
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
|
|
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
|
|
|
|
p = StableDiffusionProcessingControl(
|
|
prompt = prompt,
|
|
negative_prompt = negative,
|
|
styles = styles,
|
|
steps = steps,
|
|
sampler_name = processing.get_sampler_name(sampler_index),
|
|
hr_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,
|
|
cfg_scale = cfg_scale,
|
|
clip_skip = clip_skip,
|
|
image_cfg_scale = image_cfg_scale,
|
|
diffusers_guidance_rescale = diffusers_guidance_rescale,
|
|
sag_scale = sag_scale,
|
|
full_quality = full_quality,
|
|
restore_faces = restore_faces,
|
|
tiling = tiling,
|
|
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,
|
|
n_iter = batch_count,
|
|
batch_size = batch_size,
|
|
inpaint_full_res = masking.opts.mask_only,
|
|
# inpaint_full_res_padding = masking.opts.mask_padding,
|
|
inpainting_mask_invert = 1 if masking.opts.invert else 0,
|
|
inpainting_fill = 1,
|
|
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,
|
|
outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_control_samples,
|
|
outpath_grids=shared.opts.outdir_grids or shared.opts.outdir_control_grids,
|
|
)
|
|
processing.process_init(p)
|
|
|
|
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
|
|
else:
|
|
del p.width
|
|
del p.height
|
|
|
|
t0 = time.time()
|
|
num_units = 0
|
|
for u in units:
|
|
if u.type != unit_type:
|
|
continue
|
|
num_units += 1
|
|
debug(f'Control unit: i={num_units} type={u.type} enabled={u.enabled}')
|
|
if not u.enabled:
|
|
continue
|
|
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:
|
|
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
|
|
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}')
|
|
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:
|
|
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))
|
|
p.ops.append('control')
|
|
debug(f'Control active: process={len(active_process)} model={len(active_model)}')
|
|
|
|
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.extra_generation_params["Control model"] = (active_model[0].model_id or '') if active_model[0].model is not None else None
|
|
has_models = selected_models is not None
|
|
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]
|
|
p.extra_generation_params["Control model"] = ', '.join([(m.model_id or '') 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)
|
|
p.extra_generation_params["Control conditioning"] = control_conditioning
|
|
else:
|
|
pass
|
|
|
|
debug(f'Control: run type={unit_type} models={has_models}')
|
|
if unit_type == 't2i adapter' and has_models:
|
|
p.extra_generation_params["Control mode"] = '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 mode"] = 'ControlNet'
|
|
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
|
|
instance = controlnet.ControlNetPipeline(selected_models, shared.sd_model)
|
|
pipe = instance.pipeline
|
|
elif unit_type == 'xs' and has_models:
|
|
p.extra_generation_params["Control mode"] = '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 mode"] = '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':
|
|
p.extra_generation_params["Control mode"] = '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
|
|
"""
|
|
try:
|
|
pipe = diffusers.AutoPipelineForText2Image.from_pipe(shared.sd_model) # use set_diffuser_pipe
|
|
except Exception as e:
|
|
shared.log.warning(f'Control pipeline create: {e}')
|
|
pipe = shared.sd_model
|
|
"""
|
|
|
|
|
|
debug(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
|
|
|
|
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(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)
|
|
|
|
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():
|
|
msg = f'Control: video open failed: path={inputs}'
|
|
shared.log.error(msg)
|
|
restore_pipeline()
|
|
return msg
|
|
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:
|
|
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:
|
|
msg = f'Control: video open failed: path={inputs} {e}'
|
|
shared.log.error(msg)
|
|
restore_pipeline()
|
|
return msg
|
|
|
|
while status:
|
|
processed_image = None
|
|
if frame is not None:
|
|
inputs = [Image.fromarray(frame)] # cv2 to pil
|
|
for i, input_image in enumerate(inputs):
|
|
debug(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
|
|
restore_pipeline()
|
|
return 'Control 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('Control Init image: same as control')
|
|
init_image = input_image
|
|
elif inits is None:
|
|
debug('Control Init image: none')
|
|
init_image = None
|
|
elif isinstance(inits[i], str):
|
|
debug(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(f'Control Init image: {i % len(inits) + 1} of {len(inits)}')
|
|
init_image = inits[i % len(inits)]
|
|
index += 1
|
|
if video is not None and index % (video_skip_frames + 1) != 0:
|
|
continue
|
|
|
|
# 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(f'Control resize: op=before image={input_image} width={width_before} height={height_before} mode={resize_mode_before} name={resize_name_before}')
|
|
input_image = images.resize_image(resize_mode_before, input_image, width_before, height_before, resize_name_before)
|
|
if input_image is not None and init_image is not None and init_image.size != input_image.size:
|
|
debug(f'Control resize init: image={p.override} 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(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(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
|
|
if len(active_process) > 0:
|
|
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]
|
|
image_mode = 'L' if unit_type == 't2i adapter' and len(active_model) > i and ('Canny' in active_model[i].model_id or 'Sketch' in active_model[i].model_id) else 'RGB' # t2iadapter canny and sketch work in grayscale only
|
|
debug(f'Control: i={i+1} process="{process.processor_id}" input={masked_image} override={process.override}')
|
|
processed_image = process(
|
|
image_input=masked_image,
|
|
mode=image_mode,
|
|
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(f'Control processed: {len(processed_images)}')
|
|
if len(processed_images) > 0:
|
|
p.extra_generation_params["Control process"] = [p.processor_id for p in active_process if p.processor_id is not None]
|
|
if len(p.extra_generation_params["Control process"]) == 0:
|
|
p.extra_generation_params["Control process"] = None
|
|
if any(img is None for img in processed_images):
|
|
msg = 'Control: attempting process but output is none'
|
|
shared.log.error(f'{msg}: {processed_images}')
|
|
restore_pipeline()
|
|
return msg
|
|
if len(processed_images) > 1:
|
|
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)
|
|
else:
|
|
processed_image = processed_images[0]
|
|
if isinstance(selected_models, list) and len(processed_images) == len(selected_models):
|
|
debug(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):
|
|
msg = f'Control: number of inputs does not match: input={len(processed_images)} models={len(selected_models)}'
|
|
shared.log.error(msg)
|
|
restore_pipeline()
|
|
return msg
|
|
elif selected_models is not None:
|
|
if len(processed_images) > 1:
|
|
debug('Control: using blended image for single model')
|
|
p.init_images = [processed_image]
|
|
else:
|
|
debug('Control processed: using input direct')
|
|
processed_image = input_image
|
|
|
|
if unit_type == 'reference':
|
|
p.ref_image = p.override or input_image
|
|
p.task_args.pop('image', None)
|
|
p.task_args['ref_image'] = p.ref_image
|
|
debug(f'Control: process=None image={p.ref_image}')
|
|
if p.ref_image is None:
|
|
msg = 'Control: attempting reference mode but image is none'
|
|
shared.log.error(msg)
|
|
restore_pipeline()
|
|
return msg
|
|
elif unit_type == 'controlnet' and input_type == 1: # Init image same as control
|
|
p.task_args['control_image'] = p.init_images # switch image and control_image
|
|
p.init_images = [p.override or input_image] * len(active_model)
|
|
p.task_args['strength'] = p.denoising_strength
|
|
elif unit_type == 'controlnet' and input_type == 2: # Separate init image
|
|
if init_image is None:
|
|
shared.log.warning('Control: separate init image not provided')
|
|
init_image = input_image
|
|
p.task_args['control_image'] = p.init_images # switch image and control_image
|
|
p.init_images = [init_image] * len(active_model)
|
|
p.task_args['strength'] = p.denoising_strength
|
|
|
|
if is_generator:
|
|
image_txt = f'{processed_image.width}x{processed_image.height}' if processed_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(f'Control yield: {msg}')
|
|
yield (None, processed_image, f'Control {msg}')
|
|
t2 += time.time() - t2
|
|
|
|
# determine txt2img, img2img, inpaint pipeline
|
|
if unit_type == 'reference': # 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]
|
|
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]
|
|
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
|
|
else:
|
|
p.init_hr()
|
|
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
|
|
elif 'control_image' in p.task_args:
|
|
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:
|
|
p.task_args['image'] = p.init_images # need to set explicitly for txt2img
|
|
if unit_type == 'lite':
|
|
instance.apply(selected_models, p.init_images, control_conditioning)
|
|
if hasattr(p, 'init_images') and p.init_images is None: # delete as its set via task_args
|
|
del p.init_images
|
|
|
|
# ip adapter apply is run in processing.process_images
|
|
p.ip_adapter_name = ip_adapter
|
|
p.ip_adapter_scale = ip_scale
|
|
p.ip_adapter_image = ip_image or input_image
|
|
|
|
# pipeline
|
|
output = None
|
|
if pipe is not None: # run new pipeline
|
|
debug(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)} class={pipe.__class__}')
|
|
debug(f'Control exec pipeline: p={vars(p)}')
|
|
debug(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)}')
|
|
p.scripts = scripts.scripts_control
|
|
p.script_args = input_script_args
|
|
processed = p.scripts.run(p, *input_script_args)
|
|
if processed is None:
|
|
processed: processing.Processed = processing.process_images(p) # run actual pipeline
|
|
output = processed.images if processed is not None else None
|
|
# 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
|
|
if output is not None and len(output) > 0:
|
|
output_image = output[0]
|
|
if output_image is not None:
|
|
|
|
# resize after
|
|
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':
|
|
debug(f'Control resize: op=after image={output_image} width={width_after} height={height_after} mode={resize_mode_after} name={resize_name_after}')
|
|
output_image = images.resize_image(resize_mode_after, output_image, width_after, height_after, resize_name_after)
|
|
|
|
output_images.append(output_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, processed_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(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()
|
|
|
|
shared.log.info(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)}')
|
|
except Exception as e:
|
|
shared.log.error(f'Control pipeline failed: type={unit_type} units={len(active_model)} error={e}')
|
|
errors.display(e, 'Control')
|
|
|
|
shared.sd_model = original_pipeline
|
|
pipe = None
|
|
devices.torch_gc()
|
|
|
|
if len(output_images) == 0:
|
|
output_images = None
|
|
image_txt = 'images=None'
|
|
elif len(output_images) == 1:
|
|
output_images = output_images[0]
|
|
image_txt = f'| Images 1 | Size {output_images.width}x{output_images.height}' if output_image is not None else 'None'
|
|
else:
|
|
image_txt = f'| Images {len(output_images)} | Size {output_images[0].width}x{output_images[0].height}' if output_image is not None else 'None'
|
|
|
|
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)
|
|
image_txt = f'| Frames {len(output_images)} | Size {output_images[0].width}x{output_images[0].height}'
|
|
|
|
image_txt += f' | {util.dict2str(p.extra_generation_params)}'
|
|
if hasattr(instance, 'restore'):
|
|
instance.restore()
|
|
restore_pipeline()
|
|
debug(f'Control ready: {image_txt}')
|
|
if is_generator:
|
|
yield (output_images, processed_image, f'Control ready {image_txt}', output_filename)
|
|
else:
|
|
return (output_images, processed_image, f'Control ready {image_txt}', output_filename)
|