add asset_count to BoardDTO and split it out from image count

invokeai/app/services/board_image_records/board_image_records_base.py

@@ -49,3 +49,11 @@ class BoardImageRecordStorageBase(ABC):
     ) -> int:
         """Gets the number of images for a board."""
         pass
+
+    @abstractmethod
+    def get_asset_count_for_board(
+        self,
+        board_id: str,
+    ) -> int:
+        """Gets the number of assets for a board."""
+        pass

invokeai/app/services/board_image_records/board_image_records_sqlite.py

@@ -3,6 +3,8 @@ from typing import Optional, cast
 
 from invokeai.app.services.board_image_records.board_image_records_base import BoardImageRecordStorageBase
 from invokeai.app.services.image_records.image_records_common import (
+    ASSETS_CATEGORIES,
+    IMAGE_CATEGORIES,
     ImageCategory,
     ImageRecord,
     deserialize_image_record,
@@ -151,15 +153,38 @@ class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
 
     def get_image_count_for_board(self, board_id: str) -> int:
         with self._db.transaction() as cursor:
+            # Convert the enum values to unique list of strings
+            category_strings = [c.value for c in set(IMAGE_CATEGORIES)]
+            # Create the correct length of placeholders
+            placeholders = ",".join("?" * len(category_strings))
             cursor.execute(
-                """--sql
+                f"""--sql
                     SELECT COUNT(*)
                     FROM board_images
                     INNER JOIN images ON board_images.image_name = images.image_name
-                    WHERE images.is_intermediate = FALSE
+                    WHERE images.is_intermediate = FALSE AND images.image_category IN ( {placeholders} )
                     AND board_images.board_id = ?;
                     """,
-                (board_id,),
+                (board_id),
+            )
+            count = cast(int, cursor.fetchone()[0])
+        return count
+
+    def get_asset_count_for_board(self, board_id: str) -> int:
+        with self._db.transaction() as cursor:
+            # Convert the enum values to unique list of strings
+            category_strings = [c.value for c in set(ASSETS_CATEGORIES)]
+            # Create the correct length of placeholders
+            placeholders = ",".join("?" * len(category_strings))
+            cursor.execute(
+                f"""--sql
+                    SELECT COUNT(*)
+                    FROM board_images
+                    INNER JOIN images ON board_images.image_name = images.image_name
+                    WHERE images.is_intermediate = FALSE AND images.image_category IN ( {placeholders} )
+                    AND board_images.board_id = ?;
+                    """,
+                (board_id),
             )
             count = cast(int, cursor.fetchone()[0])
         return count

invokeai/app/services/boards/boards_common.py

@@ -12,15 +12,18 @@ class BoardDTO(BoardRecord):
     """The URL of the thumbnail of the most recent image in the board."""
     image_count: int = Field(description="The number of images in the board.")
     """The number of images in the board."""
+    asset_count: int = Field(description="The number of assets in the board.")
+    """The number of assets in the board."""
     video_count: int = Field(description="The number of videos in the board.")
     """The number of videos in the board."""
 
 
-def board_record_to_dto(board_record: BoardRecord, cover_image_name: Optional[str], image_count: int, video_count: int) -> BoardDTO:
+def board_record_to_dto(board_record: BoardRecord, cover_image_name: Optional[str], image_count: int, asset_count: int, video_count: int) -> BoardDTO:
     """Converts a board record to a board DTO."""
     return BoardDTO(
         **board_record.model_dump(exclude={"cover_image_name"}),
         cover_image_name=cover_image_name,
         image_count=image_count,
+        asset_count=asset_count,
         video_count=video_count,
     )

invokeai/app/services/boards/boards_default.py

@@ -17,7 +17,7 @@ class BoardService(BoardServiceABC):
         board_name: str,
     ) -> BoardDTO:
         board_record = self.__invoker.services.board_records.save(board_name)
-        return board_record_to_dto(board_record, None, 0)
+        return board_record_to_dto(board_record, None, 0, 0, 0)
 
     def get_dto(self, board_id: str) -> BoardDTO:
         board_record = self.__invoker.services.board_records.get(board_id)
@@ -27,8 +27,9 @@ class BoardService(BoardServiceABC):
         else:
             cover_image_name = None
         image_count = self.__invoker.services.board_image_records.get_image_count_for_board(board_id)
+        asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(board_id)
         video_count = 0 # noop for OSS
-        return board_record_to_dto(board_record, cover_image_name, image_count, video_count)
+        return board_record_to_dto(board_record, cover_image_name, image_count, asset_count, video_count)
 
     def update(
         self,
@@ -43,8 +44,9 @@ class BoardService(BoardServiceABC):
             cover_image_name = None
 
         image_count = self.__invoker.services.board_image_records.get_image_count_for_board(board_id)
+        asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(board_id)
         video_count = 0 # noop for OSS
-        return board_record_to_dto(board_record, cover_image_name, image_count, video_count)
+        return board_record_to_dto(board_record, cover_image_name, image_count, asset_count, video_count)
 
     def delete(self, board_id: str) -> None:
         self.__invoker.services.board_records.delete(board_id)
@@ -69,8 +71,9 @@ class BoardService(BoardServiceABC):
                 cover_image_name = None
 
             image_count = self.__invoker.services.board_image_records.get_image_count_for_board(r.board_id)
+            asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(r.board_id)
             video_count = 0 # noop for OSS
-            board_dtos.append(board_record_to_dto(r, cover_image_name, image_count, video_count))
+            board_dtos.append(board_record_to_dto(r, cover_image_name, image_count, asset_count, video_count))
 
         return OffsetPaginatedResults[BoardDTO](items=board_dtos, offset=offset, limit=limit, total=len(board_dtos))
 
@@ -87,7 +90,8 @@ class BoardService(BoardServiceABC):
                 cover_image_name = None
 
             image_count = self.__invoker.services.board_image_records.get_image_count_for_board(r.board_id)
+            asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(r.board_id)
             video_count = 0 # noop for OSS
-            board_dtos.append(board_record_to_dto(r, cover_image_name, image_count, video_count))
+            board_dtos.append(board_record_to_dto(r, cover_image_name, image_count, asset_count, video_count))
 
         return board_dtos

invokeai/app/services/image_records/image_records_common.py

@@ -58,6 +58,9 @@ class ImageCategory(str, Enum, metaclass=MetaEnum):
     """OTHER: The image is some other type of image with a specialized purpose. To be used by external nodes."""
 
 
+IMAGE_CATEGORIES: list[ImageCategory] = [ImageCategory.GENERAL];
+ASSETS_CATEGORIES: list[ImageCategory] = [ImageCategory.CONTROL, ImageCategory.MASK, ImageCategory.USER, ImageCategory.OTHER];
+
 class InvalidImageCategoryException(ValueError):
     """Raised when a provided value is not a valid ImageCategory.