import torch from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler import base64 from io import BytesIO from PIL import ImageDraw, ImageOps, Image, ImageFont from scipy.spatial import KDTree from webcolors import CSS3_HEX_TO_NAMES, hex_to_rgb, hex_to_name import cv2 import numpy as np class EndpointHandler(): def __init__(self, path=""): model_id = "timbrooks/instruct-pix2pix" 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=info.pop("punchline_text", data.get("punchline_text", "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) #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) 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 = f'{color_name} shadows, worst quality, low quality, low res, blurry, watermark, cropped, jpeg artifacts, error, sketch ,duplicate, ugly, monochrome, horror, 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) color_mapping = { 'tomato': 'red', 'chocolate': 'brown-black', 'darkgoldenrod': 'yellow', } color_name = names[index] mapped_color = color_mapping.get(color_name, color_name) return mapped_color # Helpers. def draw_text(self, img, text, font=cv2.FONT_HERSHEY_TRIPLEX, 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) text_w, text_h = text_size cv2.putText(img, text, (x, y + text_h), font, font_scale, text_color, font_thickness) return text_size def smooth_corners(self, image, radius, alpha=255): factor = 5 # Factor to increase the image size for anti-aliasing radius = radius * factor size = (image.width, image.height) # Create the rounded rectangle mask image with a white background mask = Image.new('RGBA', (size[0] * factor, size[1] * factor), (255, 255, 255, 0)) # Create a corner image corner = Image.new('RGBA', (radius, radius), (0, 0, 0, 0)) draw = ImageDraw.Draw(corner) draw.pieslice((0, 0, radius * 2, radius * 2), 180, 270, fill=(0, 0, 0, alpha + 55)) # Paste and rotate the corners as needed mx, my = (size[0] * factor, size[1] * factor) mask.paste(corner, (0, 0), corner) mask.paste(corner.rotate(90), (0, my - radius), corner.rotate(90)) mask.paste(corner.rotate(180), (mx - radius, my - radius), corner.rotate(180)) mask.paste(corner.rotate(270), (mx - radius, 0), corner.rotate(270)) # Draw the inner rectangles draw = ImageDraw.Draw(mask) draw.rectangle([(radius, 0), (mx - radius, my)], fill=(0, 0, 0, alpha)) draw.rectangle([(0, radius), (mx, my - radius)], fill=(0, 0, 0, alpha)) # Resize the mask to the original image size with anti-aliasing mask = mask.resize(size, Image.ANTIALIAS) # Use the mask to composite the original image with smoothed corners result_image = Image.new('RGBA', size) result_image.paste(image, (0, 0), mask) return result_image def add_logo_to_image(self, base_image_path, logo_path, corner_radius=30): base_image = Image.open(base_image_path) smoothed_image = self.smooth_corners(base_image, radius=30) # Create a white background with the same size as the original image white_background = Image.new("RGB", base_image.size, "white") # Paste the smoothed_image over the white background result_image = Image.new("RGBA", base_image.size) result_image.paste(white_background, (0, 0)) result_image.paste(smoothed_image, (0, 0), smoothed_image) logo_path = logo_path # Load the logo image logo = Image.open(logo_path) # Calculate the desired logo height desired_height = base_image.height / 3 # Calculate the scaling factor to achieve the desired height scaling_factor = desired_height / logo.height # Calculate the new width while maintaining the aspect ratio new_width = int(logo.width * scaling_factor) # Resize the logo image resized_logo = logo.resize((new_width, int(desired_height))) # Get the dimensions of the original image image_width, image_height = base_image.size # Create a new white background image with the same width as the original image white_background = Image.new("RGB", (image_width, resized_logo.height), "white") # Calculate the position to paste the logo in the middle x_logo = (image_width - resized_logo.width) // 2 y_logo = 0 # Paste the logo onto the white background white_background.paste(resized_logo, (x_logo, y_logo)) # Create an empty image with the same width as the original image and height enough to accommodate both logo and base_image result_logo_image = Image.new("RGB", (image_width, image_height + white_background.height)) # Paste the base_image below the logo result_logo_image.paste(white_background, (0, 0)) result_logo_image.paste(result_image, (0, white_background.height)) # Save or display the final image return result_logo_image def create_template(self, logo_and_image, punchline_text, punchline_text_color, punchline_text_max_width, spacing_image_text): result_logo_image_resized = logo_and_image.resize((325, 425)) # Create a white background with a size of 645x645 whitespace = Image.new("RGB", (645, 645), "white") # Calculate the position to paste the result_logo_image in the middle x_logo = (whitespace.width - result_logo_image_resized.width) // 2 y_logo = 10 whitespace.paste(result_logo_image_resized, (x_logo, y_logo)) # Define the text, position, font, and color font = cv2.FONT_HERSHEY_TRIPLEX font_scale = 1 font_thickness = 2 text_size = cv2.getTextSize(punchline_text, font, font_scale, font_thickness)[0] # Reduce spacing between the image and the first line y_logo += spacing_image_text x = (whitespace.size[1] - punchline_text_max_width) // 2 y = y_logo + 2 * text_size[1] + result_logo_image_resized.height lines = [] # List to store lines of text # Split the text into words words = punchline_text.split() current_line = "" for word in words: # Calculate the width of the current line current_line_size = cv2.getTextSize(current_line + " " + word, font, font_scale, font_thickness)[0] if current_line_size[0] > punchline_text_max_width: # Start a new line if the current line exceeds the maximum width lines.append(current_line) current_line = word else: if current_line: current_line += " " current_line += word # Append the last linea lines.append(current_line) w_1 = whitespace.size[1] # Iterate through the lines and add them to the image for line in lines: text_size = cv2.getTextSize(line, font, font_scale, font_thickness)[0] x = (w_1 - text_size[0]) // 2 whitespace = np.asarray(whitespace) # Convert the position to a tuple of integers position = (int(x), int(y)) whitespace = cv2.putText(whitespace, line, position, font, font_scale, punchline_text_color, thickness=2) y += 1.5 * text_size[1] # Adjust the vertical position for the next line punchline_text_height = len(lines) * text_size[1] # Define the desired vertical spacing between the punchline text and the button template = Image.fromarray(whitespace) return template, punchline_text_height, len(lines) def draw_button_with_text(self, button_color, button_text, button_font=cv2.FONT_HERSHEY_TRIPLEX, button_font_scale = 0.7, button_font_thickness=1, button_text_color=(0, 0, 255)): # Create a blank 300x80 white image (background) background_image = Image.new("RGB", (300, 80), "white") button_image_pil = self.smooth_corners(background_image, 20, alpha=255) button_image_pil = ImageOps.colorize(button_image_pil.convert("L"), "white", button_color) button_width, button_height = button_image_pil.size text_size, _ = cv2.getTextSize(button_text, button_font, button_font_scale, button_font_thickness) text_w, text_h = text_size text_x = (button_width - text_w) // 2 if text_w <= button_width else 0 text_y = (button_height - text_h) // 2 button_image_np = np.asarray(button_image_pil) text_size = self.draw_text(button_image_np, button_text, button_font, (text_x, text_y), button_font_scale, button_font_thickness, button_text_color) # Convert the NumPy array back to a PIL image button_image_pil = Image.fromarray(button_image_np) return button_image_pil def concat_template_with_button(self, template, button_image, line_count, spacing_between_punchline_and_button, punchline_text_height): result_template = template.copy() # Calculate the position for the button button_width, button_height = button_image.size button_x = (template.width - button_width) // 2 # Adjust the button width as needed line_count = line_count - 1 if line_count > 1 else 1 button_y = 425 + 10 + punchline_text_height + line_count* 33 + spacing_between_punchline_and_button # Paste the button image onto the template result_template.paste(button_image, (button_x, int(button_y))) return result_template def create_image_template(self, 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): punchline_text_color = tuple(int(punchline_text_color[i:i+2], 16) for i in (1, 3, 5)) button_color = tuple(int(button_color[i:i+2], 16) for i in (1, 3, 5)) button_text_color = tuple(int(button_text_color[i:i+2], 16) for i in (1, 3, 5)) result_image = self.add_logo_to_image(base_image_path, logo_path, corner_radius=corner_radius) # Create the image template with the punchline text template, punchline_text_height, line_count = self.create_template( result_image, punchline_text, punchline_text_color, punchline_text_max_width, spacing_image_text, ) # Create the button image button_image = self.draw_button_with_text( button_color, button_text, button_font, button_font_scale, button_font_thickness, button_text_color, ) # Concatenate the template with the button result_template = self.concat_template_with_button( template, button_image, line_count, spacing_between_punchline_and_button, punchline_text_height ) return result_template if __name__=="__main__": my_handler=EndpointHandler(path='.')