Update handler.py

handler.py

@@ -1,24 +1,42 @@
 import torch
 import numpy as np
-from PIL import Image
-from transformers import AutoTokenizer
+from typing import Dict, Any
 from pulid.pipeline_v1_1 import PuLIDPipeline
 from pulid.utils import resize_numpy_image_long
 from pulid import attention_processor as attention
 
+# Disable gradients for inference
 torch.set_grad_enabled(False)
 
 class EndpointHandler:
-    def __init__(self, model_dir=None):
-        # Initialize the model and tokenizer
+    def __init__(self, model_dir: str = None):
+        """
+        Initializes the model and necessary components.
+
+        Args:
+            model_dir (str): Directory containing the model weights.
+        """
         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, inputs):
-        # Extracts image and parameters from the input data
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Handles inference requests.
+
+        Args:
+            data (Dict[str, Any]): Input data for inference.
+
+        Returns:
+            Dict[str, Any]: Results containing the generated image and debug information.
+        """
+        # Preprocess inputs
+        inputs = data.get("inputs", [])
+        if not inputs or len(inputs) < 14:
+            raise ValueError("Invalid inputs. Expected 14 elements in the input list.")
+
         id_image = inputs[0]
         supp_images = inputs[1:4]
         prompt = inputs[4]
@@ -36,7 +54,7 @@ class EndpointHandler:
         if seed == -1:
             seed = torch.Generator(device="cpu").seed()
 
-        # Handle the ortho settings
+        # Handle orthogonal settings
         if ortho == 'v2':
             self.attention.ORTHO = False
             self.attention.ORTHO_v2 = True
@@ -47,7 +65,7 @@ class EndpointHandler:
             self.attention.ORTHO = False
             self.attention.ORTHO_v2 = False
 
-        # Process the images
+        # Process images
         if id_image is not None:
             id_image = resize_numpy_image_long(id_image, 1024)
             supp_id_image_list = [
@@ -59,19 +77,14 @@ class EndpointHandler:
             uncond_id_embedding = None
             id_embedding = None
 
-        return (prompt, neg_prompt, scale, seed, steps, H, W, id_scale, num_zero, uncond_id_embedding, id_embedding)
-
-    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(inputs)
-
-        # Run the inference pipeline
+        # Generate image
         img = self.pipeline.inference(
             prompt, (1, H, W), neg_prompt, id_embedding, uncond_id_embedding, id_scale, scale, steps, seed
         )[0]
 
+        # Prepare response
         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],
-        }
+        }