mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
|
|
|
|
from fastapi import Body, HTTPException, Path, Response
|
|
from fastapi.routing import APIRouter
|
|
|
|
from invokeai.app.services.batch_manager_storage import BatchSession, BatchSessionNotFoundException
|
|
|
|
# Importing * is bad karma but needed here for node detection
|
|
from ...invocations import * # noqa: F401 F403
|
|
from ...services.batch_manager import Batch, BatchProcessResponse
|
|
from ...services.graph import Graph
|
|
from ..dependencies import ApiDependencies
|
|
|
|
batches_router = APIRouter(prefix="/v1/batches", tags=["sessions"])
|
|
|
|
|
|
@batches_router.post(
|
|
"/",
|
|
operation_id="create_batch",
|
|
responses={
|
|
200: {"model": BatchProcessResponse},
|
|
400: {"description": "Invalid json"},
|
|
},
|
|
)
|
|
async def create_batch(
|
|
graph: Graph = Body(description="The graph to initialize the session with"),
|
|
batch: Batch = Body(description="Batch config to apply to the given graph"),
|
|
) -> BatchProcessResponse:
|
|
"""Creates a batch process"""
|
|
return ApiDependencies.invoker.services.batch_manager.create_batch_process(batch, graph)
|
|
|
|
|
|
@batches_router.put(
|
|
"/b/{batch_process_id}/invoke",
|
|
operation_id="start_batch",
|
|
responses={
|
|
202: {"description": "Batch process started"},
|
|
404: {"description": "Batch session not found"},
|
|
},
|
|
)
|
|
async def start_batch(
|
|
batch_process_id: str = Path(description="ID of Batch to start"),
|
|
) -> Response:
|
|
"""Executes a batch process"""
|
|
try:
|
|
ApiDependencies.invoker.services.batch_manager.run_batch_process(batch_process_id)
|
|
return Response(status_code=202)
|
|
except BatchSessionNotFoundException:
|
|
raise HTTPException(status_code=404, detail="Batch session not found")
|
|
|
|
|
|
@batches_router.delete(
|
|
"/b/{batch_process_id}",
|
|
operation_id="cancel_batch",
|
|
responses={202: {"description": "The batch is canceled"}},
|
|
)
|
|
async def cancel_batch(
|
|
batch_process_id: str = Path(description="The id of the batch process to cancel"),
|
|
) -> Response:
|
|
"""Cancels a batch process"""
|
|
ApiDependencies.invoker.services.batch_manager.cancel_batch_process(batch_process_id)
|
|
return Response(status_code=202)
|
|
|
|
|
|
@batches_router.get(
|
|
"/incomplete",
|
|
operation_id="list_incomplete_batches",
|
|
responses={200: {"model": list[BatchProcessResponse]}},
|
|
)
|
|
async def list_incomplete_batches() -> list[BatchProcessResponse]:
|
|
"""Lists incomplete batch processes"""
|
|
return ApiDependencies.invoker.services.batch_manager.get_incomplete_batch_processes()
|
|
|
|
|
|
@batches_router.get(
|
|
"/",
|
|
operation_id="list_batches",
|
|
responses={200: {"model": list[BatchProcessResponse]}},
|
|
)
|
|
async def list_batches() -> list[BatchProcessResponse]:
|
|
"""Lists all batch processes"""
|
|
return ApiDependencies.invoker.services.batch_manager.get_batch_processes()
|
|
|
|
|
|
@batches_router.get(
|
|
"/b/{batch_process_id}",
|
|
operation_id="get_batch",
|
|
responses={200: {"model": BatchProcessResponse}},
|
|
)
|
|
async def get_batch(
|
|
batch_process_id: str = Path(description="The id of the batch process to get"),
|
|
) -> BatchProcessResponse:
|
|
"""Gets a Batch Process"""
|
|
return ApiDependencies.invoker.services.batch_manager.get_batch(batch_process_id)
|
|
|
|
|
|
@batches_router.get(
|
|
"/b/{batch_process_id}/sessions",
|
|
operation_id="get_batch_sessions",
|
|
responses={200: {"model": list[BatchSession]}},
|
|
)
|
|
async def get_batch_sessions(
|
|
batch_process_id: str = Path(description="The id of the batch process to get"),
|
|
) -> list[BatchSession]:
|
|
"""Gets a list of batch sessions for a given batch process"""
|
|
return ApiDependencies.invoker.services.batch_manager.get_sessions(batch_process_id)
|