moved to new arible statefulschema

handler.py

@@ -52,12 +52,12 @@ def load_models():
         "lllyasviel/control_v11f1e_sd15_tile",
         torch_dtype=torch.float16,
     )
-    
+
     controlnet_brightness = ControlNetModel.from_pretrained(
         "ioclab/control_v1p_sd15_brightness",
         torch_dtype=torch.float16,
     )
-    
+
     pipe = StableDiffusionControlNetPipeline.from_pretrained(
         MODEL_ID,
         controlnet=[
@@ -68,12 +68,12 @@ def load_models():
         cache_dir="cache",
         # local_files_only=True,
     )
-    
+
     pipe.to("cuda")
 
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     pipe.enable_vae_slicing()
-   
+
     return pipe
 
 
@@ -130,7 +130,20 @@ def generate_image_with_conditioning_scale(**inputs):
         image=[qr_image] * 2,
     ).images
 
-    return [{"data": image_to_base64(image), "format": "png"} for image in images]
+    # return [{"data": image_to_base64(image), "format": "png"} for image in images]
+
+    return {
+        "fields": [
+            {
+                "name": "output",
+                "type": "Image",
+                "value": [
+                    f"data:image/png;base64,{image_to_base64(image)}"
+                    for image in images
+                ],
+            }
+        ]
+    }
 
 
 def generate_image(pipe, inputs):
@@ -161,6 +174,7 @@ class EndpointHandler:
     def __init__(self, path=""):
         self._model = load_models()
 
-    def __call__(self, model_input: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, model_input: Dict[str, Any]):
         images = generate_image(self._model, model_input)
+        yield {"fields": [{"name": "progress", "type": "Progress"}]}
         return images