From 24d4ea055af77d312806b8cbad606ab4889227bb Mon Sep 17 00:00:00 2001
From: kqlio67 <kqlio67@users.noreply.github.com>
Date: Fri, 6 Sep 2024 01:34:12 +0300
Subject: [PATCH] refactor(image-models): replace providers with ModelUtils

---
 g4f/client/image_models.py | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/g4f/client/image_models.py b/g4f/client/image_models.py
index db2ce09a2fb..edaa4592242 100644
--- a/g4f/client/image_models.py
+++ b/g4f/client/image_models.py
@@ -2,18 +2,15 @@
 
 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