Spaces:
Runtime error
Runtime error
| 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 | |