add include mask in output, improve model offload compatibility

modules/control/run.py

@@ -482,6 +482,8 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                         debug(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)} class={pipe.__class__}')
                         debug(f'Control exec pipeline: p={vars(p)}')
                         debug(f'Control exec pipeline: args={p.task_args} image={p.task_args.get("image", None)} control={p.task_args.get("control_image", None)} mask={p.task_args.get("mask_image", None) or p.image_mask} ref={p.task_args.get("ref_image", None)}')
+                        if sd_models.get_diffusers_task(pipe) != sd_models.DiffusersTaskType.TEXT_2_IMAGE: # force vae back to gpu if not in txt2img mode
+                            sd_models.move_model(pipe.vae, devices.device)
                         p.scripts = scripts.scripts_control
                         p.script_args = input_script_args
                         processed = p.scripts.run(p, *input_script_args)
@@ -508,6 +510,10 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                                 output_image = images.resize_image(resize_mode_after, output_image, width_after, height_after, resize_name_after)
 
                             output_images.append(output_image)
+                            if shared.opts.include_mask:
+                                if processed_image is not None and isinstance(processed_image, Image.Image):
+                                    output_images.append(processed_image)
+
                             if is_generator:
                                 image_txt = f'{output_image.width}x{output_image.height}' if output_image is not None else 'None'
                                 if video is not None:

modules/processing.py

@@ -396,6 +396,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     if not p.disable_extra_networks:
         extra_networks.deactivate(p, extra_network_data)
 
+    if shared.opts.include_mask:
+        if getattr(p, 'image_mask', None) is not None and isinstance(p.image_mask, Image.Image):
+            output_images.append(p.image_mask)
+
     processed = Processed(
         p,
         images_list=output_images,

modules/processing_diffusers.py

@@ -345,7 +345,6 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 p.task_args['sag_scale'] = p.sag_scale
             else:
                 shared.log.warning(f'SAG incompatible scheduler: current={sd_model.scheduler.__class__.__name__} supported={supported}')
-
         if shared.opts.cuda_compile_backend == "olive-ai":
             sd_model = olive_check_parameters_changed(p, is_refiner_enabled())
         if sd_model.__class__.__name__ == "OnnxRawPipeline":
@@ -362,12 +361,6 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         shared.sd_model = orig_pipeline
         return results
 
-    if shared.opts.diffusers_move_base:
-        sd_models.move_model(shared.sd_model, devices.device)
-
-    # recompile if a parameter changes
-    sd_models_compile.openvino_recompile_model(p, hires=False, refiner=False)
-
     # pipeline type is set earlier in processing, but check for sanity
     is_control = getattr(p, 'is_control', False) is True
     has_images = len(getattr(p, 'init_images' ,[])) > 0
@@ -378,6 +371,9 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         if len(getattr(p, 'init_images' ,[])) == 0:
             p.init_images = [TF.to_pil_image(torch.rand((3, getattr(p, 'height', 512), getattr(p, 'width', 512))))]
 
+    sd_models.move_model(shared.sd_model, devices.device)
+    sd_models_compile.openvino_recompile_model(p, hires=False, refiner=False) # recompile if a parameter changes
+
     use_refiner_start = is_txt2img() and is_refiner_enabled() and not p.is_hr_pass and p.refiner_start > 0 and p.refiner_start < 1
     use_denoise_start = not is_txt2img() and p.refiner_start > 0 and p.refiner_start < 1
 

modules/sd_models.py

@@ -742,6 +742,12 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model'):
 
 def move_model(model, device=None, force=False):
     if model is not None:
+        if getattr(model, 'vae', None) is not None and get_diffusers_task(model) != DiffusersTaskType.TEXT_2_IMAGE:
+            if device == devices.device: # force vae back to gpu if not in txt2img mode
+                model.vae.to(device)
+                if hasattr(model.vae, '_hf_hook'):
+                    debug_move(f'Model move: to={device} class={model.vae.__class__} function={sys._getframe(1).f_code.co_name}') # pylint: disable=protected-access
+                    model.vae._hf_hook.execution_device = device # pylint: disable=protected-access
         if getattr(model, 'has_accelerate', False) and not force:
             return
         debug_move(f'Model move: to={device} class={model.__class__} function={sys._getframe(1).f_code.co_name}') # pylint: disable=protected-access

modules/shared.py

@@ -507,6 +507,7 @@ options_templates.update(options_section(('saving-images', "Image Options"), {
     "img_max_size_mp": OptionInfo(250, "Maximum image size (MP)", gr.Slider, {"minimum": 100, "maximum": 2000, "step": 1}),
     "webp_lossless": OptionInfo(False, "WebP lossless compression"),
     "save_selected_only": OptionInfo(True, "Save only saves selected image"),
+    "include_mask": OptionInfo(False, "Include mask in outputs"),
     "samples_save_zip": OptionInfo(True, "Create ZIP archive"),
 
     "image_sep_metadata": OptionInfo("<h2>Metadata/Logging</h2>", "", gr.HTML),