diff --git a/modules/modular.py b/modules/modular.py new file mode 100644 index 000000000..bdcd1c470 --- /dev/null +++ b/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 diff --git a/modules/processing_args.py b/modules/processing_args.py index 971adad44..7770f210d 100644 --- a/modules/processing_args.py +++ b/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) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 4de4e0a47..75f74a11b 100644 --- a/modules/processing_diffusers.py +++ b/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') diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 15ee19822..7906261b9 100644 --- a/modules/processing_helpers.py +++ b/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 diff --git a/modules/sd_models.py b/modules/sd_models.py index 4423c780d..9b9ec7d68 100644 --- a/modules/sd_models.py +++ b/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: