Update handler.py

handler.py

@@ -8,30 +8,29 @@ from pulid import attention_processor as attention
 
 torch.set_grad_enabled(False)
 
-# Initialize the model and tokenizer
-class ModelHandler:
-    def __init__(self):
-        # Set default model parameters
+class EndpointHandler:
+    def __init__(self, model_dir=None):
+        # Initialize the model and tokenizer
         self.pipeline = PuLIDPipeline(sdxl_repo='RunDiffusion/Juggernaut-XL-v9', sampler='dpmpp_sde')
         self.default_cfg = 7.0
         self.default_steps = 25
         self.attention = attention
         self.pipeline.debug_img_list = []
 
-    def preprocess(self, input_data):
+    def preprocess(self, inputs):
         # Extracts image and parameters from the input data
-        id_image = input_data[0]
-        supp_images = input_data[1:4]
-        prompt = input_data[4]
-        neg_prompt = input_data[5]
-        scale = input_data[6]
-        seed = int(input_data[7])
-        steps = int(input_data[8])
-        H = int(input_data[9])
-        W = int(input_data[10])
-        id_scale = input_data[11]
-        num_zero = int(input_data[12])
-        ortho = input_data[13]
+        id_image = inputs[0]
+        supp_images = inputs[1:4]
+        prompt = inputs[4]
+        neg_prompt = inputs[5]
+        scale = inputs[6]
+        seed = int(inputs[7])
+        steps = int(inputs[8])
+        H = int(inputs[9])
+        W = int(inputs[10])
+        id_scale = inputs[11]
+        num_zero = int(inputs[12])
+        ortho = inputs[13]
 
         # Set seed if needed
         if seed == -1:
@@ -62,21 +61,17 @@ class ModelHandler:
 
         return (prompt, neg_prompt, scale, seed, steps, H, W, id_scale, num_zero, uncond_id_embedding, id_embedding)
 
-    def predict(self, input_data):
+    def predict(self, inputs):
         # Preprocess the input data
-        (prompt, neg_prompt, scale, seed, steps, H, W, id_scale, num_zero, uncond_id_embedding, id_embedding) = self.preprocess(input_data)
+        (prompt, neg_prompt, scale, seed, steps, H, W, id_scale, num_zero, uncond_id_embedding, id_embedding) = self.preprocess(inputs)
 
         # Run the inference pipeline
         img = self.pipeline.inference(
             prompt, (1, H, W), neg_prompt, id_embedding, uncond_id_embedding, id_scale, scale, steps, seed
         )[0]
 
-        return np.array(img), str(seed), self.pipeline.debug_img_list
-
-
-# Instantiate the model handler
-handler = ModelHandler()
-
-def handler_function(input_data):
-    # Predict using the handler
-    return handler.predict(input_data)
+        return {
+            "image": np.array(img).tolist(),
+            "seed": str(seed),
+            "debug_images": [np.array(debug_img).tolist() for debug_img in self.pipeline.debug_img_list],
+        }