Update handler.py

handler.py

@@ -3,7 +3,7 @@ from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDisc
 import base64
 from io import BytesIO
 from PIL import Image
-import json
+from webcolors import hex_to_name, CSS3_HEX_TO_NAMES
 
 class EndpointHandler():
     def __init__(self, path=""):
@@ -16,16 +16,51 @@ class EndpointHandler():
 
         info=data['inputs']
         image=info.pop("image",data)
-        prompt=info.pop("text",data)
+
+        # color code:
+        color_code=info.pop("color_code",data)
+        
+        # prompt
+        text_prompt=info.pop("text",data)
+
+        # image
         image=base64.b64decode(image)
         raw_images = Image.open(BytesIO(image)).convert('RGB')
+        raw_images = raw_images.convert("RGB")
+        raw_images = raw_images.resize((512, 512))
         
-        images = self.pipe(prompt, image=raw_images, num_inference_steps=25, image_guidance_scale=1).images
+        result_prompt, negative_prompt = self.build_prompt(text_prompt, color_code)
+
+        images = self.pipe(result_prompt, negative_prompt = negative_prompt, image=raw_images, num_inference_steps=25, image_guidance_scale=1).images
         img=images[0]
         img.save("./1.png")
         with open('./1.png','rb') as img_file:
             encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
         return {'image':encoded_string}
+
+
+    def build_prompt(self, text_prompt, color_code):
+        color_name = self.hex_to_name(color_code)
+        q_prompt = "No blur, high quality, no distortion over objects, text remains clear and undistorted. Do not modify or alter the logo."
+        base_prompts = text_prompt + q_prompt                 
+        coloring_prompt = f" with a {color_name} color applied to only the designated area or key element in the picture by avoiding it becoming the dominant color of the image, leaving the text, logo, and shadows untouched."
+        result_prompt = f"{base_prompts}{coloring_prompt}"
+        negative_prompt = 'blue shadows, worst quality, low quality, low res, blurry, watermark, banner, extra digits, cropped, jpeg artifacts, signature, username, error, sketch ,duplicate, ugly, monochrome, horror, geometry, mutation, disgusting'
+
+        return result_prompt, negative_prompt 
+
+    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 = []
+
+        for color_hex, color_name in CSS3_HEX_TO_NAMES.items():
+            names.append(color_name)
+            rgb_values.append(hex_to_rgb(color_hex))
+
+        kdt_db = KDTree(rgb_values)
+        distance, index = kdt_db.query(rgb_tuple)
+        return names[index]