import logging import math from pathlib import Path import re from time import time from typing import List, Optional, Tuple from PIL import Image, ImageDraw, ImageFont from PIL.ImageFont import FreeTypeFont from PIL import ImageFilter logging.basicConfig(level="DEBUG") logger = logging.getLogger(__name__) BG_COMP = "bg_composant" PROMPT_VAR = "prompt_variable" GENERATION_VAR = "generation_var" BOX_COMP = "box_comp" def wrap_text(font: FreeTypeFont, text: str, max_width: int, direction: str = "ltr") -> str: """ Wraps the text at the given width. :param font: Font to use. :param text: Text to fit. :param max_width: Maximum width of the final text, in pixels. :param max_height: Maximum height height of the final text, in pixels. :param spacing: The number of pixels between lines. :param direction: Direction of the text. It can be 'rtl' (right to left), 'ltr' (left to right) or 'ttb' (top to bottom). Requires libraqm. :return: The wrapped text. """ words = text.split() lines: list[str] = [""] curr_line_width = 0 for word in words: if curr_line_width == 0: word_width = font.getlength(word, direction) lines[-1] = word curr_line_width = word_width else: new_line_width = font.getlength(f"{lines[-1]} {word}", direction) if new_line_width > max_width: # Word is too long to fit on the current line word_width = font.getlength(word, direction) # Put the word on the next line lines.append(word) curr_line_width = word_width else: # Put the word on the current line lines[-1] = f"{lines[-1]} {word}" curr_line_width = new_line_width return "\n".join(lines) # tr.wrap(my_str, width=30) def _fit_paragraph( paragraph, font: FreeTypeFont, max_width: int, max_height: int, line_height, num_lines: List[str], spacing: int = 4, direction: str = "ltr", ): paragraph_lines: list[str] = [""] curr_line_width = 0 # This is a very bad splitter for Chinese words = list(paragraph) if re.search(u'[\u4e00-\u9fff]',paragraph) else paragraph.split(" ") for word in words: if curr_line_width == 0: word_width = font.getlength(word, direction) if word_width > max_width: # Word is longer than max_width return None paragraph_lines[-1] = word curr_line_width = word_width else: new_line_width = font.getlength(f"{paragraph_lines[-1]} {word}", direction) if new_line_width > max_width: # Word is too long to fit on the current line word_width = font.getlength(word, direction) new_num_lines = num_lines + len(paragraph_lines) + 1 new_text_height = (new_num_lines * line_height) + (new_num_lines * spacing) if word_width > max_width or new_text_height > max_height: # Word is longer than max_width, and # adding a new line would make the text too tall return None # Put the word on the next line paragraph_lines.append(word) curr_line_width = word_width else: # Put the word on the current line paragraph_lines[-1] = f"{paragraph_lines[-1]} {word}" curr_line_width = new_line_width return paragraph_lines def try_fit_text( font: FreeTypeFont, prompt: str, generation: str, max_width: int, max_height: int, spacing: int = 4, direction: str = "ltr", ) -> Optional[str]: """ Attempts to wrap the text into a rectangle. Tries to fit the text into a box using the given font at decreasing sizes, based on ``scale_factor``. Makes ``max_iterations`` attempts. :param font: Font to use. :param text: Text to fit. :param max_width: Maximum width of the final text, in pixels. :param max_height: Maximum height height of the final text, in pixels. :param spacing: The number of pixels between lines. :param direction: Direction of the text. It can be 'rtl' (right to left), 'ltr' (left to right) or 'ttb' (top to bottom). Requires libraqm. :return: If able to fit the text, the wrapped text. Otherwise, ``None``. """ line_height = font.size if line_height > max_height: # The line height is already too big return None prompt_lines: list[str] = [] paragraphs = prompt.split("\n") for paragraph in paragraphs: paragraph_lines = _fit_paragraph( paragraph=paragraph, font=font, max_width=max_width, max_height=max_height, line_height=line_height, spacing=spacing, num_lines=len(prompt_lines), direction=direction, ) if paragraph_lines is None: return None prompt_lines.extend(paragraph_lines) generation_lines: list[str] = [] paragraphs = f"{prompt_lines[-1]}{generation}".split("\n") for paragraph in paragraphs: paragraph_lines = _fit_paragraph( paragraph=paragraph, font=font, max_width=max_width, max_height=max_height, line_height=line_height, spacing=spacing, num_lines=len(prompt_lines) + len(generation_lines), direction=direction, ) if paragraph_lines is None: return None generation_lines.extend(paragraph_lines) generation_lines[0] = generation_lines[0][len(prompt_lines[-1]):] return "\n".join(prompt_lines), "\n".join(generation_lines) # pylint: disable=too-many-arguments def fit_text( font, prompt: str, generation: str, max_width: int, max_height: int, spacing: int = 4, scale_factor: float = 0.8, max_iterations: int = 5, direction: str = "ltr", ) -> Tuple[FreeTypeFont, str]: """ Automatically determines text wrapping and appropriate font size. Tries to fit the text into a box using the given font at decreasing sizes, based on ``scale_factor``. Makes ``max_iterations`` attempts. If unable to find an appropriate font size within ``max_iterations`` attempts, wraps the text at the last attempted size. :param font: Font to use. :param text: Text to fit. :param max_width: Maximum width of the final text, in pixels. :param max_height: Maximum height height of the final text, in pixels. :param spacing: The number of pixels between lines. :param scale_factor: :param max_iterations: Maximum number of attempts to try to fit the text. :param direction: Direction of the text. It can be 'rtl' (right to left), 'ltr' (left to right) or 'ttb' (top to bottom). Requires libraqm. :return: The font at the appropriate size and the wrapped text. """ initial_font_size = font.size # logger.debug('Trying to fit text "%s"', text) for i in range(max_iterations): trial_font_size = int(initial_font_size * pow(scale_factor, i)) trial_font = font.font_variant(size=trial_font_size) logger.debug("Trying font size %i", trial_font_size) wrapped = try_fit_text(trial_font, prompt, generation, max_width, max_height, spacing, direction) if wrapped is not None: logger.debug("Successfully fit text") return (trial_font, wrapped) # Give up and wrap the text at the last size logger.debug("Gave up trying to fit text; just wrapping text") wrapped = wrap_text(trial_font, prompt, max_width, direction) + wrap_text( trial_font, generation, max_width, direction ) return (trial_font, wrapped) def main( prompt, generation, width, height, assets_path, font_path, colors, frame_to_box_margin, text_to_text_box_margin, init_font_size, right_align = False ): # prompt_color = "#ffffff" # text_color = "#FF57A0" # text_box = ((500, 500, 2300, 1800)) # margin = 50 # input_color = "#CDD2E3" right_align_params = {"direction":'rtl',"align":'right',"features":'rtla'} if right_align else {} text_box_margin_r = frame_to_box_margin text_box_margin_t = frame_to_box_margin text_box_margin_l = frame_to_box_margin text_box_margin_b = int(height / 4.5) text_box = ( text_box_margin_l, text_box_margin_t, width - text_box_margin_r, height - text_box_margin_b, ) background = Image.new("RGB", (width, height), color=colors[BG_COMP]) # Get assets assets_path = Path(assets_path) flower_1 = Image.open(assets_path / "image_1.png").copy() flower_2 = Image.open(assets_path / "image_2.png").copy() shadow = Image.open(assets_path / "image_3.png").copy() bloom_logo = Image.open(assets_path / "image_4.png").copy() input_info = Image.open(assets_path / "image_7.png").copy() output_info = Image.open(assets_path / "image_6.png").copy() details = Image.open(assets_path / "image_5.png").copy() flower_1_offsets = (int(width - flower_1.width * 2 / 3), int(-flower_1.height / 3)) background.paste(flower_1, flower_1_offsets, flower_1) flower_2_offsets = ( -int(flower_2.width * 2 / 5), int(height / 2 - flower_2.height / 2), ) background.paste(flower_2, flower_2_offsets, flower_2) bloom_offsets = ( frame_to_box_margin, int(height - bloom_logo.height - frame_to_box_margin), ) background.paste(bloom_logo, bloom_offsets, bloom_logo) input_info_offsets = ( width - details.width - text_to_text_box_margin - input_info.width - output_info.width , int(height - details.height - frame_to_box_margin), ) background.paste(input_info, input_info_offsets, input_info) output_info_offsets =( width - details.width - text_to_text_box_margin - input_info.width - output_info.width , int(height - details.height - frame_to_box_margin + input_info.height + text_to_text_box_margin), ) background.paste(output_info, output_info_offsets, output_info) details_offsets = ( width - frame_to_box_margin - details.width , int(height - details.height - frame_to_box_margin), ) background.paste(details, details_offsets, details) box_margin = ( text_box[0] + text_to_text_box_margin, text_box[1] + text_to_text_box_margin, text_box[2] - text_to_text_box_margin, text_box[3] - text_to_text_box_margin, ) # text_box = ImageText(box_margin) drawing = ImageDraw.Draw(background, "RGBA") # Text box for main text input_color = colors[BOX_COMP][1:] input_color = tuple(int(input_color[i : i + 2], 16) for i in (0, 2, 4)) + ( int(255 * 0.99), ) # RGB + A drawing.rounded_rectangle(text_box, outline="#000", fill=input_color, radius=47.9) # Adapt text size to box # font, (prompt_a, generation_a) = adapt_text_to_ratio(box_margin, prompt, generation) # generation_a = adapt_text_to_ratio(box_margin, generation, prompt) # font = adapt_font_to_text(box_margin, prompt_a + generation_a, font_path=font_path) init_font = ImageFont.truetype(font_path, size=init_font_size) font, (prompt_a, generation_a) = fit_text( prompt=prompt, generation=generation, font=init_font, max_height=box_margin[3] - box_margin[1], max_width=box_margin[2] - box_margin[0], max_iterations=50, scale_factor=0.95, direction="rtl" if right_align else "ltr" ) # Prompt, main, then the last line prompt_s = prompt_a.split("\n") prompt_main = "\n".join(prompt_s[:-1]) prompt_lastline = prompt_s[-1] drawing.multiline_text( (box_margin[0], box_margin[1]), prompt_main, colors[PROMPT_VAR], font, **right_align_params ) end_prompt_main = font.getsize_multiline(prompt_main) end_prompt_last = font.getsize_multiline(prompt_lastline) drawing.multiline_text( ((box_margin[2] - end_prompt_last[0]) if right_align else box_margin[0], box_margin[1] + end_prompt_main[1]), prompt_lastline, colors[PROMPT_VAR], font, **right_align_params ) # Generated text, first line, then the rest generation_split = generation_a.split("\n") generation_firstline = generation_split[0] generation_main = "\n".join(generation_split[1:]) drawing.multiline_text( # margin x + length(last line of prompt), margin Y + length(main part of prompt) (box_margin[0], box_margin[1] + end_prompt_main[1]) if right_align else \ (box_margin[0] + end_prompt_last[0], box_margin[1] + end_prompt_main[1]), generation_firstline, colors[GENERATION_VAR], font, **right_align_params ) drawing.multiline_text( (box_margin[0], box_margin[1] + end_prompt_main[1] + end_prompt_last[1]), generation_main, colors[GENERATION_VAR], font, **right_align_params ) final_font_size = font.size return final_font_size, background