[feat] add model to replicate

edenai · Oct 30, 2023 · 8554dc6 · 8554dc6
1 parent 81d6ca9
commit 8554dc6
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 16 deletions.
diff --git a/edenai_apis/apis/replicate/config.py b/edenai_apis/apis/replicate/config.py
@@ -1,4 +1,9 @@
 get_model_id = {
     "llama-2-70b" : "14ce4448d5e7e9ed0c37745ac46eca157aab09061f0c179ac2b323b5de56552b",
     "llama-2-70b-chat" :"58d078176e02c219e11eb4da5a02a7830a283b14cf8f94537af893ccff5ee781"
+}
+get_model_id_image = {
+    "anime-style" : "09a5805203f4c12da649ec1923bb7729517ca25fcac790e640eaa9ed66573b65",
+    "classic" : "c0259010b93e7a4102a4ba946d70e06d7d0c7dc007201af443cfc8f943ab1d3c",
+    "vintedois-diffusion" : "28cea91bdfced0e2dc7fda466cc0a46501c0edc84905b2120ea02e0707b967fd",
 }
diff --git a/edenai_apis/apis/replicate/info.json b/edenai_apis/apis/replicate/info.json
@@ -6,7 +6,13 @@
                   "256x256",
                   "512x512",
                   "1024x1024"
-                ]
+                ],
+                "models": [
+                  "anime-style",
+                  "vintedois-diffusion",
+                  "classic"
+                ],
+                "default_model" :  "classic"
               },
             "version" : "v1"
         }

diff --git a/edenai_apis/apis/replicate/outputs/image/generation_output.json b/edenai_apis/apis/replicate/outputs/image/generation_output.json
diff --git a/edenai_apis/apis/replicate/replicate_api.py b/edenai_apis/apis/replicate/replicate_api.py
@@ -16,7 +16,7 @@
 from edenai_apis.utils.exception import ProviderException
 from edenai_apis.utils.types import ResponseType
 import base64
-from .config import get_model_id
+from .config import get_model_id, get_model_id_image
 
 
 class ReplicateApi(ProviderInterface, ImageInterface, TextInterface):
@@ -96,7 +96,6 @@ def __get_response(
         status = response_dict["status"]
         while status != "succeeded":
             response = requests.get(url_get_response, headers=self.headers)
-
             try:
                 response_dict = response.json()
             except requests.JSONDecodeError:
@@ -120,17 +119,21 @@ def image__generation(
     ) -> ResponseType[GenerationDataClass]:
         url = f"{self.base_url}/predictions"
         size = resolution.split("x")
+        version = get_model_id_image[model]
+
         payload = {
             "input": {
                 "prompt": text,
                 "width": int(size[0]),
                 "height": int(size[1]),
             },
-            "version": "c0259010b93e7a4102a4ba946d70e06d7d0c7dc007201af443cfc8f943ab1d3c",
+            "version": version,
         }
 
         response_dict = ReplicateApi.__get_response(self, url, payload)
         image_url = response_dict.get("output")
+        if isinstance(image_url, list):
+            image_url = image_url[0]
         image_bytes = base64.b64encode(requests.get(image_url).content)
 
         return ResponseType[GenerationDataClass](