v 2.1.5

Add `GridAnnotation` node
Remove old debug `print`s
Update typing to python3.8
Add support Annotations for `BaseImagesGridNode`
Improved and simplified code
Remove `grid_annotations` util
Add support annotation for `create_images_grid` util
Add static path and font

__init__.py

@@ -1,4 +1,10 @@
-from .src import LatentCombineNode, ImagesGridByColumnsNode, ImagesGridByRowsNode, ImageCombineNode
+from .src import (
+    LatentCombineNode,
+    ImagesGridByColumnsNode,
+    ImagesGridByRowsNode,
+    ImageCombineNode,
+    GridAnnotationNode,
+)
 
 
 NODE_CLASS_MAPPINGS = {
@@ -6,4 +12,5 @@ NODE_CLASS_MAPPINGS = {
     "ImagesGridByColumns": ImagesGridByColumnsNode,
     "ImagesGridByRows": ImagesGridByRowsNode,
     "ImageCombine": ImageCombineNode,
+    "GridAnnotation": GridAnnotationNode,
 }

src/__init__.py

@@ -1,3 +1,4 @@
 from .nodes.images_grid import ImagesGridByColumnsNode, ImagesGridByRowsNode
 from .nodes.latent_combine import LatentCombineNode
 from .nodes.image_combine import ImageCombineNode
+from .nodes.grid_annotation import GridAnnotationNode

src/base.py

@@ -1,9 +1,12 @@
 import typing as t
+from pathlib import Path
+
+
+STATIC_PATH = Path(__file__).parent.parent / "static"
+
+Image = t.Any
 
 
 class BaseNode():
     CATEGORY: str = "ImagesGrid"
     FUNCTION: str = "execute"
-
-
-Image = t.Any

src/nodes/grid_annotation.py

@@ -0,0 +1,49 @@
+import typing as t
+
+from PIL import ImageFont
+
+from ..base import BaseNode, STATIC_PATH
+from ..utils import Annotation
+
+
+class GridAnnotationNode(BaseNode):
+    RETURN_TYPES: tuple[str] = ("GRID_ANNOTATION",)
+
+    @classmethod
+    def INPUT_TYPES(cls) -> dict[str, t.Any]:
+        return {
+            "required": {
+                "column_texts": ("STRING", {"multiline": False}),
+                "row_texts": ("STRING", {"multiline": False}),
+                "font_size": ("INT", {"default": 30, "min": 1}),
+            },
+        }
+
+    def execute(
+        self,
+        column_texts: str,
+        row_texts: str,
+        font_size: int,
+    ) -> tuple[Annotation]:
+        font = ImageFont.truetype(str(STATIC_PATH / "Roboto-Regular.ttf"), size=font_size)
+        column_texts_list = self._set_value_to_texts_list(
+            self._get_texts_from_string(column_texts),
+        )
+        row_texts_list = self._set_value_to_texts_list(
+            self._get_texts_from_string(row_texts),
+        )
+
+        result = Annotation(column_texts=column_texts_list, row_texts=row_texts_list, font=font)
+        return (result,)
+
+    def _get_texts_from_string(self, string: str) -> list[str]:
+        return [
+            result
+            for i in string.split(";")
+            if (result := i.strip()) != ""
+        ]
+
+    def _set_value_to_texts_list(self, texts_list: list[str]) -> list[str]:
+        if not texts_list:
+            return ["None"]
+        return texts_list

src/nodes/image_combine.py

@@ -6,10 +6,10 @@ from ..base import BaseNode, Image
 
 
 class ImageCombineNode(BaseNode):
-    RETURN_TYPES: t.Tuple[str] = ("IMAGE",)
+    RETURN_TYPES: tuple[str] = ("IMAGE",)
 
     @classmethod
-    def INPUT_TYPES(cls) -> t.Dict[str, t.Any]:
+    def INPUT_TYPES(cls) -> dict[str, t.Any]:
         return {
             "required": {
                 "image_1": ("IMAGE",),
@@ -21,12 +21,7 @@ class ImageCombineNode(BaseNode):
         self,
         image_1: Image,
         image_2: Image,
-    ) -> t.Tuple[Image]:
+    ) -> tuple[Image]:
-        print(image_1.size())
-        print(image_2.size())
-        print(image_1)
-
         result = torch.cat((image_1, image_2), 0)
-        print(result.size())
 
         return (result,)

src/nodes/images_grid.py

@@ -6,25 +6,42 @@ from ..utils import (
     pillow_to_tensor,
     create_images_grid_by_columns,
     create_images_grid_by_rows,
+    Annotation,
 )
 
 
 class BaseImagesGridNode(BaseNode):
-    RETURN_TYPES: t.Tuple[str] = ("IMAGE",)
+    RETURN_TYPES: tuple[str] = ("IMAGE",)
 
     @classmethod
-    def _create_input_types(cls, coordinate_name: str) -> t.Dict[str, t.Any]:
+    def _create_input_types(cls, coordinate_name: str) -> dict[str, t.Any]:
         return {
             "required": {
                 "images": ("IMAGE",),
                 "gap": ("INT", {"default": 0, "min": 0}),
                 coordinate_name: ("INT", {"default": 1, "min": 1}),
+            },
+            "optional": {
+                "annotation": ("GRID_ANNOTATION",),
             }
         }
 
-    def _create_execute(self, images, function, function_kw) -> t.Tuple[Image]:
+    def _create_execute(
+        self,
+        function: t.Callable,
+        \
+        images: Image,
+        gap: int,
+        annotation: Annotation | None = None,
+        **kw,
+    ) -> tuple[Image]:
         pillow_images = [tensor_to_pillow(i) for i in images]
-        pillow_grid = function(images=pillow_images, **function_kw)
+        pillow_grid = function(
+            images=pillow_images,
+            gap=gap,
+            annotation=annotation,
+            **kw,
+        )
         tensor_grid = pillow_to_tensor(pillow_grid)
 
         return (tensor_grid,)
@@ -32,17 +49,17 @@ class BaseImagesGridNode(BaseNode):
 
 class ImagesGridByColumnsNode(BaseImagesGridNode):
     @classmethod
-    def INPUT_TYPES(cls) -> t.Dict[str, t.Any]:
+    def INPUT_TYPES(cls) -> dict[str, t.Any]:
         return cls._create_input_types("max_columns")
 
-    def execute(self, images: Image, **kw) -> tuple[Image]:
+    def execute(self, **kw) -> tuple[Image]:
-        return self._create_execute(images, create_images_grid_by_columns, kw)
+        return self._create_execute(create_images_grid_by_columns, **kw)
 
 
 class ImagesGridByRowsNode(BaseImagesGridNode):
     @classmethod
-    def INPUT_TYPES(cls) -> t.Dict[str, t.Any]:
+    def INPUT_TYPES(cls) -> dict[str, t.Any]:
         return cls._create_input_types("max_rows")
 
-    def execute(self, images: Image, **kw) -> tuple[Image]:
+    def execute(self, **kw) -> tuple[Image]:
-        return self._create_execute(images, create_images_grid_by_rows, kw)
+        return self._create_execute(create_images_grid_by_rows, **kw)

src/nodes/latent_combine.py

@@ -6,10 +6,10 @@ from ..base import BaseNode, Image
 
 
 class LatentCombineNode(BaseNode):
-    RETURN_TYPES: t.Tuple[str] = ("LATENT",)
+    RETURN_TYPES: tuple[str] = ("LATENT",)
 
     @classmethod
-    def INPUT_TYPES(cls) -> t.Dict[str, t.Any]:
+    def INPUT_TYPES(cls) -> dict[str, t.Any]:
         return {
             "required": {
                 "latent_1": ("LATENT",),
@@ -19,9 +19,9 @@ class LatentCombineNode(BaseNode):
 
     def execute(
         self,
-        latent_1: t.Dict[str, t.Any],
+        latent_1: dict[str, t.Any],
-        latent_2: t.Dict[str, t.Any],
+        latent_2: dict[str, t.Any],
-    ) -> t.Tuple[t.Dict[str, t.Any]]:
+    ) -> tuple[dict[str, t.Any]]:
         samples = torch.cat((latent_1["samples"], latent_2["samples"]), 0)
 
         return ({"samples": samples},)

src/utils/__init__.py

@@ -1,6 +1,6 @@
 from .images_grid import (
     create_images_grid_by_columns,
     create_images_grid_by_rows,
+    Annotation,
 )
-from grid_annotations import create_grid_annotations
 from .tensor_convert import tensor_to_pillow, pillow_to_tensor

src/utils/grid_annotations.py

@@ -1,104 +0,0 @@
-import typing as t
-from dataclasses import dataclass
-
-from PIL import Image, ImageDraw, ImageFont
-
-
-@dataclass
-class GridImagesInfo():
-    image: Image.Image
-    gap: int
-    one_image_size: t.Tuple[int, int]
-
-
-def create_grid_annotations(
-    grid_info: GridImagesInfo,
-    column_texts,
-    row_texts,
-    font: ImageFont.ImageFont,
-    margin: int = 5,
-) -> Image.Image:
-    grid = grid_info.image
-    left_padding = int(max(map(font.getlength, row_texts))) + 2*margin  # type: ignore
-    top_padding = font.size + 2*margin  # type: ignore
-
-    image = Image.new(
-        "RGB",
-        (grid.size[0] + left_padding, grid.size[1] + top_padding),
-        color="white",
-    )
-    draw = ImageDraw.Draw(image)
-    draw.font = font  # type: ignore
-
-    image.paste(grid, (image.size[0] - grid.size[0], image.size[1] - grid.size[1]))
-    _draw_column_text(
-        draw=draw,
-        texts=column_texts,
-        grid_info=grid_info,
-        left_padding=left_padding,
-        top_padding=top_padding,
-    )
-    _draw_row_text(
-        draw=draw,
-        texts=row_texts,
-        grid_info=grid_info,
-        left_padding=left_padding,
-        top_padding=top_padding,
-    )
-
-    return image
-
-
-def _draw_column_text(
-    draw: ImageDraw.ImageDraw,
-    texts: t.List[str],
-    grid_info: GridImagesInfo,
-    left_padding: int,
-    top_padding: int,
-) -> None:
-    i = 0
-    x0 = left_padding
-    y0 = 0
-    x1 = left_padding + grid_info.one_image_size[0]
-    y1 = top_padding
-    while x0 != grid_info.image.size[0] + left_padding:
-        _draw_center_text(draw, (x0, y0, x1, y1), texts[i])
-        x0 += grid_info.one_image_size[0] + grid_info.gap
-        x1 += grid_info.one_image_size[0] + grid_info.gap
-        i += 1
-
-
-def _draw_row_text(
-    draw: ImageDraw.ImageDraw,
-    texts: t.List[str],
-    grid_info: GridImagesInfo,
-    left_padding: int,
-    top_padding: int,
-) -> None:
-    i = 0
-    x0 = 0
-    y0 = top_padding
-    x1 = left_padding
-    y1 = top_padding + grid_info.one_image_size[1]
-    while y0 != grid_info.image.size[1] + top_padding:
-        _draw_center_text(draw, (x0, y0, x1, y1), texts[i])
-        y0 += grid_info.one_image_size[1] + grid_info.gap
-        y1 += grid_info.one_image_size[1] + grid_info.gap
-        i += 1
-
-
-def _draw_center_text(
-    draw: ImageDraw.ImageDraw,
-    xy: t.Tuple[int, int, int, int],
-    text: str,
-    fill: t.Any = "black",
-) -> None:
-    _, _, *text_size = draw.textbbox((0, 0), text)
-    draw.text(
-        (
-            (xy[2] - text_size[0] + xy[0]) / 2,
-            (xy[3] - text_size[1] + xy[1]) / 2,
-        ),
-        text,
-        fill=fill,
-    )

src/utils/images_grid.py

@@ -1,39 +1,81 @@
 import typing as t
+from dataclasses import dataclass
+from contextlib import suppress
 
 from PIL import Image, ImageDraw, ImageFont
 
 
+@dataclass
+class Annotation():
+    column_texts: list[str]
+    row_texts: list[str]
+    font: ImageFont.FreeTypeFont
+
+
+@dataclass
+class _GridInfo():
+    image: Image.Image
+    gap: int
+    one_image_size: tuple[int, int]
+
+
 def create_images_grid_by_columns(
-    images: t.List[Image.Image],
+    images: list[Image.Image],
     gap: int,
     max_columns: int,
+    annotation: Annotation | None = None,
 ) -> Image.Image:
     max_rows = (len(images) + max_columns - 1) // max_columns
-    return _create_images_grid(images=images, gap=gap, max_columns=max_columns, max_rows=max_rows)
+    return _create_images_grid(images, gap, max_columns, max_rows, annotation)
 
 
 def create_images_grid_by_rows(
-    images: t.List[Image.Image],
+    images: list[Image.Image],
     gap: int,
     max_rows: int,
+    annotation: Annotation | None = None,
 ) -> Image.Image:
     max_columns = (len(images) + max_rows - 1) // max_rows
-    return _create_images_grid(images=images, gap=gap, max_columns=max_columns, max_rows=max_rows)
+    return _create_images_grid(images, gap, max_columns, max_rows, annotation)
 
 
 def _create_images_grid(
-    images: t.List[Image.Image],
+    images: list[Image.Image],
     gap: int,
     max_columns: int,
     max_rows: int,
+    annotation: Annotation | None,
 ) -> Image.Image:
     size = images[0].size
-
     grid_width = size[0] * max_columns + (max_columns - 1) * gap
     grid_height = size[1] * max_rows + (max_rows - 1) * gap
 
     grid_image = Image.new("RGB", (grid_width, grid_height), color="white")
 
+    _arrange_images_on_grid(grid_image, images=images, size=size, max_columns=max_columns, gap=gap)
+
+    if annotation is None:
+        return grid_image
+    return _create_grid_annotations(
+        grid_info=_GridInfo(
+            image=grid_image,
+            gap=gap,
+            one_image_size=size,
+        ),
+        column_texts=annotation.column_texts,
+        row_texts=annotation.row_texts,
+        font=annotation.font,
+    )
+
+
+def _arrange_images_on_grid(
+    grid_image: Image.Image,
+    /,
+    images: list[Image.Image],
+    size: tuple[int, int],
+    max_columns: int,
+    gap: int,
+):
     for i, image in enumerate(images):
         image = image.crop((0, 0, size[0], size[1]))
         x = (i % max_columns) * (size[0] + gap)
@@ -41,24 +83,112 @@ def _create_images_grid(
 
         grid_image.paste(image, (x, y))
 
-    return grid_image
+
+def _create_grid_annotations(
+    grid_info: _GridInfo,
+    column_texts,
+    row_texts,
+    font: ImageFont.FreeTypeFont,
+) -> Image.Image:
+    if not column_texts or not row_texts:
+        raise ValueError("Column text or row text is empty")
+
+    grid = grid_info.image
+    margin = font.size // 2
+    left_padding = int(max(map(font.getlength, row_texts))) + 2*margin
+    top_padding = font.size + 2*margin
+
+    image = Image.new(
+        "RGB",
+        (grid.size[0] + left_padding, grid.size[1] + top_padding),
+        color="white",
+    )
+    draw = ImageDraw.Draw(image)
+    draw.font = font  # type: ignore
+
+    _paste_image_to_lower_left_corner(image, grid)
+    _draw_column_text(
+        draw=draw,
+        texts=column_texts,
+        grid_info=grid_info,
+        left_padding=left_padding,
+        top_padding=top_padding,
+    )
+    _draw_row_text(
+        draw=draw,
+        texts=row_texts,
+        grid_info=grid_info,
+        left_padding=left_padding,
+        top_padding=top_padding,
+    )
+
+    return image
+
+
+def _draw_column_text(
+    draw: ImageDraw.ImageDraw,
+    texts: list[str],
+    grid_info: _GridInfo,
+    left_padding: int,
+    top_padding: int,
+) -> None:
+    i = 0
+    x0 = left_padding
+    y0 = 0
+    x1 = left_padding + grid_info.one_image_size[0]
+    y1 = top_padding
+    while x0 != grid_info.image.size[0] + left_padding + grid_info.gap:
+        i  = _draw_text_by_xy((x0, y0, x1, y1), i, draw=draw, texts=texts)
+        x0 += grid_info.one_image_size[0] + grid_info.gap
+        x1 += grid_info.one_image_size[0] + grid_info.gap
+
+
+def _draw_row_text(
+    draw: ImageDraw.ImageDraw,
+    texts: list[str],
+    grid_info: _GridInfo,
+    left_padding: int,
+    top_padding: int,
+) -> None:
+    i = 0
+    x0 = 0
+    y0 = top_padding
+    x1 = left_padding
+    y1 = top_padding + grid_info.one_image_size[1]
+    while y0 != grid_info.image.size[1] + top_padding + grid_info.gap:
+        i  = _draw_text_by_xy((x0, y0, x1, y1), i, draw=draw, texts=texts)
+        y0 += grid_info.one_image_size[1] + grid_info.gap
+        y1 += grid_info.one_image_size[1] + grid_info.gap
+
+
+def _draw_text_by_xy(
+    xy: tuple[int, int, int, int],
+    index: int,
+    \
+    draw: ImageDraw.ImageDraw,
+    texts: list[str],
+) -> int:
+    with suppress(IndexError):
+        _draw_center_text(draw, xy, texts[index])
+    return index + 1
 
 
 def _draw_center_text(
-  text: str,
+    draw: ImageDraw.ImageDraw,
-  draw: ImageDraw.ImageDraw,
+    xy: tuple[int, int, int, int],
-  font: ImageFont.ImageFont,
+    text: str,
-  fill: int = 128,
+    fill: t.Any = "black",
-):
+) -> None:
-  image = draw.im  # type: ignore
+    _, _, *text_size = draw.textbbox((0, 0), text)
-  _, _, *text_size = draw.textbbox((0, 0), text, font=font)
+    draw.text(
-  draw.text(
+        (
-    (
+            (xy[2] - text_size[0] + xy[0]) / 2,
-      (image.size[0]-text_size[0])/2,
+            (xy[3] - text_size[1] + xy[1]) / 2,
-      (image.size[1]-text_size[1])/2,
+        ),
-    ),
+        text,
-    text,
+        fill=fill,
-    font=font,
+    )
-    fill=fill,
+
-  )
+
-  return image
+def _paste_image_to_lower_left_corner(base: Image.Image, image: Image.Image) -> None:
+    base.paste(image, (base.size[0] - image.size[0], base.size[1] - image.size[1]))