automatic/modules/control/run.py

import os
import time
import math
from typing import List, Union
import cv2
import numpy as np
import diffusers
from PIL import Image
from modules.control import util
from modules.control import unit
from modules.control import processors
from modules.control import controlnets # lllyasviel ControlNet
from modules.control import controlnetsxs # VisLearn ControlNet-XS
from modules.control import adapters # TencentARC T2I-Adapter
from modules.control import reference # ControlNet-Reference
from modules import devices, shared, errors, processing, images, sd_models, sd_samplers


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


class ControlProcessing(processing.StableDiffusionProcessingImg2Img):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.strength = None
        self.adapter_conditioning_scale = None
        self.adapter_conditioning_factor = None
        self.guess_mode = None
        self.controlnet_conditioning_scale = None
        self.control_guidance_start = None
        self.control_guidance_end = None
        self.reference_attn = None
        self.reference_adain = None
        self.attention_auto_machine_weight = None
        self.gn_auto_machine_weight = None
        self.style_fidelity = None
        self.ref_image = None
        self.image = None
        self.query_weight = None
        self.adain_weight = None
        self.adapter_conditioning_factor = 1.0
        self.attention = 'Attention'
        self.fidelity = 0.5
        self.override = None

    def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): # abstract
        pass


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, unit_type: str, is_generator: bool,
                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, full_quality, restore_faces, tiling,
                hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
                resize_mode, resize_name, width, height, scale_by, selected_scale_tab, resize_time,
                denoising_strength, batch_count, batch_size,
                video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate,
        ):
    global pipe, original_pipeline # pylint: disable=global-statement
    debug(f'Control {unit_type}: input={inputs}')
    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[controlnets.ControlNet, controlnetsxs.ControlNetXS, adapters.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
    width = 8 * math.ceil(width / 8)
    height = 8 * math.ceil(height / 8)

    p = ControlProcessing(
        prompt = prompt,
        negative_prompt = negative,
        styles = styles,
        steps = steps,
        sampler_name = sd_samplers.samplers[sampler_index].name,
        latent_sampler = sd_samplers.samplers[sampler_index].name,
        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,
        full_quality = full_quality,
        restore_faces = restore_faces,
        tiling = tiling,
        hdr_clamp = hdr_clamp,
        hdr_boundary = hdr_boundary,
        hdr_threshold = hdr_threshold,
        hdr_center = hdr_center,
        hdr_channel_shift = hdr_channel_shift,
        hdr_full_shift = hdr_full_shift,
        hdr_maximize = hdr_maximize,
        hdr_max_center = hdr_max_center,
        hdr_max_boundry = hdr_max_boundry,
        resize_mode = resize_mode,
        resize_name = resize_name,
        scale_by = scale_by,
        selected_scale_tab = selected_scale_tab,
        denoising_strength = denoising_strength,
        n_iter = batch_count,
        batch_size = batch_size,
    )
    if resize_mode != 0 or inputs is None or inputs == [None]:
        p.width = width # pylint: disable=attribute-defined-outside-init
        p.height = height # pylint: disable=attribute-defined-outside-init
        if selected_scale_tab == 1:
            width = int(width * scale_by)
            height = int(height * scale_by)
    else:
        del p.width
        del p.height


    t0 = time.time()
    for u in units:
        if not u.enabled or u.type != unit_type:
            continue
        if unit_type == 'adapter' and u.adapter.model is not None:
            active_process.append(u.process)
            active_model.append(u.adapter)
            active_strength.append(u.strength)
            p.adapter_conditioning_factor = u.factor
            shared.log.debug(f'Control T2I-Adapter unit: 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(u.strength)
            active_start.append(u.start)
            active_end.append(u.end)
            p.guess_mode = u.guess
            shared.log.debug(f'Control ControlNet unit: 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(u.strength)
            active_start.append(u.start)
            active_end.append(u.end)
            p.guess_mode = u.guess
            shared.log.debug(f'Control ControlNet-XS unit: 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 = u.query_weight
            p.adain_weight = u.adain_weight
            p.fidelity = u.fidelity
            shared.log.debug('Control Reference unit')
        else:
            active_process.append(u.process)
            # active_model.append(model)
            active_strength.append(u.strength)
    """
    if (len(active_process) == 0) and (unit_type != 'reference'):
        msg = 'Control: no active units'
        shared.log.warning(msg)
        restore_pipeline()
        return msg
    """
    p.ops.append('control')

    has_models = False
    selected_models: List[Union[controlnets.ControlNetModel, controlnetsxs.ControlNetXSModel, adapters.AdapterModel]] = None
    if unit_type == 'adapter' or unit_type == 'controlnet' or unit_type == 'xs':
        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
        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
        use_conditioning = active_strength[0] if len(active_strength) == 1 else list(active_strength) # strength or list[strength]
    else:
        pass
    shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) # reset current pipeline

    debug(f'Control: run type={unit_type} models={has_models}')
    if unit_type == 'adapter' and has_models:
        p.extra_generation_params["Control mode"] = 'Adapter'
        p.extra_generation_params["Control conditioning"] = use_conditioning
        p.task_args['adapter_conditioning_scale'] = use_conditioning
        instance = adapters.AdapterPipeline(selected_models, shared.sd_model)
        pipe = instance.pipeline
    elif unit_type == 'controlnet' and has_models:
        p.extra_generation_params["Control mode"] = 'ControlNet'
        p.extra_generation_params["Control conditioning"] = use_conditioning
        p.task_args['controlnet_conditioning_scale'] = use_conditioning
        p.task_args['control_guidance_start'] = active_start[0] if len(active_start) == 1 else list(active_start)
        p.task_args['control_guidance_end'] = active_end[0] if len(active_end) == 1 else list(active_end)
        p.task_args['guess_mode'] = p.guess_mode
        instance = controlnets.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.extra_generation_params["Control conditioning"] = use_conditioning
        p.controlnet_conditioning_scale = use_conditioning
        p.control_guidance_start = active_start[0] if len(active_start) == 1 else list(active_start)
        p.control_guidance_end = active_end[0] if len(active_end) == 1 else list(active_end)
        instance = controlnetsxs.ControlNetPipeline(selected_models, shared.sd_model)
        pipe = instance.pipeline
    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
    else: # run in img2img mode
        if len(active_strength) > 0:
            p.strength = active_strength[0]
        pipe = diffusers.AutoPipelineForImage2Image.from_pipe(shared.sd_model) # use set_diffuser_pipe
    debug(f'Control pipeline: class={pipe.__class__} 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
    if pipe is not None:
        shared.sd_model = pipe
        if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
            shared.sd_model.to(shared.device)
        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
                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()
                    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):
                    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(input_image)
                        except Exception as e:
                            shared.log.error(f'Control: image open failed: path={input_image} error={e}')
                            continue
                    index += 1
                    if video is not None and index % (video_skip_frames + 1) != 0:
                        continue

                    # resize
                    if resize_mode != 0 and input_image is not None:
                        p.extra_generation_params["Control resize"] = f'{resize_time}: {resize_name}'
                    if resize_mode != 0 and input_image is not None and resize_time == 'Before':
                        debug(f'Control resize: image={input_image} width={width} height={height} mode={resize_mode} name={resize_name} sequence={resize_time}')
                        input_image = images.resize_image(resize_mode, input_image, width, height, resize_name)

                    # process
                    if input_image is None:
                        p.image = None
                        processed_image = None
                        debug('Control: process=None image=None')
                    elif len(active_process) == 0 and unit_type == 'reference':
                        p.ref_image = p.override or input_image
                        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
                        processed_image = p.ref_image
                    elif len(active_process) == 1:
                        p.image = active_process[0](input_image)
                        p.task_args['image'] = p.image
                        p.extra_generation_params["Control process"] = active_process[0].processor_id
                        debug(f'Control: process={active_process[0].processor_id} image={p.image}')
                        if p.image is None:
                            msg = 'Control: attempting process but output is none'
                            shared.log.error(msg)
                            restore_pipeline()
                            return msg
                        processed_image = p.image
                    else:
                        if len(active_process) > 0:
                            p.image = [p(input_image) for p in active_process] # list[image]
                        else:
                            p.image = [input_image]
                        p.task_args['image'] = p.image
                        p.extra_generation_params["Control process"] = [p.processor_id for p in active_process]
                        debug(f'Control: process={[p.processor_id for p in active_process]} image={p.image}')
                        if any(img is None for img in p.image):
                            msg = 'Control: attempting process but output is none'
                            shared.log.error(msg)
                            restore_pipeline()
                            return msg
                        processed_image = [np.array(i) for i in p.image]
                        processed_image = util.blend(processed_image) # blend all processed images into one
                        processed_image = Image.fromarray(processed_image)
                    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

                    # pipeline
                    output = None
                    if pipe is not None: # run new pipeline
                        if not has_models and (unit_type == 'controlnet' or unit_type == 'adapter' or unit_type == 'xs'): # run in img2img mode
                            if p.image is None:
                                if hasattr(p, 'init_images'):
                                    del p.init_images
                                shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) # reset current pipeline
                            else:
                                p.init_images = [processed_image] # pylint: disable=attribute-defined-outside-init
                                shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) # reset current pipeline
                        else:
                            if hasattr(p, 'init_images'):
                                del p.init_images
                            shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) # reset current pipeline
                        debug(f'Control exec pipeline: class={pipe.__class__} p={vars(p)}')
                        debug(f'Control exec pipeline: class={pipe.__class__} args={p.task_args}')
                        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
                            if resize_mode != 0 and resize_time == 'After':
                                debug(f'Control resize: image={input_image} width={width} height={height} mode={resize_mode} name={resize_name} sequence={resize_time}')
                                output_image = images.resize_image(resize_mode, output_image, width, height, resize_name)
                            elif hasattr(p, 'width') and hasattr(p, 'height'):
                                output_image = output_image.resize((p.width, p.height), Image.Resampling.LANCZOS)

                            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()
                    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)}'
    debug(f'Control ready: {image_txt}')
    restore_pipeline()
    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)