Update handler.py

handler.py

@@ -35,9 +35,12 @@ class EndpointHandler:
         image_data = data.get("inputs")
         if image_data is None:
             return [{"error": "Missing 'inputs' key"}]
+        # Wake up function to restart the server
+        elif image_data == "wake": 
+            return [{"status": "woken"}]
 
         # Generation args
-        args = data.get("generation_args")
+        args = data.get("generation_args", {})
         generation_args = self.default_args.copy()
         for k in self.default_args:
             if k in args and args[k] is not None:
@@ -46,7 +49,7 @@ class EndpointHandler:
         # Decode base64 image
         try:
             image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
-            image.thumbnail([768,768], Image.LANCZOS)
+            image.thumbnail([512,512], Image.LANCZOS)
         except Exception as e:
             return [{"error": f"Image decoding failed: {str(e)}"}]