Address 2 bugs

1. missing mask sketch
2. excessive Img2Img load time

app.py

@@ -423,15 +423,21 @@ def generate_image_lowmem(
         mask_parameters = {}
         # Load the mask image if provided
         if (pipeline_name == "FluxFillPipeline"):
-            mask_image = open_image(mask_image).convert("RGBA")
-            mask_condition_type = constants.condition_type[5]
-            guidance_scale = 30
-            num_inference_steps=50
-            max_sequence_length=512
-            print(f"\nAdded mask image.\n {mask_image.size}")
-            mask_parameters ={
-                "mask_image": mask_image,
-            }
+            try:
+                mask_image = open_image(mask_image).convert("RGBA")
+                mask_condition_type = constants.condition_type[5]
+                guidance_scale = 30
+                num_inference_steps=50
+                max_sequence_length=512
+                print(f"\nAdded mask image.\n {mask_image.size}")
+                mask_parameters ={
+                    "mask_image": mask_image,
+                }
+            except Exception as e:
+                print(f"Error loading mask image: {e}")
+                mask_image = None
+                gr.Warning("Please sketch a mask image to use the Fill model.")
+                raise Exception("Please sketch a mask image to use the Fill model.")
 
         # Set the random seed for reproducibility
         generator = torch.Generator(device=device).manual_seed(seed)

utils/image_utils.py

@@ -61,6 +61,8 @@ def get_image_from_dict(image_path):
             image_path = image_path.get('composite')
         elif 'image' in image_path:
             image_path = image_path.get('image')
+        elif 'background' in image_path:
+            image_path = image_path.get('background')
         else:
             print("\n Unknown image dictionary.\n")
             raise UserWarning("Unknown image dictionary.")