imageforge/backend/app/providers/factory.py

from __future__ import annotations

from .a1111_provider import A1111Provider
from .diffusion_provider import DiffusionProvider
from .dummy_provider import DummyProvider
from .interface import IImageProvider
from .localai_provider import LocalAIProvider
from .zimageturbo_provider import ZImageTurboProvider


class ProviderRegistry:
    def __init__(self) -> None:
        self._providers: dict[str, IImageProvider] = {
            "dummy": DummyProvider(),
            "localai": LocalAIProvider(),
            "diffusion": DiffusionProvider(),
            "a1111": A1111Provider(),
            "zimageturbo": ZImageTurboProvider(),
        }

    def get(self, provider_id: str) -> IImageProvider:
        provider = self._providers.get(provider_id)
        if provider is None:
            raise KeyError(f"Unknown provider: {provider_id}")
        return provider

    def list(self) -> list[IImageProvider]:
        return list(self._providers.values())