major controlnet work, xinsir promax and tiling support

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/control/run.py

@@ -7,6 +7,7 @@ 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
@@ -83,6 +84,7 @@ def control_run(state: str = '',
                 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
 
@@ -149,6 +151,7 @@ def control_run(state: str = '',
         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
 
@@ -242,7 +245,7 @@ def control_run(state: str = '',
             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}')
+            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)
@@ -250,8 +253,12 @@ def control_run(state: str = '',
             active_start.append(float(u.start))
             active_end.append(float(u.end))
             p.guess_mode = u.guess
-            p.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}')
+            if isinstance(u.mode, str):
+                p.control_mode = 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)
@@ -291,6 +298,7 @@ def control_run(state: str = '',
             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.lower()
             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
@@ -305,29 +313,30 @@ def control_run(state: str = '',
         has_models = any(u.enabled for u in units if u.type == 'reference')
     else:
         pass
+    p.is_tile = p.is_tile and has_models
 
     def set_pipe():
         global pipe, instance # pylint: disable=global-statement
         pipe = None
         if has_models:
             p.ops.append('control')
-            p.extra_generation_params["Control mode"] = unit_type # overriden later with pretty-print
+            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 model"] = ';'.join([(m.model_id or '') for m in active_model if m.model is not None])
             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 mode"] = 'T2I-Adapter'
+            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 mode"] = 'ControlNet'
+            p.extra_generation_params["Control type"] = '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
@@ -335,7 +344,7 @@ def control_run(state: str = '',
             instance = controlnet.ControlNetPipeline(selected_models, shared.sd_model, p=p)
             pipe = instance.pipeline
         elif unit_type == 'xs' and has_models:
-            p.extra_generation_params["Control mode"] = 'ControlNet-XS'
+            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
@@ -344,14 +353,14 @@ def control_run(state: str = '',
             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.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 mode"] = 'Reference'
+            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
@@ -393,6 +402,8 @@ def control_run(state: str = '',
     else:
         original_pipeline = None
 
+    possible = sd_models.get_call(pipe).keys()
+
     try:
         with devices.inference_context():
             if isinstance(inputs, str): # only video, the rest is a list
@@ -562,19 +573,29 @@ def control_run(state: str = '',
                             return [], '', '', 'Reference mode without image'
                     elif unit_type == 'controlnet' and has_models:
                         if input_type == 0: # Control only
-                            if shared.sd_model_type in ['f1', 'sd3'] and 'control_image' not in p.task_args:
-                                p.task_args['control_image'] = p.init_images # some controlnets mandate this
+                            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'] = p.control_mode
+                            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.task_args['control_image'] = p.init_images # switch image and control_image
-                            p.task_args['strength'] = p.denoising_strength
+                            if 'control_image' in possible:
+                                p.task_args['control_image'] = p.init_images # switch image and control_image
+                            if 'strength' in possible:
+                                p.task_args['strength'] = p.denoising_strength
                             p.init_images = [p.override or input_image] * len(active_model)
                         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
-                            p.task_args['control_image'] = p.init_images # switch image and control_image
-                            p.task_args['strength'] = p.denoising_strength
+                            if 'control_image' in possible:
+                                p.task_args['control_image'] = p.init_images # switch image and control_image
+                            if 'strength' in possible:
+                                p.task_args['strength'] = p.denoising_strength
                             p.init_images = [init_image] * len(active_model)
 
                     if is_generator:
@@ -607,11 +628,11 @@ def control_run(state: str = '',
                             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:
+                        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:
+                            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':
@@ -624,9 +645,14 @@ def control_run(state: str = '',
 
                     # final check
                     if has_models:
-                        if unit_type in ['controlnet', 't2i adapter', 'lite', 'xs'] and p.task_args.get('image', None) is None and getattr(p, 'init_images', None) is None:
+                        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:
-                                yield terminate(f'Mode={p.extra_generation_params.get("Control mode", None)} input image is none')
+                                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
@@ -656,11 +682,17 @@ def control_run(state: str = '',
                                 script_runner.initialize_scripts(False)
                             p.script_args = script.init_default_script_args(script_runner)
 
-                        processed = p.scripts.run(p, *p.script_args)
+                        # actual processing
+                        if p.is_tile:
+                            processed: processing.Processed = tile.run_tiling(p, input_image)
+                        if processed is 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
                         processed = p.scripts.after(p, processed, *p.script_args)
                         output = None
                         if processed is not None:

modules/control/tile.py

@@ -0,0 +1,75 @@
+from PIL import Image
+from modules import shared, processing, images, sd_models
+
+
+def get_tile(image: Image.Image, x: int, y: int, sx: int, sy: int) -> Image.Image:
+    return image.crop((
+        (x + 0) * image.width // sx,
+        (y + 0) * image.height // sy,
+        (x + 1) * image.width // sx,
+        (y + 1) * image.height // sy
+    ))
+
+
+def set_tile(image: Image.Image, x: int, y: int, tiled: Image.Image):
+    image.paste(tiled, (x * tiled.width, y * tiled.height))
+    return image
+
+
+def run_tiling(p: processing.StableDiffusionProcessing, input_image: Image.Image) -> processing.Processed:
+    # prepare images
+    sx, sy = p.control_tile.split('x')
+    sx = int(sx)
+    sy = int(sy)
+    if sx <= 0 or sy <= 0:
+        raise ValueError('Control: invalid tile size')
+    control_image = p.task_args.get('control_image', None) or p.task_args.get('image', None)
+    control_upscaled = None
+    if isinstance(control_image, list) and len(control_image) > 0:
+        control_upscaled = images.resize_image(resize_mode=1 if sx==sy else 5,
+                                               im=control_image[0],
+                                               width=8 * int(sx * control_image[0].width) // 8,
+                                               height=8 * int(sy * control_image[0].height) // 8,
+                                               context='add with forward'
+                                              )
+    init_image = p.override or input_image
+    init_upscaled = None
+    if init_image is not None:
+        init_upscaled = images.resize_image(resize_mode=1 if sx==sy else 5,
+                                            im=init_image,
+                                            width=8 * int(sx * init_image.width) // 8,
+                                            height=8 * int(sy * init_image.height) // 8,
+                                            context='add with forward'
+                                           )
+
+    # stop processing from restoring pipeline on each iteration
+    orig_restore_pipeline = getattr(shared.sd_model, 'restore_pipeline', None)
+    shared.sd_model.restore_pipeline = None
+
+    # run tiling
+    for x in range(sx):
+        for y in range(sy):
+            shared.log.info(f'Control Tile: tile={x+1}-{sx}/{y+1}-{sy} target={control_upscaled}')
+            shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
+            p.init_images = None
+            p.task_args['control_mode'] = p.control_mode
+            p.task_args['strength'] = p.denoising_strength
+            if init_upscaled is not None:
+                p.task_args['image'] = [get_tile(init_upscaled, x, y, sx, sy)]
+            if control_upscaled is not None:
+                p.task_args['control_image'] = [get_tile(control_upscaled, x, y, sx, sy)]
+            processed: processing.Processed = processing.process_images(p) # run actual pipeline
+            if processed is None or len(processed.images) == 0:
+                continue
+            control_upscaled = set_tile(control_upscaled, x, y, processed.images[0])
+
+    # post-process
+    p.width = control_upscaled.width
+    p.height = control_upscaled.height
+    processed.images = [control_upscaled]
+    processed.info = processed.infotext(p, 0)
+    processed.infotexts = [processed.info]
+    shared.sd_model.restore_pipeline = orig_restore_pipeline
+    if hasattr(shared.sd_model, 'restore_pipeline') and shared.sd_model.restore_pipeline is not None:
+        shared.sd_model.restore_pipeline()
+    return processed

modules/control/unit.py

@@ -16,6 +16,22 @@ unit_types = ['t2i adapter', 'controlnet', 'xs', 'lite', 'reference', 'ip']
 
 
 class Unit(): # mashup of gradio controls and mapping to actual implementation classes
+    def update_choices(self, model_id=None):
+        name = model_id or self.model_name
+        if name == 'InstantX Union':
+            self.choices = ['canny', 'tile', 'depth', 'blur', 'pose', 'gray', 'lq']
+        elif name == 'Shakker-Labs Union':
+            self.choices = ['canny', 'tile', 'depth', 'blur', 'pose', 'gray', 'lq']
+        elif name == 'Xinsir Union XL':
+            self.choices = ['openpose', 'depth', 'scribble', 'canny', 'normal']
+        elif name == 'Xinsir ProMax XL':
+            self.choices = ['openpose', 'depth', 'scribble', 'canny', 'normal', 'segment', 'tile', 'repaint']
+        else:
+            self.choices = ['default']
+
+    def __str__(self):
+        return f'Unit: type={self.type} enabled={self.enabled} strength={self.strength} start={self.start} end={self.end} mode={self.mode} tile={self.tile}'
+
     def __init__(self,
                  # values
                  index: int = None,
@@ -38,6 +54,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                  control_start = None,
                  control_end = None,
                  control_mode = None,
+                 control_tile = None,
                  result_txt = None,
                  extra_controls: list = [],
         ):
@@ -70,6 +87,10 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
         self.fidelity = 0.5
         self.query_weight = 1.0
         self.adain_weight = 1.0
+        # control mode
+        self.choices = ['default']
+        # control tile
+        self.tile = '1x1'
 
         def reset():
             if self.process is not None:
@@ -92,10 +113,16 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
             self.end = max(start, end)
 
         def control_mode_change(mode):
-            self.mode = mode - 1 if mode > 0 else None
+            self.mode = self.choices.index(mode) if mode is not None and mode in self.choices else 0
 
-        def control_mode_show(model_id):
-            return gr.update(visible='union' in model_id.lower())
+        def control_tile_change(tile):
+            self.tile = tile
+
+        def control_choices(model_id):
+            self.update_choices(model_id)
+            mode_visible = 'union' in model_id.lower() or 'promax' in model_id.lower()
+            tile_visible = 'union' in model_id.lower() or 'promax' in model_id.lower() or 'tile' in model_id.lower()
+            return [gr.update(visible=mode_visible, choices=self.choices), gr.update(visible=tile_visible)]
 
         def adapter_extra(c1):
             self.factor = c1
@@ -172,7 +199,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                 else:
                     self.controls.append(model_id)
                     model_id.change(fn=self.controlnet.load, inputs=[model_id], outputs=[result_txt], show_progress=True)
-                    model_id.change(fn=control_mode_show, inputs=[model_id], outputs=[control_mode], show_progress=False)
+                    model_id.change(fn=control_choices, inputs=[model_id], outputs=[control_mode, control_tile], show_progress=False)
             if extra_controls is not None and len(extra_controls) > 0:
                 extra_controls[0].change(fn=controlnet_extra, inputs=extra_controls)
         elif self.type == 'xs':
@@ -231,3 +258,6 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
         if control_mode is not None:
             self.controls.append(control_mode)
             control_mode.change(fn=control_mode_change, inputs=[control_mode])
+        if control_tile is not None:
+            self.controls.append(control_tile)
+            control_tile.change(fn=control_tile_change, inputs=[control_tile])

modules/control/units/controlnet.py

@@ -52,17 +52,20 @@ predefined_sdxl = {
     'Depth Mid XL': 'diffusers/controlnet-depth-sdxl-1.0-mid',
     'OpenPose XL': 'thibaud/controlnet-openpose-sdxl-1.0/bin',
     'Xinsir Union XL': 'xinsir/controlnet-union-sdxl-1.0',
+    'Xinsir ProMax XL': 'brad-twinkl/controlnet-union-sdxl-1.0-promax',
     'Xinsir OpenPose XL': 'xinsir/controlnet-openpose-sdxl-1.0',
     'Xinsir Canny XL': 'xinsir/controlnet-canny-sdxl-1.0',
     'Xinsir Depth XL': 'xinsir/controlnet-depth-sdxl-1.0',
     'Xinsir Scribble XL': 'xinsir/controlnet-scribble-sdxl-1.0',
     'Xinsir Anime Painter XL': 'xinsir/anime-painter',
+    'Xinsir Tile XL': 'xinsir/controlnet-tile-sdxl-1.0',
     'NoobAI Canny XL': 'Eugeoter/noob-sdxl-controlnet-canny',
     'NoobAI Lineart Anime XL': 'Eugeoter/noob-sdxl-controlnet-lineart_anime',
     'NoobAI Depth XL': 'Eugeoter/noob-sdxl-controlnet-depth',
     'NoobAI Normal XL': 'Eugeoter/noob-sdxl-controlnet-normal',
     'NoobAI SoftEdge XL': 'Eugeoter/noob-sdxl-controlnet-softedge_hed',
     'NoobAI OpenPose XL': 'einar77/noob-openpose',
+    'TTPlanet Tile Realistic XL': 'Yakonrus/SDXL_Controlnet_Tile_Realistic_v2',
     # 'StabilityAI Canny R128': 'stabilityai/control-lora/control-LoRAs-rank128/control-lora-canny-rank128.safetensors',
     # 'StabilityAI Depth R128': 'stabilityai/control-lora/control-LoRAs-rank128/control-lora-depth-rank128.safetensors',
     # 'StabilityAI Recolor R128': 'stabilityai/control-lora/control-LoRAs-rank128/control-lora-recolor-rank128.safetensors',
@@ -166,30 +169,30 @@ class ControlNet():
         self.model_id = None
 
     def get_class(self, model_id:str=''):
-        import modules.shared
-        if modules.shared.sd_model_type == 'sd':
+        from modules import shared
+        if shared.sd_model_type == 'none':
+            _load = shared.sd_model # trigger a load
+        if shared.sd_model_type == 'sd':
             from diffusers import ControlNetModel as cls # pylint: disable=reimported
             config = 'lllyasviel/control_v11p_sd15_canny'
-        elif modules.shared.sd_model_type == 'sdxl':
-            # TODO ControlNetUnion
-            """
+        elif shared.sd_model_type == 'sdxl':
             if 'union' in model_id.lower():
                 from diffusers import ControlNetUnionModel as cls
                 config = 'xinsir/controlnet-union-sdxl-1.0'
+            elif 'promax' in model_id.lower():
+                from diffusers import ControlNetUnionModel as cls
+                config = 'brad-twinkl/controlnet-union-sdxl-1.0-promax'
             else:
                 from diffusers import ControlNetModel as cls # pylint: disable=reimported # sdxl shares same model class
                 config = 'Eugeoter/noob-sdxl-controlnet-canny'
-            """
-            from diffusers import ControlNetModel as cls # pylint: disable=reimported # sdxl shares same model class
-            config = 'Eugeoter/noob-sdxl-controlnet-canny'
-        elif modules.shared.sd_model_type == 'f1':
+        elif shared.sd_model_type == 'f1':
             from diffusers import FluxControlNetModel as cls
             config = 'InstantX/FLUX.1-dev-Controlnet-Union'
-        elif modules.shared.sd_model_type == 'sd3':
+        elif shared.sd_model_type == 'sd3':
             from diffusers import SD3ControlNetModel as cls
             config = 'InstantX/SD3-Controlnet-Canny'
         else:
-            log.error(f'Control {what}: type={modules.shared.sd_model_type} unsupported model')
+            log.error(f'Control {what}: type={shared.sd_model_type} unsupported model')
             return None, None
         return cls, config
 
@@ -299,7 +302,7 @@ class ControlNetPipeline():
                  controlnet: Union[ControlNetModel, list[ControlNetModel]],
                  pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline, FluxPipeline, StableDiffusion3Pipeline],
                  dtype = None,
-                 p: StableDiffusionProcessingControl = None,
+                 p: StableDiffusionProcessingControl = None, # pylint: disable=unused-argument
                 ):
         t0 = time.time()
         self.orig_pipeline = pipeline
@@ -314,14 +317,11 @@ class ControlNetPipeline():
             return
         elif detect.is_sdxl(pipeline) and len(controlnets) > 0:
             from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetUnionPipeline
-            # TODO ControlNetUnion
-            """
             if controlnet.__class__.__name__ == 'ControlNetUnionModel':
                 cls = StableDiffusionXLControlNetUnionPipeline
+                controlnets = controlnets[0] # using only first one
             else:
                 cls = StableDiffusionXLControlNetPipeline
-            """
-            cls = StableDiffusionXLControlNetPipeline
             self.pipeline = cls(
                 vae=pipeline.vae,
                 text_encoder=pipeline.text_encoder,