scripts/batch_hijack.py

import os
from copy import copy
from enum import Enum
from typing import Tuple, List

from modules import img2img, processing, shared, script_callbacks
from scripts import external_code


class BatchHijack:
    def __init__(self):
        self.is_batch = False
        self.batch_index = 0
        self.batch_size = 1
        self.init_seed = None
        self.init_subseed = None
        self.process_batch_callbacks = [self.on_process_batch]
        self.process_batch_each_callbacks = []
        self.postprocess_batch_each_callbacks = [self.on_postprocess_batch_each]
        self.postprocess_batch_callbacks = [self.on_postprocess_batch]

    def img2img_process_batch_hijack(self, p, *args, **kwargs):
        cn_is_batch, batches, output_dir, _ = get_cn_batches(p)
        if not cn_is_batch:
            return getattr(img2img, '__controlnet_original_process_batch')(p, *args, **kwargs)

        self.dispatch_callbacks(self.process_batch_callbacks, p, batches, output_dir)

        try:
            return getattr(img2img, '__controlnet_original_process_batch')(p, *args, **kwargs)
        finally:
            self.dispatch_callbacks(self.postprocess_batch_callbacks, p)

    def processing_process_images_hijack(self, p, *args, **kwargs):
        if self.is_batch:
            # we are in img2img batch tab, do a single batch iteration
            return self.process_images_cn_batch(p, *args, **kwargs)

        cn_is_batch, batches, output_dir, input_file_names = get_cn_batches(p)
        if not cn_is_batch:
            # we are not in batch mode, fallback to original function
            return getattr(processing, '__controlnet_original_process_images_inner')(p, *args, **kwargs)

        output_images = []
        try:
            self.dispatch_callbacks(self.process_batch_callbacks, p, batches, output_dir)

            for batch_i in range(self.batch_size):
                processed = self.process_images_cn_batch(p, *args, **kwargs)
                if shared.opts.data.get('controlnet_show_batch_images_in_ui', False):
                    output_images.extend(processed.images[processed.index_of_first_image:])

                if output_dir:
                    self.save_images(output_dir, input_file_names[batch_i], processed.images[processed.index_of_first_image:])

                if shared.state.interrupted:
                    break

        finally:
            self.dispatch_callbacks(self.postprocess_batch_callbacks, p)

        if output_images:
            processed.images = output_images
        else:
            processed = processing.Processed(p, [], p.seed)

        return processed

    def process_images_cn_batch(self, p, *args, **kwargs):
        self.dispatch_callbacks(self.process_batch_each_callbacks, p)
        old_detectmap_output = shared.opts.data.get('control_net_no_detectmap', False)
        try:
            shared.opts.data.update({'control_net_no_detectmap': True})
            processed = getattr(processing, '__controlnet_original_process_images_inner')(p, *args, **kwargs)
        finally:
            shared.opts.data.update({'control_net_no_detectmap': old_detectmap_output})

        self.dispatch_callbacks(self.postprocess_batch_each_callbacks, p, processed)

        # do not go past control net batch size
        if self.batch_index >= self.batch_size:
            shared.state.interrupted = True

        return processed

    def save_images(self, output_dir, init_image_path, output_images):
        os.makedirs(output_dir, exist_ok=True)
        for n, processed_image in enumerate(output_images):
            filename = os.path.basename(init_image_path)

            if n > 0:
                left, right = os.path.splitext(filename)
                filename = f"{left}-{n}{right}"

            if processed_image.mode == 'RGBA':
                processed_image = processed_image.convert("RGB")
            processed_image.save(os.path.join(output_dir, filename))

    def do_hijack(self):
        script_callbacks.on_script_unloaded(self.undo_hijack)
        hijack_function(
            module=img2img,
            name='process_batch',
            new_name='__controlnet_original_process_batch',
            new_value=self.img2img_process_batch_hijack,
        )
        hijack_function(
            module=processing,
            name='process_images_inner',
            new_name='__controlnet_original_process_images_inner',
            new_value=self.processing_process_images_hijack
        )

    def undo_hijack(self):
        unhijack_function(
            module=img2img,
            name='process_batch',
            new_name='__controlnet_original_process_batch',
        )
        unhijack_function(
            module=processing,
            name='process_images_inner',
            new_name='__controlnet_original_process_images_inner',
        )

    def adjust_job_count(self, p):
        if shared.state.job_count == -1:
            shared.state.job_count = p.n_iter
        shared.state.job_count *= self.batch_size

    def on_process_batch(self, p, batches, output_dir, *args):
        print('controlnet batch mode')
        self.is_batch = True
        self.batch_index = 0
        self.batch_size = len(batches)
        processing.fix_seed(p)
        if shared.opts.data.get('controlnet_increment_seed_during_batch', False):
            self.init_seed = p.seed
            self.init_subseed = p.subseed
        self.adjust_job_count(p)
        p.do_not_save_grid = True
        p.do_not_save_samples = bool(output_dir)

    def on_postprocess_batch_each(self, p, *args):
        self.batch_index += 1
        if shared.opts.data.get('controlnet_increment_seed_during_batch', False):
            p.seed = p.seed + len(p.all_prompts)
            p.subseed = p.subseed + len(p.all_prompts)

    def on_postprocess_batch(self, p, *args):
        self.is_batch = False
        self.batch_index = 0
        self.batch_size = 1
        if shared.opts.data.get('controlnet_increment_seed_during_batch', False):
            p.seed = self.init_seed
            p.all_seeds = [self.init_seed]
            p.subseed = self.init_subseed
            p.all_subseeds = [self.init_subseed]

    def dispatch_callbacks(self, callbacks, *args):
        for callback in callbacks:
            callback(*args)


def hijack_function(module, name, new_name, new_value):
    # restore original function in case of reload
    unhijack_function(module=module, name=name, new_name=new_name)
    setattr(module, new_name, getattr(module, name))
    setattr(module, name, new_value)


def unhijack_function(module, name, new_name):
    if hasattr(module, new_name):
        setattr(module, name, getattr(module, new_name))
        delattr(module, new_name)


class InputMode(Enum):
    SIMPLE = "simple"
    BATCH = "batch"


def get_cn_batches(p: processing.StableDiffusionProcessing) -> Tuple[bool, List[List[str]], str, List[str]]:
    units = external_code.get_all_units_in_processing(p)
    units = [copy(unit) for unit in units if getattr(unit, 'enabled', False)]
    any_unit_is_batch = False
    output_dir = ''
    input_file_names = []
    for unit in units:
        if getattr(unit, 'input_mode', InputMode.SIMPLE) == InputMode.BATCH:
            any_unit_is_batch = True
            output_dir = getattr(unit, 'output_dir', '')
            if isinstance(unit.batch_images, str):
                unit.batch_images = shared.listfiles(unit.batch_images)
                input_file_names = unit.batch_images

    if any_unit_is_batch:
        cn_batch_size = min(len(getattr(unit, 'batch_images', []))
                         for unit in units
                         if getattr(unit, 'input_mode', InputMode.SIMPLE) == InputMode.BATCH)
    else:
        cn_batch_size = 1

    batches = [[] for _ in range(cn_batch_size)]
    for i in range(cn_batch_size):
        for unit in units:
            if getattr(unit, 'input_mode', InputMode.SIMPLE) == InputMode.SIMPLE:
                batches[i].append(unit.image)
            else:
                batches[i].append(unit.batch_images[i])

    return any_unit_is_batch, batches, output_dir, input_file_names


instance = BatchHijack()