Prepare for upgrade

.gitattributes

@@ -25,3 +25,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ttf filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

DejaVuSans.ttf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7da195a74c55bef988d0d48f9508bd5d849425c1770dba5d7bfc6ce9ed848954
+size 757076

GoNotoCurrent.ttf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89dea97df90b5fa98dc97b12a9fb2ad072a5bc345fba4b70de7ad1f1bb8e49bd
+size 14804740

screenshot.py

@@ -0,0 +1,387 @@
+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