diff options
| -rw-r--r-- | g4f/client/image_models.py | 19 |
1 files changed, 8 insertions, 11 deletions
diff --git a/g4f/client/image_models.py b/g4f/client/image_models.py index db2ce09a..edaa4592 100644 --- a/g4f/client/image_models.py +++ b/g4f/client/image_models.py @@ -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 |