diff --git a/modules/control/run.py b/modules/control/run.py index 29e6099b8..e44968ca5 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -75,8 +75,11 @@ def set_pipe(p, has_models, unit_type, selected_models, active_model, active_str p.task_args['control_guidance_start'] = control_guidance_start p.task_args['control_guidance_end'] = control_guidance_end p.task_args['guess_mode'] = p.guess_mode - instance = controlnet.ControlNetPipeline(selected_models, shared.sd_model, p=p) - pipe = instance.pipeline + if 'Flex' not in shared.sd_model.__class__.__name__: + instance = controlnet.ControlNetPipeline(selected_models, shared.sd_model, p=p) + pipe = instance.pipeline + else: + pipe = shared.sd_model elif unit_type == 'xs' and has_models: p.extra_generation_params["Control type"] = 'ControlNet-XS' p.controlnet_conditioning_scale = control_conditioning @@ -189,7 +192,11 @@ def check_enabled(p, unit_type, units, active_model, active_strength, active_sta control_conditioning = None control_guidance_start = None control_guidance_end = None - if unit_type == 't2i adapter' or unit_type == 'controlnet' or unit_type == 'xs' or unit_type == 'lite': + if 'Flex' in p.sd_model.__class__.__name__: + has_models = True + selected_models = [None] + p.guess_mode = False + elif unit_type == 't2i adapter' or unit_type == 'controlnet' or unit_type == 'xs' or unit_type == 'lite': if len(active_model) == 0: selected_models = None elif len(active_model) == 1: @@ -609,18 +616,30 @@ def control_run(state: str = '', p.task_args['strength'] = p.denoising_strength p.init_images = None elif input_type == 1: # Init image same as control - if 'control_image' in possible: + p.init_images = [p.override or input_image] * max(1, len(active_model)) + if 'inpaint_image' in possible: # flex + p.task_args['inpaint_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images + p.task_args['inpaint_mask'] = Image.new('L', p.task_args['inpaint_image'].size, int(p.denoising_strength * 255)) + p.task_args['control_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images + p.task_args['width'] = p.width + p.task_args['height'] = p.height + elif 'control_image' in possible: p.task_args['control_image'] = p.init_images # switch image and control_image if 'control_mode' in possible: p.task_args['control_mode'] = getattr(p, 'control_mode', None) if 'strength' in possible: p.task_args['strength'] = p.denoising_strength - p.init_images = [p.override or input_image] * len(active_model) elif input_type == 2: # Separate init image if init_image is None: shared.log.warning('Control: separate init image not provided') init_image = input_image - if 'control_image' in possible: + if 'inpaint_image' in possible: # flex + p.task_args['inpaint_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images + p.task_args['inpaint_mask'] = Image.new('L', p.task_args['inpaint_image'].size, int(p.denoising_strength * 255)) + p.task_args['control_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images + p.task_args['width'] = p.width + p.task_args['height'] = p.height + elif 'control_image' in possible: p.task_args['control_image'] = p.init_images # switch image and control_image if 'control_mode' in possible: p.task_args['control_mode'] = getattr(p, 'control_mode', None) diff --git a/modules/model_flex.py b/modules/model_flex.py index 721ac785a..9ef5f44e1 100644 --- a/modules/model_flex.py +++ b/modules/model_flex.py @@ -81,6 +81,10 @@ def load_flex(checkpoint_info, diffusers_load_config={}): **load_args, ) sd_hijack_te.init_hijack(pipe) + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flex2"] = Flex2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["fluxcfgzero"] = Flex2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["fluxcfgzero"] = Flex2Pipeline + del text_encoder_2 del transformer diff --git a/modules/processing_args.py b/modules/processing_args.py index a7bb4831e..863a77fb4 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -284,7 +284,14 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t args["prior_guidance_scale"] = p.cfg_scale if 'decoder_guidance_scale' in possible: args["decoder_guidance_scale"] = p.image_cfg_scale - + if 'Flex' in model.__class__.__name__: + if p.init_images is not None and len(p.init_images) > 0: + args['inpaint_image'] = p.init_images[0] if isinstance(p.init_images, list) else p.init_images + args['inpaint_mask'] = Image.new('L', args['inpaint_image'].size, 1) + args['control_image'] = args['inpaint_image'].convert('L').convert('RGB') # will be interpreted as depth + args['control_strength'] = p.denoising_strength + args['width'] = p.width + args['height'] = p.height # set callbacks if 'prior_callback_steps' in possible: # Wuerstchen / Cascade args['prior_callback_steps'] = 1 diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index 599a77534..db26ef328 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -9,11 +9,13 @@ from modules import shared, devices, processing_correction, extra_networks, time p = None debug = os.environ.get('SD_CALLBACK_DEBUG', None) is not None debug_callback = shared.log.trace if debug else lambda *args, **kwargs: None +warned = False def set_callbacks_p(processing): - global p # pylint: disable=global-statement + global p, warned # pylint: disable=global-statement p = processing + warned = False def prompt_callback(step, kwargs): @@ -129,7 +131,10 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {} except Exception: pass except Exception as e: - shared.log.error(f'Callback: {e}') + global warned # pylint: disable=global-statement + if not warned: + shared.log.error(f'Callback: {e}') + warned = True # from modules import errors # errors.display(e, 'Callback') if shared.cmd_opts.profile and shared.profiler is not None: diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 71e31f663..7d417525b 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -456,6 +456,8 @@ def calculate_base_steps(p, use_denoise_start, use_refiner_start): if not is_txt2img(): if use_denoise_start and shared.sd_model_type == 'sdxl': steps = p.steps // (1 - p.refiner_start) + elif 'Flex' in shared.sd_model.__class__.__name__: + steps = p.steps elif shared.sd_model_type == 'omnigen': steps = p.steps elif p.denoising_strength > 0: