tracing and control improvements

modules/control/adapters.py

@@ -1,12 +1,13 @@
 import os
 import time
-from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, StableDiffusionAdapterPipeline, StableDiffusionXLAdapterPipeline
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, MultiAdapter, StableDiffusionAdapterPipeline, StableDiffusionXLAdapterPipeline
 from modules.shared import log
 from modules import errors
 
 
 what = 'T2I-Adapter'
-debug = log.debug if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug = log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: CONTROL')
 predefined_sd15 = {
     'Canny': 'TencentARC/t2iadapter_canny_sd15v2',
     'Depth': 'TencentARC/t2iadapter_depth_sd15v2',
@@ -37,11 +38,12 @@ def list_models(refresh=False):
     models = {}
     if modules.shared.sd_model_type == 'none':
         models = ['None']
-    if modules.shared.sd_model_type == 'sdxl':
+    elif modules.shared.sd_model_type == 'sdxl':
         models = ['None'] + sorted(predefined_sdxl)
-    if modules.shared.sd_model_type == 'sd':
+    elif modules.shared.sd_model_type == 'sd':
         models = ['None'] + sorted(predefined_sd15)
     else:
+        log.warning(f'Control {what} model list failed: unknown model type')
         models = ['None'] + sorted(list(predefined_sd15) + list(predefined_sdxl))
     debug(f'Control list {what}: path={cache_dir} models={models}')
     return models
@@ -102,6 +104,8 @@ class AdapterPipeline():
         if pipeline is None:
             log.error(f'Control {what} pipeline: model not loaded')
             return
+        # if isinstance(adapter, list) and len(adapter) > 1: # TODO use MultiAdapter
+        #    adapter = MultiAdapter(adapter)
         if isinstance(pipeline, StableDiffusionXLPipeline):
             self.pipeline = StableDiffusionXLAdapterPipeline(
                 vae=pipeline.vae,
@@ -111,7 +115,7 @@ class AdapterPipeline():
                 tokenizer_2=pipeline.tokenizer_2,
                 unet=pipeline.unet,
                 scheduler=pipeline.scheduler,
-                adapter=adapter, # can be a list
+                adapter=adapter,
             ).to(pipeline.device)
         elif isinstance(pipeline, StableDiffusionPipeline):
             self.pipeline = StableDiffusionAdapterPipeline(
@@ -123,7 +127,7 @@ class AdapterPipeline():
                 requires_safety_checker=False,
                 safety_checker=None,
                 feature_extractor=None,
-                adapter=adapter, # can be a list
+                adapter=adapter,
             ).to(pipeline.device)
         else:
             log.error(f'Control {what} pipeline: class={pipeline.__class__.__name__} unsupported model type')

modules/control/controlnets.py

@@ -8,7 +8,8 @@ from modules import errors
 
 
 what = 'ControlNet'
-debug = log.debug if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug = log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: CONTROL')
 predefined_sd15 = {
     'OpenPose': "lllyasviel/control_v11p_sd15_openpose",
     'Canny': "lllyasviel/control_v11p_sd15_canny",
@@ -65,7 +66,7 @@ def list_models(refresh=False):
     elif modules.shared.sd_model_type == 'sd':
         models = ['None'] + sorted(predefined_sd15) + sorted(find_models())
     else:
-        log.error('Control model list failed: unknown model type')
+        log.warning(f'Control {what} model list failed: unknown model type')
         models = ['None'] + sorted(predefined_sd15) + sorted(predefined_sdxl) + sorted(find_models())
     debug(f'Control list {what}: path={cache_dir} models={models}')
     return models

modules/control/controlnetsxs.py

@@ -15,7 +15,8 @@ except Exception:
 
 
 what = 'ControlNet-XS'
-debug = log.debug if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug = log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: CONTROL')
 predefined_sd15 = {
 }
 predefined_sdxl = {

modules/control/processors.py

@@ -26,7 +26,8 @@ from modules.control.proc.zoe import ZoeDetector
 
 models = {}
 cache_dir = 'models/control/processors'
-debug = log.debug if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug = log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: CONTROL')
 config = {
     # pose models
     'OpenPose': {'class': OpenposeDetector, 'checkpoint': True, 'params': {'include_body': True, 'include_hand': False, 'include_face': False}},
@@ -115,9 +116,11 @@ class Processor():
         self.load_config = { 'cache_dir': cache_dir }
         from_config = config.get(processor_id, {}).get('load_config', None)
         if load_config is not None:
-            self.load_config.update(load_config)
+            for k, v in load_config.items():
+                self.load_config[k] = v
         if from_config is not None:
-            self.load_config.update(from_config)
+            for k, v in from_config.items():
+                self.load_config[k] = v
         if processor_id is not None:
             self.load()
 
@@ -137,9 +140,11 @@ class Processor():
                 return
             from_config = config.get(processor_id, {}).get('load_config', None)
             if from_config is not None:
-                self.load_config.update(from_config)
+                for k, v in from_config.items():
+                    self.load_config[k] = v
             cls = config[processor_id]['class']
             log.debug(f'Control processor loading: id="{processor_id}" class={cls.__name__}')
+            debug(f'Control processor config={self.load_config}')
             if 'DWPose' in processor_id:
                 det_ckpt = 'https://download.openmmlab.com/mmdetection/v2.0/yolox/yolox_l_8x8_300e_coco/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth'
                 if 'Tiny' == config['DWPose']['model']:

modules/control/run.py

@@ -16,7 +16,8 @@ from modules.control import reference # ControlNet-Reference
 from modules import devices, shared, errors, processing, images, sd_models, sd_samplers
 
 
-debug = shared.log.debug if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
+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
 
@@ -192,7 +193,7 @@ def control_run(units: List[unit.Unit], inputs, unit_type: str, is_generator: bo
         pass
     shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) # reset current pipeline
 
-    if not has_models and (unit_type == 'reference' or unit_type == 'controlnet' or unit_type == 'xs'): # run in img2img mode
+    if not has_models and (unit_type == 'reference' or unit_type == 'adapter' or unit_type == 'controlnet' or unit_type == 'xs'): # 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
@@ -229,7 +230,7 @@ def control_run(units: List[unit.Unit], inputs, unit_type: str, is_generator: bo
         instance = reference.ReferencePipeline(shared.sd_model)
         pipe = instance.pipeline
     else:
-        shared.log.error('Control: unknown unit type')
+        shared.log.error(f'Control: unknown unit type: {unit_type}')
         pipe = None
     debug(f'Control pipeline: class={pipe.__class__} args={vars(p)}')
     t1, t2, t3 = time.time(), 0, 0
@@ -352,7 +353,7 @@ def control_run(units: List[unit.Unit], inputs, unit_type: str, is_generator: bo
                     # pipeline
                     output = None
                     if pipe is not None: # run new pipeline
-                        if not has_models and (unit_type == 'reference' or unit_type == 'controlnet'): # run in img2img mode
+                        if not has_models and (unit_type == 'reference' or 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
@@ -365,6 +366,7 @@ def control_run(units: List[unit.Unit], inputs, unit_type: str, is_generator: bo
                             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__} args={vars(p)}')
                         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
@@ -375,22 +377,22 @@ def control_run(units: List[unit.Unit], inputs, unit_type: str, is_generator: bo
                     # 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)
 
-                        # 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
+                            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()

modules/control/test.py

@@ -11,6 +11,8 @@ def test_processors(image):
     from PIL import ImageDraw, ImageFont
     images = []
     for processor_id in processors.list_models():
+        if shared.state.interrupted:
+            continue
         shared.log.info(f'Testing processor: {processor_id}')
         processor = processors.Processor(processor_id)
         if processor is None:

modules/generation_parameters_copypaste.py

@@ -17,6 +17,7 @@ type_of_gr_update = type(gr.update())
 paste_fields = {}
 registered_param_bindings = []
 debug = shared.log.trace if os.environ.get('SD_PASTE_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: PASTE')
 
 
 class ParamBinding:

modules/img2img.py

@@ -10,6 +10,7 @@ from modules.memstats import memory_stats
 
 
 debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: PROCESS')
 
 
 def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args):

modules/paths.py

@@ -33,6 +33,7 @@ sd_default_config = os.path.join(sd_configs_path, "v1-inference.yaml")
 sd_model_file = cli.ckpt or os.path.join(script_path, 'model.ckpt') # not used
 default_sd_model_file = sd_model_file # not used
 debug = log.trace if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: PATH')
 paths = {}
 
 if os.environ.get('SD_PATH_DEBUG', None) is not None:

modules/processing.py

@@ -45,6 +45,7 @@ from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile
 opt_C = 4
 opt_f = 8
 debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: PROCESS')
 
 
 def setup_color_correction(image):

modules/processing_correction.py

@@ -9,6 +9,7 @@ from modules import shared
 
 
 debug = shared.log.trace if os.environ.get('SD_HDR_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: HDR')
 
 
 def soft_clamp_tensor(tensor, threshold=0.8, boundary=4):

modules/prompt_parser.py

@@ -71,6 +71,7 @@ re_attention_v1 = re.compile(r"""
 
 debug_output = os.environ.get('SD_PROMPT_DEBUG', None)
 debug = log.trace if debug_output is not None else lambda *args, **kwargs: None
+debug('Trace: PROMPT')
 
 
 def get_learned_conditioning_prompt_schedules(prompts, steps):

modules/prompt_parser_diffusers.py

@@ -7,7 +7,7 @@ from compel.embeddings_provider import BaseTextualInversionManager, EmbeddingsPr
 from modules import shared, prompt_parser, devices
 
 debug = shared.log.trace if os.environ.get('SD_PROMPT_DEBUG', None) is not None else lambda *args, **kwargs: None
-
+debug('Trace: PROMPT')
 CLIP_SKIP_MAPPING = {
     None: ReturnedEmbeddingsType.LAST_HIDDEN_STATES_NORMALIZED,
     1: ReturnedEmbeddingsType.LAST_HIDDEN_STATES_NORMALIZED,

modules/sd_samplers.py

@@ -4,6 +4,7 @@ from modules.sd_samplers_common import samples_to_image_grid, sample_to_image #
 
 
 debug = shared.log.trace if os.environ.get('SD_SAMPLER_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: SAMPLER')
 all_samplers = []
 all_samplers = []
 all_samplers_map = {}

modules/sd_samplers_diffusers.py

@@ -5,6 +5,7 @@ from modules import sd_samplers_common
 
 
 debug = shared.log.trace if os.environ.get('SD_SAMPLER_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: SAMPLER')
 
 try:
     from diffusers import (

modules/txt2img.py

@@ -6,6 +6,7 @@ from modules.ui import plaintext_to_html
 
 
 debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: PROCESS')
 
 
 def txt2img(id_task,

modules/ui_control.py

@@ -15,6 +15,8 @@ max_units = 10
 units: list[unit.Unit] = [] # main state variable
 input_source = None
 debug = os.environ.get('SD_CONTROL_DEBUG', None) is not None
+if debug:
+    shared.log.trace('Control debug enabled')
 
 
 def initialize():

modules/ui_extra_networks.py

@@ -25,6 +25,7 @@ dir_cache = {} # key=path, value=(mtime, listdir(path))
 refresh_time = 0
 extra_pages = shared.extra_networks
 debug = shared.log.trace if os.environ.get('SD_EN_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug('Trace: EN')
 card_full = '''
     <div class='card' onclick={card_click} title='{name}' data-tab='{tabname}' data-page='{page}' data-name='{name}' data-filename='{filename}' data-tags='{tags}' data-mtime='{mtime}' data-size='{size}' data-search='{search}'>
         <div class='overlay'>

webui.py

@@ -317,6 +317,7 @@ def webui(restart=False):
             shared.log.debug(f'Registered callbacks: {k}={len(v)} {[c.script for c in v]}')
     log.info(f"Startup time: {timer.startup.summary()}")
     debug = log.trace if os.environ.get('SD_SCRIPT_DEBUG', None) is not None else lambda *args, **kwargs: None
+    debug('Trace: SCRIPTS')
     debug('Loaded scripts:')
     for m in modules.scripts.scripts_data:
         debug(f'  {m}')