Feat/model as tool (#2744)

api/core/model_runtime/entities/model_entities.py

@@ -17,7 +17,7 @@ class ModelType(Enum):
     SPEECH2TEXT = "speech2text"
     MODERATION = "moderation"
     TTS = "tts"
-    # TEXT2IMG = "text2img"
+    TEXT2IMG = "text2img"
 
     @classmethod
     def value_of(cls, origin_model_type: str) -> "ModelType":
@@ -36,6 +36,8 @@ class ModelType(Enum):
             return cls.SPEECH2TEXT
         elif origin_model_type == 'tts' or origin_model_type == cls.TTS.value:
             return cls.TTS
+        elif origin_model_type == 'text2img' or origin_model_type == cls.TEXT2IMG.value:
+            return cls.TEXT2IMG
         elif origin_model_type == cls.MODERATION.value:
             return cls.MODERATION
         else:
@@ -59,10 +61,11 @@ class ModelType(Enum):
             return 'tts'
         elif self == self.MODERATION:
             return 'moderation'
+        elif self == self.TEXT2IMG:
+            return 'text2img'
         else:
             raise ValueError(f'invalid model type {self}')
 
-
 class FetchFrom(Enum):
     """
     Enum class for fetch from.

api/core/model_runtime/model_providers/__base/text2img_model.py

@@ -0,0 +1,48 @@
+from abc import abstractmethod
+from typing import IO, Optional
+
+from core.model_runtime.entities.model_entities import ModelType
+from core.model_runtime.model_providers.__base.ai_model import AIModel
+
+
+class Text2ImageModel(AIModel):
+    """
+    Model class for text2img model.
+    """
+    model_type: ModelType = ModelType.TEXT2IMG
+
+    def invoke(self, model: str, credentials: dict, prompt: str, 
+               model_parameters: dict, user: Optional[str] = None) \
+            -> list[IO[bytes]]:
+        """
+        Invoke Text2Image model
+
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt: prompt for image generation
+        :param model_parameters: model parameters
+        :param user: unique user id
+
+        :return: image bytes
+        """
+        try:
+            return self._invoke(model, credentials, prompt, model_parameters, user)
+        except Exception as e:
+            raise self._transform_invoke_error(e)
+
+    @abstractmethod
+    def _invoke(self, model: str, credentials: dict, prompt: str, 
+                model_parameters: dict, user: Optional[str] = None) \
+            -> list[IO[bytes]]:
+        """
+        Invoke Text2Image model
+
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt: prompt for image generation
+        :param model_parameters: model parameters
+        :param user: unique user id
+
+        :return: image bytes
+        """
+        raise NotImplementedError