mask sanitization

hydra_node/sanitize.py

@@ -134,6 +134,38 @@ def sanitize_stable_diffusion(request, config):
             traceback.print_exc()
             return False, "Error while opening and cleaning image"
 
+    if request.masks is not None:
+        masks = request.masks
+        images = []
+        for x in range(len(masks)):
+            image = masks[x]["mask"]
+            try:
+                image = base64.b64decode(image.encode('utf-8'))
+
+            except Exception as e:
+                traceback.print_exc()
+                return False, "image is not valid base64"
+
+            try:
+                from PIL import Image
+                image = Image.open(BytesIO(image))
+                image.verify()
+
+            except Exception as e:
+                traceback.print_exc()
+                return False, "image is not valid"
+
+            #image is valid, load it again(still check again, verify() can't be sure as it doesn't decode.)
+            try:
+                image = Image.open(BytesIO(image))
+                #image = image.convert('RGB')
+                image = image.resize((request.width//request.downsampling_factor, request.height//request.downsampling_factor), resample=Image.Resampling.LANCZOS)
+            except Exception as e:
+                traceback.print_exc()
+                return False, "Error while opening and cleaning image"
+            
+            masks[x]["mask"] = image
+
     return True, request
 
 def sanitize_dalle_mini(request):