From a26421bcd8b8580650fa05b3fb4f8fdfa0ef9921 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <hlohaus@users.noreply.github.com>
Date: Sun, 21 Apr 2024 15:15:55 +0200
Subject: Add image model list

---
 g4f/gui/server/api.py     | 28 ++++++++++++++++++++++++----
 g4f/gui/server/backend.py |  4 ++++
 2 files changed, 28 insertions(+), 4 deletions(-)

(limited to 'g4f/gui')

diff --git a/g4f/gui/server/api.py b/g4f/gui/server/api.py
index 47c25eb7..435700ea 100644
--- a/g4f/gui/server/api.py
+++ b/g4f/gui/server/api.py
@@ -16,7 +16,8 @@ conversations: dict[dict[str, BaseConversation]] = {}
 
 class Api():
 
-    def get_models(self) -> list[str]:
+    @staticmethod
+    def get_models() -> list[str]:
         """
         Return a list of all models.
 
@@ -27,7 +28,8 @@ class Api():
         """
         return models._all_models
 
-    def get_provider_models(self, provider: str) -> list[dict]:
+    @staticmethod
+    def get_provider_models(provider: str) -> list[dict]:
         if provider in __map__:
             provider: ProviderType = __map__[provider]
             if issubclass(provider, ProviderModelMixin):
@@ -40,7 +42,24 @@ class Api():
             else:
                 return [];
 
-    def get_providers(self) -> list[str]:
+    @staticmethod
+    def get_image_models() -> list[dict]:
+        image_models = []
+        for key, provider in __map__.items():
+            if hasattr(provider, "image_models"):
+                if hasattr(provider, "get_models"):
+                    provider.get_models()
+                for model in provider.image_models:
+                    image_models.append({
+                        "provider": key,
+                        "url": provider.url,
+                        "label": provider.label if hasattr(provider, "label") else None,
+                        "image_model": model
+                    })
+        return image_models
+
+    @staticmethod
+    def get_providers() -> list[str]:
         """
         Return a list of all working providers.
         """
@@ -58,7 +77,8 @@ class Api():
             if provider.working
         }
 
-    def get_version(self):
+    @staticmethod
+    def get_version():
         """
         Returns the current and latest version of the application.
 
diff --git a/g4f/gui/server/backend.py b/g4f/gui/server/backend.py
index d30b97d9..d9e31c0e 100644
--- a/g4f/gui/server/backend.py
+++ b/g4f/gui/server/backend.py
@@ -31,6 +31,10 @@ class Backend_Api(Api):
                 'function': self.get_provider_models,
                 'methods': ['GET']
             },
+            '/backend-api/v2/image_models': {
+                'function': self.get_image_models,
+                'methods': ['GET']
+            },
             '/backend-api/v2/providers': {
                 'function': self.get_providers,
                 'methods': ['GET']
-- 
cgit v1.2.3