Update handler.py

handler.py

@@ -4,7 +4,7 @@ from PIL import Image
 from io import BytesIO
 from typing import Dict, List, Any
 
-class EndpointHandler:
+class EndpointHandler():
     def __init__(self, path=""):
         self.model_base = "AIhgenerator/nsfwxxl2"
         self.v_autoencoder = "madebyollin/sdxl-vae-fp16-fix"
@@ -21,7 +21,7 @@ class EndpointHandler:
             add_watermarker=False,
         )
         self.pipe.safety_checker = None
-        # self.pipe.to("cuda")
+        self.pipe.to("cuda")
 
         # # Load the refiner pipeline
         # self.pipe_refiner = DiffusionPipeline.from_pretrained(
@@ -32,9 +32,9 @@ class EndpointHandler:
         # )
         # self.pipe_refiner.enable_model_cpu_offload()
 
-    def __call__(self, inputs):
-        # print("height",height)
-        # inputs = data.pop("inputs", data)
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        print("data",data)
+        inputs = data.pop("inputs", data)
 
         prompt, prompt2, negative_prompt, negative_prompt2 = inputs['prompt'], inputs['prompt2'], inputs['negative_prompt'], inputs['negative_prompt2']
         print(prompt, prompt2, negative_prompt, negative_prompt2)