InvokeAI/invokeai/app/api/routers/images.py

import io
from fastapi import HTTPException, Path, Query, Request, Response, UploadFile
from fastapi.routing import APIRouter
from fastapi.responses import FileResponse
from PIL import Image
from invokeai.app.models.image import (
    ImageCategory,
    ImageType,
)
from invokeai.app.services.models.image_record import ImageDTO, ImageUrlsDTO
from invokeai.app.services.item_storage import PaginatedResults

from ..dependencies import ApiDependencies

images_router = APIRouter(prefix="/v1/images", tags=["images"])


@images_router.post(
    "/",
    operation_id="upload_image",
    responses={
        201: {"description": "The image was uploaded successfully"},
        415: {"description": "Image upload failed"},
    },
    status_code=201,
)
async def upload_image(
    file: UploadFile,
    image_type: ImageType,
    request: Request,
    response: Response,
    image_category: ImageCategory = ImageCategory.IMAGE,
) -> ImageDTO:
    """Uploads an image"""
    if not file.content_type.startswith("image"):
        raise HTTPException(status_code=415, detail="Not an image")

    contents = await file.read()

    try:
        pil_image = Image.open(io.BytesIO(contents))
    except:
        # Error opening the image
        raise HTTPException(status_code=415, detail="Failed to read image")

    try:
        image_dto = ApiDependencies.invoker.services.images_new.create(
            pil_image,
            image_type,
            image_category,
        )

        response.status_code = 201
        response.headers["Location"] = image_dto.image_url

        return image_dto
    except Exception as e:
        raise HTTPException(status_code=500, detail="Failed to create image")


@images_router.delete("/{image_type}/{image_name}", operation_id="delete_image")
async def delete_image(
    image_type: ImageType = Query(description="The type of image to delete"),
    image_name: str = Path(description="The name of the image to delete"),
) -> None:
    """Deletes an image"""

    try:
        ApiDependencies.invoker.services.images_new.delete(image_type, image_name)
    except Exception as e:
        # TODO: Does this need any exception handling at all?
        pass


@images_router.get(
    "/{image_type}/{image_name}/record",
    operation_id="get_image_record",
    response_model=ImageDTO,
)
async def get_image_record(
    image_type: ImageType = Path(description="The type of the image record to get"),
    image_name: str = Path(description="The id of the image record to get"),
) -> ImageDTO:
    """Gets an image record by id"""

    try:
        return ApiDependencies.invoker.services.images_new.get_dto(
            image_type, image_name
        )
    except Exception as e:
        raise HTTPException(status_code=404)


@images_router.get("/{image_type}/{image_name}/image", operation_id="get_image")
async def get_image(
    image_type: ImageType = Path(description="The type of the image to get"),
    image_name: str = Path(description="The id of the image to get"),
) -> FileResponse:
    """Gets an image"""

    try:
        path = ApiDependencies.invoker.services.images_new.get_path(
            image_type, image_name
        )

        return FileResponse(path)
    except Exception as e:
        raise HTTPException(status_code=404)


@images_router.get("/{image_type}/{image_name}/thumbnail", operation_id="get_thumbnail")
async def get_thumbnail(
    image_type: ImageType = Path(
        description="The type of the image whose thumbnail to get"
    ),
    image_name: str = Path(description="The id of the image whose thumbnail to get"),
) -> FileResponse:
    """Gets a thumbnail"""

    try:
        path = ApiDependencies.invoker.services.images_new.get_path(
            image_type, image_name, thumbnail=True
        )

        return FileResponse(path)
    except Exception as e:
        raise HTTPException(status_code=404)


@images_router.get(
    "/{image_type}/{image_name}/urls",
    operation_id="get_image_urls",
    response_model=ImageUrlsDTO,
)
async def get_image_urls(
    image_type: ImageType = Path(description="The type of the image whose URL to get"),
    image_name: str = Path(description="The id of the image whose URL to get"),
) -> ImageUrlsDTO:
    """Gets an image and thumbnail URL"""

    try:
        image_url = ApiDependencies.invoker.services.images_new.get_url(
            image_type, image_name
        )
        thumbnail_url = ApiDependencies.invoker.services.images_new.get_url(
            image_type, image_name, thumbnail=True
        )
        return ImageUrlsDTO(
            image_type=image_type,
            image_name=image_name,
            image_url=image_url,
            thumbnail_url=thumbnail_url,
        )
    except Exception as e:
        raise HTTPException(status_code=404)


@images_router.get(
    "/",
    operation_id="list_image_records",
    response_model=PaginatedResults[ImageDTO],
)
async def list_image_records(
    image_type: ImageType = Query(description="The type of image records to get"),
    image_category: ImageCategory = Query(
        description="The kind of image records to get"
    ),
    page: int = Query(default=0, description="The page of image records to get"),
    per_page: int = Query(
        default=10, description="The number of image records per page"
    ),
) -> PaginatedResults[ImageDTO]:
    """Gets a list of image records by type and category"""

    image_dtos = ApiDependencies.invoker.services.images_new.get_many(
        image_type,
        image_category,
        page,
        per_page,
    )

    return image_dtos