fix masking

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log for SD.Next
 
-## Update for 2024-02-13
+## Update for 2024-02-14
 
 - **improvements**:
   - **IP Adapter** major refactor  
@@ -67,6 +67,7 @@
   - control fix resize causing runtime errors  
   - control fix processor override image after processor change  
   - handle pipelines that return dict instead of object  
+  - fix inpaint mask only for diffusers  
   - fix vae dtype mismatch, thanks @Disty0  
   - fix controlnet inpaint mask  
   - fix theme list refresh  

modules/control/run.py

@@ -143,7 +143,8 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
             p.fidelity = u.fidelity
             shared.log.debug('Control Reference unit')
         else:
-            active_process.append(u.process)
+            if u.process.processor_id is not None:
+                active_process.append(u.process)
             shared.log.debug(f'Control process unit: i={num_units} process={u.process.processor_id}')
             active_strength.append(float(u.strength))
     p.ops.append('control')
@@ -347,7 +348,6 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                         p.extra_generation_params["Mask erode"] = masking.opts.mask_erode if masking.opts.mask_erode > 0 else None
                         p.extra_generation_params["Mask dilate"] = masking.opts.mask_dilate if masking.opts.mask_dilate > 0 else None
                         p.extra_generation_params["Mask model"] = masking.opts.model if masking.opts.model is not None else None
-                    if len(active_process) > 0:
                         masked_image = masking.run_mask(input_image=input_image, input_mask=mask, return_type='Masked', invert=p.inpainting_mask_invert==1) if mask is not None else input_image
                     else:
                         masked_image = input_image

modules/masking.py

@@ -376,7 +376,8 @@ def run_mask(input_image: Image.Image, input_mask: Image.Image = None, return_ty
         return None
 
     size = min(input_image.width, input_image.height)
-    debug(f'Mask args: blur={mask_blur} padding={mask_padding}')
+    if mask_blur is not None or mask_padding is not None:
+        debug(f'Mask args legacy: blur={mask_blur} padding={mask_padding}')
     if invert is not None:
         opts.invert = invert
     if mask_blur is not None: # compatibility with old img2img values which uses px values
@@ -396,24 +397,21 @@ def run_mask(input_image: Image.Image, input_mask: Image.Image = None, return_ty
     if opts.mask_erode > 0:
         try:
             kernel = np.ones((int(opts.mask_erode * size / 4) + 1, int(opts.mask_erode * size / 4) + 1), np.uint8)
-            cv2_mask = cv2.erode(mask, kernel, iterations=opts.kernel_iterations) # remove noise
+            mask = cv2.erode(mask, kernel, iterations=opts.kernel_iterations) # remove noise
-            mask = cv2_mask
             debug(f'Mask erode={opts.mask_erode:.3f} kernel={kernel.shape} mask={mask.shape}')
         except Exception as e:
             shared.log.error(f'Mask erode: {e}')
     if opts.mask_dilate > 0:
         try:
             kernel = np.ones((int(opts.mask_dilate * size / 4) + 1, int(opts.mask_dilate * size / 4) + 1), np.uint8)
-            cv2_mask = cv2.dilate(mask, kernel, iterations=opts.kernel_iterations) # expand area
+            mask = cv2.dilate(mask, kernel, iterations=opts.kernel_iterations) # expand area
-            mask = cv2_mask
             debug(f'Mask dilate={opts.mask_dilate:.3f} kernel={kernel.shape} mask={mask.shape}')
         except Exception as e:
             shared.log.error(f'Mask dilate: {e}')
     if opts.mask_blur > 0:
         try:
             sigmax, sigmay = 1 + int(opts.mask_blur * size / 4), 1 + int(opts.mask_blur * size / 4)
-            cv2_mask = cv2.GaussianBlur(mask, (0, 0), sigmaX=sigmax, sigmaY=sigmay) # blur mask
+            mask = cv2.GaussianBlur(mask, (0, 0), sigmaX=sigmax, sigmaY=sigmay) # blur mask
-            mask = cv2_mask
             debug(f'Mask blur={opts.mask_blur:.3f} x={sigmax} y={sigmay} mask={mask.shape}')
         except Exception as e:
             shared.log.error(f'Mask blur: {e}')

modules/processing_class.py

@@ -341,6 +341,11 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
                     np_mask = cv2.GaussianBlur(np_mask, (kernel_size, 1), self.mask_blur)
                     np_mask = cv2.GaussianBlur(np_mask, (1, kernel_size), self.mask_blur)
                     self.image_mask = Image.fromarray(np_mask)
+            elif shared.backend == shared.Backend.DIFFUSERS:
+                if 'control' in self.ops:
+                    self.image_mask = masking.run_mask(input_image=self.init_images, input_mask=self.image_mask, return_type='Grayscale', invert=self.inpainting_mask_invert==1) # blur/padding are handled in masking module
+                else:
+                    self.mask = masking.run_mask(input_image=self.init_images, input_mask=self.image_mask, return_type='Grayscale', invert=self.inpainting_mask_invert==1, mask_blur=self.mask_blur, mask_padding=self.inpaint_full_res_padding) # old img2img
             if self.inpaint_full_res: # mask only inpaint
                 self.mask_for_overlay = self.image_mask
                 mask = self.image_mask.convert('L')

modules/processing_diffusers.py

@@ -8,7 +8,7 @@ import numpy as np
 import torch
 import torchvision.transforms.functional as TF
 import diffusers
-from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, masking, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae
+from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae
 from modules.processing_helpers import resize_init_images, resize_hires, fix_prompts, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps
 from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed
 
@@ -124,21 +124,10 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             }
         elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images' ,[])) > 0:
             p.ops.append('inpaint')
-            if p.task_args.get('mask_image', None) is not None: # provided as override by a control module
-                p.mask = masking.run_mask(input_image=p.init_images, input_mask=p.task_args['mask_image'], return_type='Grayscale', invert=p.inpainting_mask_invert==1)
-            elif getattr(p, 'image_mask', None) is not None: # standard img2img
-                if 'control' in p.ops:
-                    p.mask = masking.run_mask(input_image=p.init_images, input_mask=p.image_mask, return_type='Grayscale', invert=p.inpainting_mask_invert==1) # blur/padding are handled in masking module
-                else:
-                    p.mask = masking.run_mask(input_image=p.init_images, input_mask=p.image_mask, return_type='Grayscale', invert=p.inpainting_mask_invert==1, mask_blur=p.mask_blur, mask_padding=p.inpaint_full_res_padding) # old img2img
-            elif getattr(p, 'mask', None) is not None: # backward compatibility
-                pass
-            else: # fallback
-                p.mask = TF.to_pil_image(torch.ones_like(TF.to_tensor(p.init_images[0]))).convert("L")
             width, height = resize_init_images(p)
             task_args = {
                 'image': p.init_images,
-                'mask_image': p.mask,
+                'mask_image': p.image_mask,
                 'strength': p.denoising_strength,
                 'height': height,
                 'width': width,
@@ -437,6 +426,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         p.extra_generation_params["Sampler Eta"] = shared.opts.scheduler_eta
     try:
         t0 = time.time()
+        img = base_args['mask_image']
+        img.save('/tmp/mask.png')
         output = shared.sd_model(**base_args) # pylint: disable=not-callable
         if isinstance(output, dict):
             output = SimpleNamespace(**output)