fix xyz with control enabled

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

CHANGELOG.md

@@ -104,6 +104,7 @@ Although upgrades and existing installations are tested and should work fine!
   - fix api `/sdapi/v1/embeddings` endpoint  
   - fix incorrect reporting of deleted and modified files  
   - fix SD3.x loader and TAESD preview  
+  - fix xyz with control enabled  
   - allow upscaling with models that have implicit VAE processing  
   - sdnq use inference context during quantization  
   - framepack improve offloading  

modules/control/processors.py

@@ -133,6 +133,9 @@ class Processor():
         if processor_id is not None:
             self.load()
 
+    def __str__(self):
+        return f' Processor(id={self.processor_id} model={self.model.__class__.__name__})' if self.processor_id and self.model else ''
+
     def reset(self, processor_id: str = None):
         if self.model is not None:
             debug(f'Control Processor unloaded: id="{self.processor_id}"')

modules/control/run.py

@@ -33,7 +33,7 @@ def restore_pipeline():
     global pipe, instance # pylint: disable=global-statement
     if instance is not None and hasattr(instance, 'restore'):
         instance.restore()
-    if original_pipeline is not None and (original_pipeline.__class__.__name__ != shared.sd_model.__class__.__name__):
+    if (original_pipeline is not None) and (original_pipeline.__class__.__name__ != shared.sd_model.__class__.__name__):
         debug_log(f'Control restored pipeline: class={shared.sd_model.__class__.__name__} to={original_pipeline.__class__.__name__}')
         shared.sd_model = original_pipeline
     pipe = None
@@ -52,6 +52,7 @@ def is_unified_model():
 
 
 def set_pipe(p, has_models, unit_type, selected_models, active_model, active_strength, control_conditioning, control_guidance_start, control_guidance_end, inits):
+    print('HERE SET PIPE')
     global pipe, instance # pylint: disable=global-statement
     pipe = None
     if has_models:
@@ -123,6 +124,7 @@ def set_pipe(p, has_models, unit_type, selected_models, active_model, active_str
 
 
 def check_active(p, unit_type, units):
+    print('HERE CHECK ACTIVE')
     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
@@ -192,6 +194,7 @@ def check_active(p, unit_type, units):
 
 
 def check_enabled(p, unit_type, units, active_model, active_strength, active_start, active_end):
+    print('HERE CHECK ENABLED')
     has_models = False
     selected_models: List[Union[controlnet.ControlNetModel, xs.ControlNetXSModel, t2iadapter.AdapterModel]] = None
     control_conditioning = None
@@ -229,6 +232,23 @@ def control_set(kwargs):
         p_extra_args[k] = v
 
 
+def init_units(units: List[unit.Unit]):
+    print('HERE INIT UNITS')
+    for u in units:
+        if not u.enabled:
+            continue
+        if u.process_name is not None and u.process_name != '' and u.process_name != 'None':
+            u.process.load(u.process_name, force=False)
+        if u.model_name is not None and u.model_name != '' and u.model_name != 'None':
+            if u.type == 't2i adapter':
+                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
+
+
 def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
                 units: List[unit.Unit] = [], inputs: List[Image.Image] = [], inits: List[Image.Image] = [], mask: Image.Image = None, unit_type: str = None, is_generator: bool = True,
                 input_type: int = 0,
@@ -249,23 +269,10 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
                 video_skip_frames: int = 0, video_type: str = 'None', video_duration: float = 2.0, video_loop: bool = False, video_pad: int = 0, video_interpolate: int = 0,
                 *input_script_args,
         ):
-    # handle optional initialization via ui
-    for u in units:
-        if not u.enabled:
-            continue
-        if u.process_name is not None and u.process_name != '' and u.process_name != 'None':
-            u.process.load(u.process_name, force=False)
-        if u.model_name is not None and u.model_name != '' and u.model_name != 'None':
-            if u.type == 't2i adapter':
-                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
-
     global pipe, original_pipeline # pylint: disable=global-statement
+
     debug_log(f'Control: type={unit_type} input={inputs} init={inits} type={input_type}')
+    init_units(units)
     if inputs is None or (type(inputs) is list and len(inputs) == 0):
         inputs = [None]
     output_images: List[Image.Image] = [] # output images

modules/control/unit.py

@@ -30,7 +30,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
             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}'
+        return f'Unit(index={self.index} enabled={self.enabled} type="{self.type}" strength={self.strength} start={self.start} end={self.end}{self.process}{self.controlnet})'
 
     def __init__(self,
                  # values

modules/control/units/controlnet.py

@@ -189,6 +189,9 @@ class ControlNet():
         if model_id is not None:
             self.load()
 
+    def __str__(self):
+        return f' ControlNet(id={self.model_id} model={self.model.__class__.__name__})' if self.model_id and self.model else ''
+
     def reset(self):
         if self.model is not None:
             debug_log(f'Control {what} model unloaded')

modules/control/units/lite.py

@@ -74,6 +74,9 @@ class ControlLLLite():
         if model_id is not None:
             self.load()
 
+    def __str__(self):
+        return f' ControlLLLite(id={self.model_id} model={self.model.__class__.__name__})' if self.model_id and self.model else ''
+
     def reset(self):
         if self.model is not None:
             debug(f'Control {what} model unloaded')

modules/control/units/t2iadapter.py

@@ -82,6 +82,9 @@ class Adapter():
         if model_id is not None:
             self.load()
 
+    def __str__(self):
+        return f' T2IAdapter(id={self.model_id} model={self.model.__class__.__name__})' if self.model_id and self.model else ''
+
     def reset(self):
         if self.model is not None:
             debug(f'Control {what} model unloaded')

modules/control/units/xs.py

@@ -70,6 +70,9 @@ class ControlNetXS():
         if model_id is not None:
             self.load()
 
+    def __str__(self):
+        return f' ControlNetXS(id={self.model_id} model={self.model.__class__.__name__})' if self.model_id and self.model else ''
+
     def reset(self):
         if self.model is not None:
             debug(f'Control {what} model unloaded')

modules/processing.py

@@ -414,9 +414,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 
             devices.torch_gc()
 
-        if hasattr(shared.sd_model, 'restore_pipeline') and shared.sd_model.restore_pipeline is not None:
-            shared.sd_model.restore_pipeline()
-        shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
+        if not p.xyz:
+            if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None):
+                shared.sd_model.restore_pipeline()
+            shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
 
         t1 = time.time()
 

modules/processing_class.py

@@ -115,6 +115,8 @@ class StableDiffusionProcessing:
                  outpath_grids=None,
                  do_not_save_samples: bool = False,
                  do_not_save_grid: bool = False,
+                 # xyz flag
+                 xyz: bool = False,
                  # scripts
                  script_args: list = [],
                  # overrides

modules/processing_diffusers.py

@@ -170,7 +170,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
         prev_job = shared.state.job
 
         # hires runs on original pipeline
-        if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None) and not shared.opts.control_hires:
+        if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None) and (not shared.opts.control_hires):
             shared.sd_model.restore_pipeline()
 
         # upscale

modules/sd_samplers.py

@@ -89,7 +89,6 @@ def create_sampler(name, model):
         return restore_default(model)
     sampler = config.constructor(model)
     if sampler.sampler is None:
-        print('HERE')
         return restore_default(model)
     is_flow = ('FlowMatch' in sampler.sampler.__class__.__name__) or (getattr(sampler.sampler.config, 'prediction_type', None) == 'flow_prediction')
 

scripts/xyz_grid.py

@@ -368,9 +368,12 @@ class Script(scripts_manager.Script):
                 include_text=include_text,
             )
 
+        if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None):
+            shared.sd_model.restore_pipeline()
+
         if not processed.images:
             return processed # something broke, no further handling needed.
-        # processed.images = (1)*grid + (z > 1 ? z : 0)*subgrids + (x*y*z)*images
+
         have_grid = 1 if include_grid else 0
         have_subgrids = len(zs) if len(zs) > 1 and include_subgrids else 0
         have_images = processed.images[have_grid+have_subgrids:]

scripts/xyz_grid_on.py

@@ -388,10 +388,13 @@ class Script(scripts_manager.Script):
                 include_text=include_text,
             )
 
+        if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None):
+            shared.sd_model.restore_pipeline()
+
         if not processed.images:
             active = False
             return processed # something broke, no further handling needed.
-        # processed.images = (1)*grid + (z > 1 ? z : 0)*subgrids + (x*y*z)*images
+
         have_grid = 1 if include_grid else 0
         have_subgrids = len(zs) if len(zs) > 1 and include_subgrids else 0
         have_images = processed.images[have_grid+have_subgrids:]