mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
143 lines
4.0 KiB
Python
143 lines
4.0 KiB
Python
from abc import ABC, abstractmethod
|
|
from logging import Logger
|
|
from typing import List, Union, Optional
|
|
from invokeai.app.services.board_image_record_storage import BoardImageRecordStorageBase
|
|
from invokeai.app.services.board_record_storage import (
|
|
BoardRecord,
|
|
BoardRecordStorageBase,
|
|
)
|
|
|
|
from invokeai.app.services.image_record_storage import (
|
|
ImageRecordStorageBase,
|
|
OffsetPaginatedResults,
|
|
)
|
|
from invokeai.app.services.models.board_record import BoardDTO
|
|
from invokeai.app.services.models.image_record import ImageDTO, image_record_to_dto
|
|
from invokeai.app.services.urls import UrlServiceBase
|
|
|
|
|
|
class BoardImagesServiceABC(ABC):
|
|
"""High-level service for board-image relationship management."""
|
|
|
|
@abstractmethod
|
|
def add_image_to_board(
|
|
self,
|
|
board_id: str,
|
|
image_name: str,
|
|
) -> None:
|
|
"""Adds an image to a board."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def remove_image_from_board(
|
|
self,
|
|
board_id: str,
|
|
image_name: str,
|
|
) -> None:
|
|
"""Removes an image from a board."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_images_for_board(
|
|
self,
|
|
board_id: str,
|
|
) -> OffsetPaginatedResults[ImageDTO]:
|
|
"""Gets images for a board."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_board_for_image(
|
|
self,
|
|
image_name: str,
|
|
) -> Optional[str]:
|
|
"""Gets an image's board id, if it has one."""
|
|
pass
|
|
|
|
|
|
class BoardImagesServiceDependencies:
|
|
"""Service dependencies for the BoardImagesService."""
|
|
|
|
board_image_records: BoardImageRecordStorageBase
|
|
board_records: BoardRecordStorageBase
|
|
image_records: ImageRecordStorageBase
|
|
urls: UrlServiceBase
|
|
logger: Logger
|
|
|
|
def __init__(
|
|
self,
|
|
board_image_record_storage: BoardImageRecordStorageBase,
|
|
image_record_storage: ImageRecordStorageBase,
|
|
board_record_storage: BoardRecordStorageBase,
|
|
url: UrlServiceBase,
|
|
logger: Logger,
|
|
):
|
|
self.board_image_records = board_image_record_storage
|
|
self.image_records = image_record_storage
|
|
self.board_records = board_record_storage
|
|
self.urls = url
|
|
self.logger = logger
|
|
|
|
|
|
class BoardImagesService(BoardImagesServiceABC):
|
|
_services: BoardImagesServiceDependencies
|
|
|
|
def __init__(self, services: BoardImagesServiceDependencies):
|
|
self._services = services
|
|
|
|
def add_image_to_board(
|
|
self,
|
|
board_id: str,
|
|
image_name: str,
|
|
) -> None:
|
|
self._services.board_image_records.add_image_to_board(board_id, image_name)
|
|
|
|
def remove_image_from_board(
|
|
self,
|
|
board_id: str,
|
|
image_name: str,
|
|
) -> None:
|
|
self._services.board_image_records.remove_image_from_board(board_id, image_name)
|
|
|
|
def get_images_for_board(
|
|
self,
|
|
board_id: str,
|
|
) -> OffsetPaginatedResults[ImageDTO]:
|
|
image_records = self._services.board_image_records.get_images_for_board(
|
|
board_id
|
|
)
|
|
image_dtos = list(
|
|
map(
|
|
lambda r: image_record_to_dto(
|
|
r,
|
|
self._services.urls.get_image_url(r.image_name),
|
|
self._services.urls.get_image_url(r.image_name, True),
|
|
board_id,
|
|
),
|
|
image_records.items,
|
|
)
|
|
)
|
|
return OffsetPaginatedResults[ImageDTO](
|
|
items=image_dtos,
|
|
offset=image_records.offset,
|
|
limit=image_records.limit,
|
|
total=image_records.total,
|
|
)
|
|
|
|
def get_board_for_image(
|
|
self,
|
|
image_name: str,
|
|
) -> Optional[str]:
|
|
board_id = self._services.board_image_records.get_board_for_image(image_name)
|
|
return board_id
|
|
|
|
|
|
def board_record_to_dto(
|
|
board_record: BoardRecord, cover_image_name: Optional[str], image_count: int
|
|
) -> BoardDTO:
|
|
"""Converts a board record to a board DTO."""
|
|
return BoardDTO(
|
|
**board_record.dict(exclude={'cover_image_name'}),
|
|
cover_image_name=cover_image_name,
|
|
image_count=image_count,
|
|
)
|