refactor(image-models): replace providers with ModelUtils

This commit is contained in:
kqlio67 2024-09-06 01:34:12 +03:00
parent 44ba42f58b
commit 24d4ea055a

View file

@ -2,18 +2,15 @@ from __future__ import annotations
from .types import Client, ImageProvider
from ..Provider.BingCreateImages import BingCreateImages
from ..Provider.needs_auth import Gemini, OpenaiChat
from ..Provider.You import You
from ..models import ModelUtils
class ImageModels():
gemini = Gemini
openai = OpenaiChat
you = You
def __init__(self, client: Client) -> None:
def __init__(self, client):
self.client = client
self.default = BingCreateImages(proxy=self.client.get_proxy())
self.models = ModelUtils.convert
def get(self, name: str, default: ImageProvider = None) -> ImageProvider:
return getattr(self, name) if hasattr(self, name) else default or self.default
def get(self, name, default=None):
model = self.models.get(name)
if model and model.best_provider:
return model.best_provider
return default