diff --git a/invokeai/app/api/routers/board_images.py b/invokeai/app/api/routers/board_images.py index eb193f6585..cb5e0ab51a 100644 --- a/invokeai/app/api/routers/board_images.py +++ b/invokeai/app/api/routers/board_images.py @@ -1,21 +1,12 @@ from fastapi import Body, HTTPException from fastapi.routing import APIRouter -from pydantic import BaseModel, Field from invokeai.app.api.dependencies import ApiDependencies +from invokeai.app.services.images.images_common import AddImagesToBoardResult, RemoveImagesFromBoardResult board_images_router = APIRouter(prefix="/v1/board_images", tags=["boards"]) -class AddImagesToBoardResult(BaseModel): - board_id: str = Field(description="The id of the board the images were added to") - added_image_names: list[str] = Field(description="The image names that were added to the board") - - -class RemoveImagesFromBoardResult(BaseModel): - removed_image_names: list[str] = Field(description="The image names that were removed from their board") - - @board_images_router.post( "/", operation_id="add_image_to_board", @@ -23,17 +14,26 @@ class RemoveImagesFromBoardResult(BaseModel): 201: {"description": "The image was added to a board successfully"}, }, status_code=201, + response_model=AddImagesToBoardResult, ) async def add_image_to_board( board_id: str = Body(description="The id of the board to add to"), image_name: str = Body(description="The name of the image to add"), -): +) -> AddImagesToBoardResult: """Creates a board_image""" try: - result = ApiDependencies.invoker.services.board_images.add_image_to_board( - board_id=board_id, image_name=image_name + added_images: set[str] = set() + affected_boards: set[str] = set() + old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" + ApiDependencies.invoker.services.board_images.add_image_to_board(board_id=board_id, image_name=image_name) + added_images.add(image_name) + affected_boards.add(board_id) + affected_boards.add(old_board_id) + + return AddImagesToBoardResult( + added_images=list(added_images), + affected_boards=list(affected_boards), ) - return result except Exception: raise HTTPException(status_code=500, detail="Failed to add image to board") @@ -45,14 +45,25 @@ async def add_image_to_board( 201: {"description": "The image was removed from the board successfully"}, }, status_code=201, + response_model=RemoveImagesFromBoardResult, ) async def remove_image_from_board( image_name: str = Body(description="The name of the image to remove", embed=True), -): +) -> RemoveImagesFromBoardResult: """Removes an image from its board, if it had one""" try: - result = ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name) - return result + removed_images: set[str] = set() + affected_boards: set[str] = set() + old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" + ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name) + removed_images.add(image_name) + affected_boards.add("none") + affected_boards.add(old_board_id) + return RemoveImagesFromBoardResult( + removed_images=list(removed_images), + affected_boards=list(affected_boards), + ) + except Exception: raise HTTPException(status_code=500, detail="Failed to remove image from board") @@ -72,16 +83,25 @@ async def add_images_to_board( ) -> AddImagesToBoardResult: """Adds a list of images to a board""" try: - added_image_names: list[str] = [] + added_images: set[str] = set() + affected_boards: set[str] = set() for image_name in image_names: try: + old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" ApiDependencies.invoker.services.board_images.add_image_to_board( - board_id=board_id, image_name=image_name + board_id=board_id, + image_name=image_name, ) - added_image_names.append(image_name) + added_images.add(image_name) + affected_boards.add(board_id) + affected_boards.add(old_board_id) + except Exception: pass - return AddImagesToBoardResult(board_id=board_id, added_image_names=added_image_names) + return AddImagesToBoardResult( + added_images=list(added_images), + affected_boards=list(affected_boards), + ) except Exception: raise HTTPException(status_code=500, detail="Failed to add images to board") @@ -100,13 +120,20 @@ async def remove_images_from_board( ) -> RemoveImagesFromBoardResult: """Removes a list of images from their board, if they had one""" try: - removed_image_names: list[str] = [] + removed_images: set[str] = set() + affected_boards: set[str] = set() for image_name in image_names: try: + old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name) - removed_image_names.append(image_name) + removed_images.add(image_name) + affected_boards.add("none") + affected_boards.add(old_board_id) except Exception: pass - return RemoveImagesFromBoardResult(removed_image_names=removed_image_names) + return RemoveImagesFromBoardResult( + removed_images=list(removed_images), + affected_boards=list(affected_boards), + ) except Exception: raise HTTPException(status_code=500, detail="Failed to remove images from board") diff --git a/invokeai/app/api/routers/images.py b/invokeai/app/api/routers/images.py index f65cafc731..a2ac6b45c8 100644 --- a/invokeai/app/api/routers/images.py +++ b/invokeai/app/api/routers/images.py @@ -17,7 +17,13 @@ from invokeai.app.services.image_records.image_records_common import ( ImageRecordChanges, ResourceOrigin, ) -from invokeai.app.services.images.images_common import ImageDTO, ImageUrlsDTO +from invokeai.app.services.images.images_common import ( + DeleteImagesResult, + ImageDTO, + ImageUrlsDTO, + StarredImagesResult, + UnstarredImagesResult, +) from invokeai.app.services.shared.pagination import OffsetPaginatedResults from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection from invokeai.app.util.controlnet_utils import heuristic_resize_fast @@ -153,18 +159,30 @@ async def create_image_upload_entry( raise HTTPException(status_code=501, detail="Not implemented") -@images_router.delete("/i/{image_name}", operation_id="delete_image") +@images_router.delete("/i/{image_name}", operation_id="delete_image", response_model=DeleteImagesResult) async def delete_image( image_name: str = Path(description="The name of the image to delete"), -) -> None: +) -> DeleteImagesResult: """Deletes an image""" + deleted_images: set[str] = set() + affected_boards: set[str] = set() + try: + image_dto = ApiDependencies.invoker.services.images.get_dto(image_name) + board_id = image_dto.board_id or "none" ApiDependencies.invoker.services.images.delete(image_name) + deleted_images.add(image_name) + affected_boards.add(board_id) except Exception: # TODO: Does this need any exception handling at all? pass + return DeleteImagesResult( + deleted_images=list(deleted_images), + affected_boards=list(affected_boards), + ) + @images_router.delete("/intermediates", operation_id="clear_intermediates") async def clear_intermediates() -> int: @@ -376,31 +394,32 @@ async def list_image_dtos( return image_dtos -class DeleteImagesFromListResult(BaseModel): - deleted_images: list[str] - - -@images_router.post("/delete", operation_id="delete_images_from_list", response_model=DeleteImagesFromListResult) +@images_router.post("/delete", operation_id="delete_images_from_list", response_model=DeleteImagesResult) async def delete_images_from_list( image_names: list[str] = Body(description="The list of names of images to delete", embed=True), -) -> DeleteImagesFromListResult: +) -> DeleteImagesResult: try: - deleted_images: list[str] = [] + deleted_images: set[str] = set() + affected_boards: set[str] = set() for image_name in image_names: try: + image_dto = ApiDependencies.invoker.services.images.get_dto(image_name) + board_id = image_dto.board_id or "none" ApiDependencies.invoker.services.images.delete(image_name) - deleted_images.append(image_name) + deleted_images.add(image_name) + affected_boards.add(board_id) except Exception: pass - return DeleteImagesFromListResult(deleted_images=deleted_images) + return DeleteImagesResult( + deleted_images=list(deleted_images), + affected_boards=list(affected_boards), + ) except Exception: raise HTTPException(status_code=500, detail="Failed to delete images") -@images_router.delete( - "/uncategorized", operation_id="delete_uncategorized_images", response_model=DeleteImagesFromListResult -) -async def delete_uncategorized_images() -> DeleteImagesFromListResult: +@images_router.delete("/uncategorized", operation_id="delete_uncategorized_images", response_model=DeleteImagesResult) +async def delete_uncategorized_images() -> DeleteImagesResult: """Deletes all images that are uncategorized""" image_names = ApiDependencies.invoker.services.board_images.get_all_board_image_names_for_board( @@ -408,14 +427,19 @@ async def delete_uncategorized_images() -> DeleteImagesFromListResult: ) try: - deleted_images: list[str] = [] + deleted_images: set[str] = set() + affected_boards: set[str] = set() for image_name in image_names: try: ApiDependencies.invoker.services.images.delete(image_name) - deleted_images.append(image_name) + deleted_images.add(image_name) + affected_boards.add("none") except Exception: pass - return DeleteImagesFromListResult(deleted_images=deleted_images) + return DeleteImagesResult( + deleted_images=list(deleted_images), + affected_boards=list(affected_boards), + ) except Exception: raise HTTPException(status_code=500, detail="Failed to delete images") @@ -424,36 +448,50 @@ class ImagesUpdatedFromListResult(BaseModel): updated_image_names: list[str] = Field(description="The image names that were updated") -@images_router.post("/star", operation_id="star_images_in_list", response_model=ImagesUpdatedFromListResult) +@images_router.post("/star", operation_id="star_images_in_list", response_model=StarredImagesResult) async def star_images_in_list( image_names: list[str] = Body(description="The list of names of images to star", embed=True), -) -> ImagesUpdatedFromListResult: +) -> StarredImagesResult: try: - updated_image_names: list[str] = [] + starred_images: set[str] = set() + affected_boards: set[str] = set() for image_name in image_names: try: - ApiDependencies.invoker.services.images.update(image_name, changes=ImageRecordChanges(starred=True)) - updated_image_names.append(image_name) + updated_image_dto = ApiDependencies.invoker.services.images.update( + image_name, changes=ImageRecordChanges(starred=True) + ) + starred_images.add(image_name) + affected_boards.add(updated_image_dto.board_id or "none") except Exception: pass - return ImagesUpdatedFromListResult(updated_image_names=updated_image_names) + return StarredImagesResult( + starred_images=list(starred_images), + affected_boards=list(affected_boards), + ) except Exception: raise HTTPException(status_code=500, detail="Failed to star images") -@images_router.post("/unstar", operation_id="unstar_images_in_list", response_model=ImagesUpdatedFromListResult) +@images_router.post("/unstar", operation_id="unstar_images_in_list", response_model=UnstarredImagesResult) async def unstar_images_in_list( image_names: list[str] = Body(description="The list of names of images to unstar", embed=True), -) -> ImagesUpdatedFromListResult: +) -> UnstarredImagesResult: try: - updated_image_names: list[str] = [] + unstarred_images: set[str] = set() + affected_boards: set[str] = set() for image_name in image_names: try: - ApiDependencies.invoker.services.images.update(image_name, changes=ImageRecordChanges(starred=False)) - updated_image_names.append(image_name) + updated_image_dto = ApiDependencies.invoker.services.images.update( + image_name, changes=ImageRecordChanges(starred=False) + ) + unstarred_images.add(image_name) + affected_boards.add(updated_image_dto.board_id or "none") except Exception: pass - return ImagesUpdatedFromListResult(updated_image_names=updated_image_names) + return UnstarredImagesResult( + unstarred_images=list(unstarred_images), + affected_boards=list(affected_boards), + ) except Exception: raise HTTPException(status_code=500, detail="Failed to unstar images") diff --git a/invokeai/app/services/images/images_common.py b/invokeai/app/services/images/images_common.py index 0464244b94..311f6e556d 100644 --- a/invokeai/app/services/images/images_common.py +++ b/invokeai/app/services/images/images_common.py @@ -1,6 +1,6 @@ from typing import Optional -from pydantic import Field +from pydantic import BaseModel, Field from invokeai.app.services.image_records.image_records_common import ImageRecord from invokeai.app.util.model_exclude_null import BaseModelExcludeNull @@ -39,3 +39,27 @@ def image_record_to_dto( thumbnail_url=thumbnail_url, board_id=board_id, ) + + +class ResultWithAffectedBoards(BaseModel): + affected_boards: list[str] = Field(description="The ids of boards affected by the delete operation") + + +class DeleteImagesResult(ResultWithAffectedBoards): + deleted_images: list[str] = Field(description="The names of the images that were deleted") + + +class StarredImagesResult(ResultWithAffectedBoards): + starred_images: list[str] = Field(description="The names of the images that were starred") + + +class UnstarredImagesResult(ResultWithAffectedBoards): + unstarred_images: list[str] = Field(description="The names of the images that were unstarred") + + +class AddImagesToBoardResult(ResultWithAffectedBoards): + added_images: list[str] = Field(description="The image names that were added to the board") + + +class RemoveImagesFromBoardResult(ResultWithAffectedBoards): + removed_images: list[str] = Field(description="The image names that were removed from their board")