Create handler.py

handler.py

@@ -0,0 +1,82 @@
+import torch
+import numpy as np
+from PIL import Image
+from transformers import AutoTokenizer
+from pulid.pipeline_v1_1 import PuLIDPipeline
+from pulid.utils import resize_numpy_image_long
+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
+        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):
+        # 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]
+
+        # Set seed if needed
+        if seed == -1:
+            seed = torch.Generator(device="cpu").seed()
+
+        # Handle the ortho settings
+        if ortho == 'v2':
+            self.attention.ORTHO = False
+            self.attention.ORTHO_v2 = True
+        elif ortho == 'v1':
+            self.attention.ORTHO = True
+            self.attention.ORTHO_v2 = False
+        else:
+            self.attention.ORTHO = False
+            self.attention.ORTHO_v2 = False
+
+        # Process the images
+        if id_image is not None:
+            id_image = resize_numpy_image_long(id_image, 1024)
+            supp_id_image_list = [
+                resize_numpy_image_long(supp_id_image, 1024) for supp_id_image in supp_images if supp_id_image is not None
+            ]
+            id_image_list = [id_image] + supp_id_image_list
+            uncond_id_embedding, id_embedding = self.pipeline.get_id_embedding(id_image_list)
+        else:
+            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, input_data):
+        # 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)
+
+        # 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)