fix lllite

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

CHANGELOG.md

@@ -26,6 +26,8 @@
   - Gated access: obfuscate and log token used for access  
   - SDXL refiner workflow  
   - Control: t2i-adapter workflow  
+  - Control: xs-controlnet workflow  
+  - Control: lllite-workflow  
 
 ## Highlights for 2025-04-28
 

modules/control/run.py

@@ -133,7 +133,7 @@ def check_active(p, unit_type, units):
         if u.type != unit_type:
             continue
         num_units += 1
-        debug_log(f'Control unit: i={num_units} type={u.type} enabled={u.enabled}')
+        debug_log(f'Control unit: i={num_units} type={u.type} enabled={u.enabled} cn={u.controlnet} proc={u.process}')
         if not u.enabled:
             if u.controlnet is not None and u.controlnet.model is not None:
                 debug_log(f'Control unit offload: model="{u.controlnet.model_id}" device={devices.cpu}')
@@ -692,7 +692,16 @@ def control_run(state: str = '',
                                 p.task_args['image'] = p.init_images # need to set explicitly for txt2img
                                 del p.init_images
                         if unit_type == 'lite':
-                            p.init_images = [input_image]
+                            if input_type == 0:
+                                shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
+                                shared.sd_model.no_task_switch = True
+                            elif input_type == 1:
+                                p.init_images = [input_image]
+                            elif input_type == 2:
+                                if init_image is None:
+                                    shared.log.warning('Control: separate init image not provided')
+                                    init_image = input_image
+                                p.init_images = [init_image]
                             instance.apply(selected_models, processed_image, control_conditioning)
                     if hasattr(p, 'init_images') and p.init_images is None: # delete empty
                         del p.init_images

modules/control/units/lite.py

@@ -125,6 +125,7 @@ class ControlLLLite():
 class ControlLLitePipeline():
     def __init__(self, pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline]):
         self.pipeline = pipeline
+        # self.pipeline.__class__.__name__ = 'ControlLLLitePipeline'
         self.nets = []
 
     def apply(self, controlnet: Union[ControlNetLLLite, list[ControlNetLLLite]], image, conditioning):

modules/control/units/xs_model.py

@@ -631,6 +631,7 @@ class ControlNetXSModel(ModelMixin, ConfigMixin):
         t_emb = t_emb.to(dtype=sample.dtype)
 
         if self.config.learn_embedding: # pylint: disable=no-member
+            base_model = base_model.to(self.control_model.device)
             ctrl_temb = self.control_model.time_embedding(t_emb, timestep_cond)
             base_temb = base_model.time_embedding(t_emb, timestep_cond)
             interpolation_param = self.config.time_embedding_mix**0.3 # pylint: disable=no-member

modules/processing_args.py

@@ -35,7 +35,7 @@ def task_specific_kwargs(p, model):
             }
     elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0:
         if shared.sd_model_type == 'sdxl' and hasattr(model, 'register_to_config'):
-            if model.__class__.__name__ in ['StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline']:
+            if model.__class__.__name__ in sd_models.i2i_pipes:
                 pass
             else:
                 model.register_to_config(requires_aesthetics_score = False)
@@ -74,7 +74,10 @@ def task_specific_kwargs(p, model):
         }
     elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0:
         if shared.sd_model_type == 'sdxl' and hasattr(model, 'register_to_config'):
-            model.register_to_config(requires_aesthetics_score = False)
+            if model.__class__.__name__ in [sd_models.i2i_pipes]:
+                pass
+            else:
+                model.register_to_config(requires_aesthetics_score = False)
         if p.detailer_enabled:
             p.ops.append('detailer')
         else:

modules/sd_models.py

@@ -31,20 +31,27 @@ debug_process = log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not Non
 diffusers_version = int(diffusers.__version__.split('.')[1])
 checkpoint_tiles = checkpoint_titles # legacy compatibility
 pipe_switch_task_exclude = [
-    'StableDiffusionReferencePipeline',
-    'StableDiffusionAdapterPipeline',
-    'AnimateDiffPipeline',
-    'AnimateDiffSDXLPipeline',
-    'OmniGenPipeline',
-    'StableDiffusion3ControlNetPipeline',
-    'InstantIRPipeline',
-    'FluxFillPipeline',
+    'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline',
     'FluxControlPipeline',
-    'PixelSmithXLPipeline',
-    'PhotoMakerStableDiffusionXLPipeline',
-    'StableDiffusionXLInstantIDPipeline',
+    'FluxFillPipeline',
+    'InstantIRPipeline',
     'LTXConditionPipeline',
-    'StableDiffusionXLAdapterPipeline',
+    'OmniGenPipeline',
+    'PhotoMakerStableDiffusionXLPipeline',
+    'PixelSmithXLPipeline',
+    'StableDiffusion3ControlNetPipeline',
+    'StableDiffusionAdapterPipeline',
+    'StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline',
+    'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline',
+    'StableDiffusionReferencePipeline',
+    'StableDiffusionXLInstantIDPipeline',
+]
+i2i_pipes = [
+    'LEditsPPPipelineStableDiffusion',
+    'LEditsPPPipelineStableDiffusionXL',
+    'OmniGenPipeline',
+    'StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline',
+    'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline',
 ]
 
 
@@ -656,13 +663,6 @@ class DiffusersTaskType(Enum):
 
 def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
     cls = pipe.__class__.__name__
-    i2i_pipes = [
-        'LEditsPPPipelineStableDiffusion',
-        'LEditsPPPipelineStableDiffusionXL',
-        'OmniGenPipeline',
-        'StableDiffusionAdapterPipeline',
-        'StableDiffusionXLAdapterPipeline',
-    ]
     if cls in i2i_pipes: # special case
         return DiffusersTaskType.IMAGE_2_IMAGE
     elif 'ImageToVideo' in cls or cls in ['LTXConditionPipeline', 'StableVideoDiffusionPipeline']: # i2v pipelines
@@ -790,6 +790,9 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE:
         clean_diffuser_pipe(pipe)
 
+    if hasattr(pipe, 'no_task_switch'):
+        del pipe.no_task_switch
+        return pipe
     if get_diffusers_task(pipe) == new_pipe_type:
         return pipe
 

modules/ui_control.py

@@ -359,7 +359,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             controlnetxs_ui_units.append(unit_ui)
                             units.append(unit.Unit(
                                 unit_type = 'xs',
-                                index = 1,
+                                index = i,
                                 enabled = enabled,
                                 result_txt = result_txt,
                                 enabled_cb = enabled_cb,