InvokeAI/invokeai/app/services/batch_manager.py

import networkx as nx
import copy

from abc import ABC, abstractmethod
from itertools import product
from pydantic import BaseModel, Field
from fastapi_events.handlers.local import local_handler
from fastapi_events.typing import Event

from invokeai.app.services.events import EventServiceBase
from invokeai.app.services.graph import Graph, GraphExecutionState
from invokeai.app.services.invoker import Invoker
from invokeai.app.services.batch_manager_storage import (
    BatchProcessStorageBase,
    BatchSessionNotFoundException,
    Batch,
    BatchProcess,
    BatchSession,
    BatchSessionChanges,
)


class BatchProcessResponse(BaseModel):
    batch_id: str = Field(description="ID for the batch")
    session_ids: list[str] = Field(description="List of session IDs created for this batch")


class BatchManagerBase(ABC):
    @abstractmethod
    def start(self, invoker: Invoker):
        pass

    @abstractmethod
    def create_batch_process(self, batches: list[Batch], graph: Graph) -> BatchProcessResponse:
        pass

    @abstractmethod
    def run_batch_process(self, batch_id: str):
        pass

    @abstractmethod
    def cancel_batch_process(self, batch_process_id: str):
        pass


class BatchManager(BatchManagerBase):
    """Responsible for managing currently running and scheduled batch jobs"""

    __invoker: Invoker
    __batch_process_storage: BatchProcessStorageBase

    def __init__(self, batch_process_storage: BatchProcessStorageBase) -> None:
        super().__init__()
        self.__batch_process_storage = batch_process_storage

    def start(self, invoker: Invoker) -> None:
        self.__invoker = invoker
        local_handler.register(event_name=EventServiceBase.session_event, _func=self.on_event)

    async def on_event(self, event: Event):
        event_name = event[1]["event"]

        match event_name:
            case "graph_execution_state_complete":
                await self.process(event, False)
            case "invocation_error":
                await self.process(event, True)

        return event

    async def process(self, event: Event, err: bool):
        data = event[1]["data"]
        batch_session = self.__batch_process_storage.get_session(data["graph_execution_state_id"])
        if not batch_session:
            return
        updateSession = BatchSessionChanges(state="error" if err else "completed")
        batch_session = self.__batch_process_storage.update_session_state(
            batch_session.batch_id,
            batch_session.session_id,
            updateSession,
        )
        batch_process = self.__batch_process_storage.get(batch_session.batch_id)
        if not batch_process.canceled:
            self.run_batch_process(batch_process.batch_id)

    def _create_batch_session(self, batch_process: BatchProcess, batch_indices: list[int]) -> GraphExecutionState:
        graph = batch_process.graph.copy(deep=True)
        batches = batch_process.batches
        g = graph.nx_graph_flat()
        sorted_nodes = nx.topological_sort(g)
        for npath in sorted_nodes:
            node = graph.get_node(npath)
            (index, batch) = next(((i, b) for i, b in enumerate(batches) if b.node_id in node.id), (None, None))
            if batch:
                batch_index = batch_indices[index]
                datum = batch.data[batch_index]
                for key in datum:
                    node.__dict__[key] = datum[key]
                graph.update_node(npath, node)

        return GraphExecutionState(graph=graph)

    def run_batch_process(self, batch_id: str):
        self.__batch_process_storage.start(batch_id)
        try:
            created_session = self.__batch_process_storage.get_created_session(batch_id)
        except BatchSessionNotFoundException:
            return
        ges = self.__invoker.services.graph_execution_manager.get(created_session.session_id)
        self.__invoker.invoke(ges, invoke_all=True)

    def _valid_batch_config(self, batch_process: BatchProcess) -> bool:
        # TODO: Check that the node_ids in the batches are unique
        # TODO: Validate data types are correct for each batch data
        return True

    def create_batch_process(self, batches: list[Batch], graph: Graph) -> BatchProcessResponse:
        batch_process = BatchProcess(
            batches=batches,
            graph=graph,
        )
        if not self._valid_batch_config(batch_process):
            return None
        batch_process = self.__batch_process_storage.save(batch_process)
        sessions = self._create_sessions(batch_process)
        return BatchProcessResponse(
            batch_id=batch_process.batch_id,
            session_ids=[session.session_id for session in sessions],
        )

    def _create_sessions(self, batch_process: BatchProcess) -> list[BatchSession]:
        batch_indices = list()
        sessions = list()
        for batch in batch_process.batches:
            batch_indices.append(list(range(len(batch.data))))
        all_batch_indices = product(*batch_indices)
        for bi in all_batch_indices:
            ges = self._create_batch_session(batch_process, bi)
            self.__invoker.services.graph_execution_manager.set(ges)
            batch_session = BatchSession(batch_id=batch_process.batch_id, session_id=ges.id, state="created")
            sessions.append(self.__batch_process_storage.create_session(batch_session))
        return sessions

    def cancel_batch_process(self, batch_process_id: str):
        self.__batch_process_storage.cancel(batch_process_id)