Update handler.py

handler.py

@@ -13,6 +13,89 @@ class EndpointHandler():
         self.pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None)
         self.pipe.to("cuda")
         self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
+
+    def __call__(self, data):
+
+        info=data['inputs']
+        
+        # Image + Logo
+        image=info.pop("image",data)
+        logo = info.pop("logo",data)
+        logo = base64.b64decode(logo)
+        logo = Image.open(BytesIO(logo)).convert('RGB')
+        
+        # Seed
+        seed = info.pop("seed", data)
+
+        # Punchline Text
+        punchline_text_max_width=info.pop("punchline_text_max_width", data.get("punchline_text_max_width", 550))
+        punchline_text_color=info.pop("punchline_text_color", data.get("punchline_text_color", "#008000"))
+        spacing_image_text=info.pop("spacing_image_text", data.get("spacing_image_text", 0))
+        
+        
+        # color code:
+        color_code=info.pop("color_code",data)
+        # inference steps:        
+        num_inference_steps = info.pop("num_inference_steps", data.get("num_inference_steps", 40))
+        # Image guidance scale
+        image_guidance_scale = info.pop("image_guidance_scale", data.get("image_guidance_scale", 1.5))
+        # Guidance scale
+        guidance_scale = info.pop("guidance_scale", data.get("guidance_scale", 7.5))
+        
+        
+        # Button color
+        button_color = info.pop("button_color", data.get("button_color", "#008000"))
+        # Button text
+        button_text = info.pop("button_text", data.get("button_text"))
+        # Button font
+        button_font = info.pop("button_font", data.get("button_font", cv2.FONT_HERSHEY_TRIPLEX))
+        # Button font scale
+        button_font_scale = info.pop("button_font_scale", data.get("button_font_scale", 0.75))
+        button_font_thickness = info.pop("button_font_thickness", data.get("button_font_thickness", 1))
+        button_text_color = info.pop("button_text_color", data.get("button_text_color", "#FFFFFF"))
+        spacing_between_punchline_and_button = info.pop("spacing_between_punchline_and_button", data.get("spacing_between_punchline_and_button", 10))
+        
+        # prompt
+        text_prompt=info.pop("prompt",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))
+        
+        result_prompt, negative_prompt = self.build_prompt(text_prompt, color_code)
+        
+        torch.manual_seed(seed)
+        images = self.pipe(result_prompt, negative_prompt = negative_prompt, image=raw_images, num_inference_steps=num_inference_steps, guidance_scale = guidance_scale, image_guidance_scale = image_guidance_scale).images
+        img=images[0]
+        img.save("./1.png")
+        logo.save("./logo.png")
+        
+        resulting_template = self.create_image_template(
+            base_image_path="./1.png",
+            logo_path="./logo.png",
+            punchline_text=punchline_text,
+            punchline_text_color=punchline_text_color,
+            punchline_text_max_width=punchline_text_max_width,
+            spacing_image_text=spacing_image_text,
+            button_color=button_color,
+            button_text=button_text,
+            button_font=button_font,
+            button_font_scale=button_font_scale,
+            button_font_thickness=button_font_thickness,
+            button_text_color=button_text_color,
+            spacing_between_punchline_and_button=spacing_between_punchline_and_button,
+            corner_radius=30)
+        
+        
+
+        resulting_template.save("./result.png")
+        
+        with open('./result.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)
@@ -48,6 +131,7 @@ class EndpointHandler():
         return mapped_color
 
         # Helpers.
+        
     def draw_text(img, text, font=cv2.FONT_HERSHEY_PLAIN, pos=(20, 45), font_scale=1, font_thickness=1, text_color=(0, 0, 255)):
         x, y = pos
         text_size, _ = cv2.getTextSize(text, font, font_scale, font_thickness)
@@ -247,21 +331,18 @@ class EndpointHandler():
     def create_image_template(
         base_image_path,
         logo_path,
-        
         punchline_text,
         punchline_text_color="#008000",
         punchline_text_max_width=550,
         spacing_image_text=0,
-        
         button_color="#008000",
         button_text="Call Action Text Here! >",
         button_font=cv2.FONT_HERSHEY_TRIPLEX,
         button_font_scale=0.7,
         button_font_thickness=1,
         button_text_color="#FFFFFF",
-        
         spacing_between_punchline_and_button=10,
-        corner_radius=30,):
+        corner_radius=30):
     
 
 
@@ -303,88 +384,7 @@ class EndpointHandler():
         return result_template
 
 
-    def __call__(self, data):
-
-        info=data['inputs']
-        
-        # Image + Logo
-        image=info.pop("image",data)
-        logo = info.pop("logo",data)
-        logo = base64.b64decode(logo)
-        logo = Image.open(BytesIO(logo)).convert('RGB')
-        
-        # Seed
-        seed = info.pop("seed", data)
-
-        # Punchline Text
-        punchline_text_max_width=info.pop("punchline_text_max_width", data.get("punchline_text_max_width", 550))
-        punchline_text_color=info.pop("punchline_text_color", data.get("punchline_text_color", "#008000"))
-        spacing_image_text=info.pop("spacing_image_text", data.get("spacing_image_text", 0))
-        
-        
-        # color code:
-        color_code=info.pop("color_code",data)
-        # inference steps:        
-        num_inference_steps = info.pop("num_inference_steps", data.get("num_inference_steps", 40))
-        # Image guidance scale
-        image_guidance_scale = info.pop("image_guidance_scale", data.get("image_guidance_scale", 1.5))
-        # Guidance scale
-        guidance_scale = info.pop("guidance_scale", data.get("guidance_scale", 7.5))
-        
-        
-        # Button color
-        button_color = info.pop("button_color", data.get("button_color", "#008000"))
-        # Button text
-        button_text = info.pop("button_text", data.get("button_text"))
-        # Button font
-        button_font = info.pop("button_font", data.get("button_font", cv2.FONT_HERSHEY_TRIPLEX))
-        # Button font scale
-        button_font_scale = info.pop("button_font_scale", data.get("button_font_scale", 0.75))
-        button_font_thickness = info.pop("button_font_thickness", data.get("button_font_thickness", 1))
-        button_text_color = info.pop("button_text_color", data.get("button_text_color", "#FFFFFF"))
-        spacing_between_punchline_and_button = info.pop("spacing_between_punchline_and_button", data.get("spacing_between_punchline_and_button", 10))
-        
-        # prompt
-        text_prompt=info.pop("prompt",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))
-        
-        result_prompt, negative_prompt = self.build_prompt(text_prompt, color_code)
-        
-        torch.manual_seed(seed)
-        images = self.pipe(result_prompt, negative_prompt = negative_prompt, image=raw_images, num_inference_steps=num_inference_steps, guidance_scale = guidance_scale, image_guidance_scale = image_guidance_scale).images
-        img=images[0]
-        img.save("./1.png")
-        logo.save("./logo.png")
-        
-        resulting_template = create_image_template(
-            base_image_path="./1.png",
-            logo_path="./logo.png",
-            punchline_text=punchline_text,
-            punchline_text_color=punchline_text_color,
-            punchline_text_max_width=punchline_text_max_width,
-            spacing_image_text=spacing_image_text,
-            button_color=button_color,
-            button_text=button_text,
-            button_font=button_font,
-            button_font_scale=button_font_scale,
-            button_font_thickness=button_font_thickness,
-            button_text_color=button_text_color,
-            spacing_between_punchline_and_button=spacing_between_punchline_and_button,
-            corner_radius=30)
-        
-        
-
-        resulting_template.save("./result.png")
-        
-        with open('./result.png','rb') as img_file:
-            encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
-        return {'image':encoded_string}