diff options
| author | Tekky <98614666+xtekky@users.noreply.github.com> | 2024-09-07 21:37:24 +0200 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-09-07 21:37:24 +0200 |
| commit | 07fa87b4d180259d1da86afb565e14ac3d60d50b (patch) | |
| tree | d474a2bf8bd79cf94bfa48cbaca3917659dd19be /g4f/client/image_models.py | |
| parent | Merge pull request #2206 from Parthsadaria/patch-1 (diff) | |
| parent | g4f/models.py g4f/Provider/MagickPen.py (diff) | |
| download | gpt4free-0.3.2.6.tar gpt4free-0.3.2.6.tar.gz gpt4free-0.3.2.6.tar.bz2 gpt4free-0.3.2.6.tar.lz gpt4free-0.3.2.6.tar.xz gpt4free-0.3.2.6.tar.zst gpt4free-0.3.2.6.zip | |
Diffstat (limited to 'g4f/client/image_models.py')
| -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 |