update

handler.py

@@ -2,16 +2,14 @@ from typing import Dict, Any
 import torch
 from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
 from PIL import Image
-import torch
-from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
-
-model_id = "timbrooks/instruct-pix2pix"
+from webcolors import CSS3_HEX_TO_NAMES, hex_to_rgb
+from scipy.spatial import KDTree
 
 class EndpointHandler:
     
-    def __init__(self):
+    def __init__(self, path=""):
         model_id = "timbrooks/instruct-pix2pix"
-        self.model = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None).to("cuda")
+        self.model = StableDiffusionInstructPix2PixPipeline.from_pretrained(path, torch_dtype=torch.float16, safety_checker=None).to("cuda")
         self.model.scheduler = EulerAncestralDiscreteScheduler.from_config(self.model.scheduler.config)
 
     def __call__(self, data: Dict[str, Any]) -> Image:
@@ -20,6 +18,7 @@ class EndpointHandler:
             data (dict):
                 The payload with the input image, text prompt, color code, and optional parameters.
         """
+        
         input_image = self.load_process_image(data['image_path'])
         text_prompt = data['text_prompt']
         color_code = data['color_code']
@@ -36,10 +35,9 @@ class EndpointHandler:
                 num_inference_steps=50,
                 guidance_scale=guidance_scale,
                 image_guidance_scale=image_guidance_scale,
-                
             ).images
             
-        return  images
+        return images
 
     def load_process_image(self, image_path):
         image = Image.open(image_path)
@@ -53,8 +51,8 @@ class EndpointHandler:
         result_prompt = f"{text_prompt}{coloring_prompt}"
         return result_prompt
 
-    def hex_to_name(hex_color):
-        rgb_tuple = tuple(int(a[i:i+2], 16) for i in (1, 3, 5))
+    def hex_to_name(self, hex_color):
+        rgb_tuple = tuple(int(hex_color[i:i+2], 16) for i in (1, 3, 5))
         names = []
         rgb_values = []
 
@@ -64,4 +62,4 @@ class EndpointHandler:
 
         kdt_db = KDTree(rgb_values)
         distance, index = kdt_db.query(rgb_tuple)
-        return names[index]
+        return names[index]