src/aamcq/generation/base.py

"""Generation backend Protocol."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Protocol

from PIL import Image


@dataclass(frozen=True)
class GenerationResult:
    image: Image.Image
    backend: str
    model_id: str
    prompt: str
    negative_prompt: str | None
    seed: int
    num_inference_steps: int
    guidance_scale: float
    height: int
    width: int


class GenerationBackend(Protocol):
    name: str

    def load(self) -> None: ...

    def generate(
        self,
        *,
        prompt: str,
        negative_prompt: str | None,
        seed: int,
        height: int,
        width: int,
    ) -> GenerationResult: ...