manually unpack base64 urls

handler.py

@@ -2,9 +2,24 @@ from typing import  Dict, List, Any
 import torch
 from PIL import Image
 from io import BytesIO
-from urllib import request
 from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DDIMScheduler
 
+import base64
+import requests
+from io import BytesIO
+from PIL import Image
+
+def load_image(image_url):
+    if image_url.startswith('data:'):
+        # Decode base64 data_uri
+        image_data = base64.b64decode(image_url.split(',')[1])
+        image = Image.open(BytesIO(image_data))
+    else:
+        # Load standard image url
+        response = requests.get(image_url)
+        image = Image.open(BytesIO(response.content))
+    return image
+
 # set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -36,38 +51,36 @@ class EndpointHandler():
         prompt = data.pop("inputs", data)
         url = data.pop("url", data)
 
-        with request.urlopen(url) as response:
-            data = response.read()
-            init_image = Image.open(data).convert("RGB")
-            init_image = Image.open(url)
-            init_image.thumbnail((512, 512))
+        init_image = load_image(url).convert("RGB")
+        init_image = Image.open(url)
+        init_image.thumbnail((512, 512))
 
 
-            params = data.pop("parameters", data)
+        params = data.pop("parameters", data)
 
-            # hyperparamters
-            num_inference_steps = params.pop("num_inference_steps", 25)
-            guidance_scale = params.pop("guidance_scale", 7.5)
-            negative_prompt = params.pop("negative_prompt", None)
-            prompt = params.pop("prompt", None)
-            height = params.pop("height", None)
-            width = params.pop("width", None)
-            manual_seed = params.pop("manual_seed", -1)
+        # hyperparamters
+        num_inference_steps = params.pop("num_inference_steps", 25)
+        guidance_scale = params.pop("guidance_scale", 7.5)
+        negative_prompt = params.pop("negative_prompt", None)
+        prompt = params.pop("prompt", None)
+        height = params.pop("height", None)
+        width = params.pop("width", None)
+        manual_seed = params.pop("manual_seed", -1)
 
-            out = None
+        out = None
 
-            generator = torch.Generator(device='cuda')
-            generator.manual_seed(manual_seed)
-            # run img2img pipeline
-            out = self.imgPipe(prompt,
-                        image=init_image,
-                        num_inference_steps=num_inference_steps,
-                        guidance_scale=guidance_scale,
-                        num_images_per_prompt=1,
-                        negative_prompt=negative_prompt,
-                        height=height,
-                        width=width
-            )
+        generator = torch.Generator(device='cuda')
+        generator.manual_seed(manual_seed)
+        # run img2img pipeline
+        out = self.imgPipe(prompt,
+                    image=init_image,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=guidance_scale,
+                    num_images_per_prompt=1,
+                    negative_prompt=negative_prompt,
+                    height=height,
+                    width=width
+        )
 
-            # return first generated PIL image
-            return out.images[0]
+        # return first generated PIL image
+        return out.images[0]