make base64 images work

hydra_node/models.py

@@ -99,12 +99,6 @@ def decode_image(image, model):
     image = custom_to_pil(image)
     return image
 
-def sanitize_image(image):
-    #Open image with PIL and get rid of alpha channel, scale to given res with center crop
-    image = Image.open(image)
-    image = image.convert('RGB')
-    return image
-
 class StableInterface(nn.Module):
     def __init__(self, model, thresholder = None):
         super().__init__()
@@ -186,9 +180,7 @@ class StableDiffusionModel(nn.Module):
         if request.image is not None:
             request.sampler = "ddim_img2img" #enforce ddim for now
             self.ddim.make_schedule(ddim_num_steps=request.steps, ddim_eta=request.ddim_eta, verbose=False)
-            image = sanitize_image(request.image)
-            image = image.resize((request.width, request.height), resample=Image.Resampling.LANCZOS)
-            start_code = encode_image(image, self.model.first_stage_model).to(self.device)
+            start_code = encode_image(request.image, self.model.first_stage_model).to(self.device)
             start_code = self.model.get_first_stage_encoding(start_code)
             print(start_code.shape)
             start_code = start_code + (torch.randn_like(start_code) * request.noise)

hydra_node/sanitize.py

 from dotmap import DotMap
 import math
+from io import BytesIO
 
 v1pp_defaults = {
     'steps': 50,
@@ -91,6 +92,24 @@ def sanitize_stable_diffusion(request):
     if request.sampler not in samplers:
         return False, "sampler should be one of {}".format(samplers)
 
+    if request.image is not None:
+        #decode from base64
+        request.image = request.image.decode('base64')
+        #check if image is valid
+        try:
+            from PIL import Image
+            image = Image.open(BytesIO(request.image))
+            image.verify()
+
+        except Exception as e:
+            return False, "image is not valid"
+        
+        #image is valid, load it again
+        image = Image.open(BytesIO(request.image))
+        image = image.convert('RGB')
+        image = image.resize((request.width, request.height), resample=Image.Resampling.LANCZOS)
+        request.image = image
+
     return True, request
 
 def sanitize_dalle_mini(request):
@@ -111,6 +130,6 @@ def sanitize_input(config, request):
 
     if config.model_name == 'stable-diffusion':
         return sanitize_stable_diffusion(request)
-        
+
     elif config.model_name == 'dalle-mini':
         return sanitize_dalle_mini(request)
\ No newline at end of file