modular pipeline prototype

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

modules/modular.py

@@ -0,0 +1,55 @@
+import time
+import diffusers
+from modules import shared
+
+
+modular_map= {
+    'StableDiffusionXLPipeline': 'StableDiffusionXLAutoBlocks',
+    'StableDiffusionXLImg2ImgPipeline': 'StableDiffusionXLAutoBlocks',
+    'FluxPipeline': 'FluxAutoBlocks',
+    'WanPipeline': 'WanAutoBlocks',
+    'QwenImagePipeline': 'QwenImageAutoBlocks',
+    'QwenImageEditPipeline': 'QwenImageEditAutoBlocks',
+}
+
+
+def is_compatible(diffusion_pipeline: diffusers.DiffusionPipeline) -> bool:
+    return diffusion_pipeline.__class__.__name__ in modular_map
+
+
+def convert_to_modular(diffusion_pipeline: diffusers.DiffusionPipeline) -> diffusers.ModularPipeline:
+    modular_pipe = None
+    try:
+        t0 = time.time()
+        modular_cls = modular_map.get(diffusion_pipeline.__class__.__name__, None)
+        if modular_cls is None:
+            raise ValueError(f'unknown: cls={diffusion_pipeline.__class__.__name__}')
+        modular_cls = getattr(diffusers, modular_cls, None)
+        if modular_cls is None:
+            raise ValueError(f'invalid: cls={diffusion_pipeline.__class__.__name__}')
+        modular_blocks = modular_cls()
+        modular_pipe = modular_blocks.init_pipeline()
+        components_dct = {k: v for k, v in diffusion_pipeline.components.items() if v is not None}
+        modular_pipe.update_components(**components_dct, **diffusion_pipeline.parameters)
+        modular_pipe.original_pipe = diffusion_pipeline
+        t1 = time.time()
+        shared.log.debug(f'Modular: from={diffusion_pipeline.__class__.__name__} to={modular_pipe.__class__.__name__} time={t1 - t0:.2f}')
+
+        """
+        for expected_input_param in modular_pipe.blocks.inputs:
+            name = expected_input_param.name
+            default = expected_input_param.default
+            kwargs_type = expected_input_param.kwargs_type
+            shared.log.trace(f'Modular input: name={name} type={kwargs_type} default={default}')
+        """
+
+    except Exception as e:
+        shared.log.error(f'Modular: {e}')
+        raise e
+    return modular_pipe
+
+
+def restore_standard(modular_pipe):
+    if hasattr(modular_pipe, 'original_pipe'):
+        shared.log.debug(f'Modular: from={modular_pipe.__class__.__name__} to={modular_pipe.original_pipe.__class__.__name__}')
+        return modular_pipe.original_pipe

modules/processing_args.py

@@ -18,6 +18,26 @@ debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
 disable_pbar = os.environ.get('SD_DISABLE_PBAR', None) is not None
 
 
+def task_modular_kwargs(p, model):
+    model_cls = model.__class__.__name__ # pylint: disable=unused-variable
+    task_args = {}
+    p.ops.append('modular')
+
+    processing_helpers.resize_init_images(p)
+    task_args['width'] = p.width
+    task_args['height'] = p.height
+    if len(getattr(p, 'init_images', [])) > 0:
+        task_args['image'] = p.init_images
+        task_args['strength'] = p.denoising_strength
+    mask_image = p.task_args.get('image_mask', None) or getattr(p, 'image_mask', None) or getattr(p, 'mask', None)
+    if mask_image is not None:
+        task_args['mask_image'] = mask_image
+
+    if debug_enabled:
+        debug_log(f'Process task specific args: {task_args}')
+    return task_args
+
+
 def task_specific_kwargs(p, model):
     model_cls = model.__class__.__name__
     vae_scale_factor = sd_vae.get_vae_scale_factor(model)
@@ -136,6 +156,16 @@ def task_specific_kwargs(p, model):
     return task_args
 
 
+def get_params(model):
+    if hasattr(model, 'blocks') and hasattr(model.blocks, 'inputs'): # modular pipeline
+        possible = [input_param.name for input_param in model.blocks.inputs]
+        return possible
+    else:
+        signature = inspect.signature(type(model).__call__, follow_wrapped=True)
+        possible = list(signature.parameters)
+        return possible
+
+
 def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:typing.Optional[list]=None, negative_prompts_2:typing.Optional[list]=None, prompt_attention:typing.Optional[str]=None, desc:typing.Optional[str]='', **kwargs):
     t0 = time.time()
     shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
@@ -151,8 +181,8 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
             model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba', disable=disable_pbar)
         else:
             model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba')
-    signature = inspect.signature(type(model).__call__, follow_wrapped=True)
-    possible = list(signature.parameters)
+
+    possible = get_params(model)
 
     if debug_enabled:
         debug_log(f'Process pipeline possible: {possible}')
@@ -357,7 +387,11 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
             args[arg] = kwargs[arg]
 
     # handle task specific args
-    task_kwargs = task_specific_kwargs(p, model)
+    if sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.MODULAR:
+        task_kwargs = task_modular_kwargs(p, model)
+    else:
+        task_kwargs = task_specific_kwargs(p, model)
+
     pipe_args = getattr(p, 'task_args', {})
     model_args = getattr(model, 'task_args', {})
     task_kwargs.update(pipe_args or {})
@@ -407,8 +441,9 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
 
     # handle implicit controlnet
     if 'control_image' in possible and 'control_image' not in args and 'image' in args:
-        debug_log('Process: set control image')
-        args['control_image'] = args['image']
+        if sd_models.get_diffusers_task(model) != sd_models.DiffusersTaskType.MODULAR:
+            debug_log('Process: set control image')
+            args['control_image'] = args['image']
 
     sd_hijack_hypertile.hypertile_set(p, hr=len(getattr(p, 'init_images', [])) > 0)
 

modules/processing_diffusers.py

@@ -96,6 +96,13 @@ def process_pre(p: processing.StableDiffusionProcessing):
     #     sd_models.move_model(shared.sd_model.unet, devices.device)
     # if hasattr(shared.sd_model, 'transformer'):
     #     sd_models.move_model(shared.sd_model.transformer, devices.device)
+
+    from modules import modular
+    if modular.is_compatible(shared.sd_model):
+        modular_pipe = modular.convert_to_modular(shared.sd_model)
+        if modular_pipe is not None:
+            shared.sd_model = modular_pipe
+
     timer.process.record('pre')
 
 

modules/processing_helpers.py

@@ -9,6 +9,7 @@ import cv2
 from PIL import Image
 from blendmodes.blend import blendLayers, BlendType
 from modules import shared, devices, images, sd_models, sd_samplers, sd_vae, sd_hijack_hypertile, processing_vae, timer
+from modules.api import helpers
 
 
 debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -16,6 +17,10 @@ debug_steps = shared.log.trace if os.environ.get('SD_STEPS_DEBUG', None) is not
 debug_steps('Trace: STEPS')
 
 
+def is_modular():
+    return sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.MODULAR
+
+
 def is_txt2img():
     return sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE
 
@@ -278,7 +283,10 @@ def validate_sample(tensor):
 def resize_init_images(p):
     if getattr(p, 'image', None) is not None and getattr(p, 'init_images', None) is None:
         p.init_images = [p.image]
+
     if getattr(p, 'init_images', None) is not None and len(p.init_images) > 0:
+        if isinstance(p.init_images[0], str):
+            p.init_images = [helpers.decode_base64_to_image(i, quiet=True) for i in p.init_images]
         vae_scale_factor = sd_vae.get_vae_scale_factor()
         tgt_width, tgt_height = vae_scale_factor * math.ceil(p.init_images[0].width / vae_scale_factor), vae_scale_factor * math.ceil(p.init_images[0].height / vae_scale_factor)
         if p.init_images[0].size != (tgt_width, tgt_height):
@@ -287,11 +295,17 @@ def resize_init_images(p):
             p.height = tgt_height
             p.width = tgt_width
             sd_hijack_hypertile.hypertile_set(p)
-        if getattr(p, 'mask', None) is not None and p.mask.size != (tgt_width, tgt_height):
+        if getattr(p, 'mask', None) is not None and p.mask is not None and p.mask.size != (tgt_width, tgt_height):
+            if isinstance(p.mask[0], str):
+                p.mask = [helpers.decode_base64_to_image(i, quiet=True) for i in p.mask]
             p.mask = images.resize_image(1, p.mask, tgt_width, tgt_height, upscaler_name=None)
-        if getattr(p, 'init_mask', None) is not None and p.init_mask.size != (tgt_width, tgt_height):
+        if getattr(p, 'init_mask', None) is not None and p.init_mask is not None and p.init_mask.size != (tgt_width, tgt_height):
+            if isinstance(p.init_mask[0], str):
+                p.init_mask = [helpers.decode_base64_to_image(i, quiet=True) for i in p.init_mask]
             p.init_mask = images.resize_image(1, p.init_mask, tgt_width, tgt_height, upscaler_name=None)
-        if getattr(p, 'mask_for_overlay', None) is not None and p.mask_for_overlay.size != (tgt_width, tgt_height):
+        if getattr(p, 'mask_for_overlay', None) is not None and p.mask_for_overlay is not None and p.mask_for_overlay.size != (tgt_width, tgt_height):
+            if isinstance(p.mask_for_overlay, str):
+                p.mask_for_overlay = helpers.decode_base64_to_image(p.mask_for_overlay, quiet=True)
             p.mask_for_overlay = images.resize_image(1, p.mask_for_overlay, tgt_width, tgt_height, upscaler_name=None)
         return tgt_width, tgt_height
     return p.width, p.height
@@ -374,6 +388,8 @@ def calculate_base_steps(p, use_denoise_start, use_refiner_start):
     cls = shared.sd_model.__class__.__name__
     if 'Flex' in cls or 'Kontext' in cls or 'Edit' in cls or 'Wan' in cls:
         steps = p.steps
+    elif is_modular():
+        steps = p.steps
     elif not is_txt2img():
         if cls in sd_models.i2i_pipes:
             steps = p.steps

modules/sd_models.py

@@ -743,6 +743,7 @@ class DiffusersTaskType(Enum):
     IMAGE_2_IMAGE = 2
     INPAINTING = 3
     INSTRUCT = 4
+    MODULAR = 5
 
 
 def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
@@ -753,6 +754,8 @@ def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
         return DiffusersTaskType.IMAGE_2_IMAGE
     elif 'Instruct' in cls:
         return DiffusersTaskType.INSTRUCT
+    elif 'Modular' in cls:
+        return DiffusersTaskType.MODULAR
     elif pipe.__class__ in diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING.values():
         return DiffusersTaskType.IMAGE_2_IMAGE
     elif pipe.__class__ in diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING.values():
@@ -950,6 +953,9 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     if get_diffusers_task(pipe) == new_pipe_type:
         return pipe
 
+    if get_diffusers_task(pipe) == DiffusersTaskType.MODULAR:
+        return pipe
+
     # skip specific pipelines
     cls = pipe.__class__.__name__
     if cls in pipe_switch_task_exclude: