Merge pull request #1057 from proteanx/veniceimage

feat: add venice.ai image generation
elizaOS · Dec 14, 2024 · 072e11b · 072e11b
2 parents 7ad8141 + c425ec8
commit 072e11b
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 4 deletions.
diff --git a/.env.example b/.env.example
@@ -230,6 +230,7 @@ VENICE_API_KEY=      # generate from venice settings
 SMALL_VENICE_MODEL=  # Default: llama-3.3-70b
 MEDIUM_VENICE_MODEL= # Default: llama-3.3-70b
 LARGE_VENICE_MODEL=  # Default: llama-3.1-405b
+IMAGE_VENICE_MODEL=  # Default: fluently-xl
 
 # fal.ai Configuration
 FAL_API_KEY=

diff --git a/agent/src/index.ts b/agent/src/index.ts
@@ -464,6 +464,7 @@ export async function createAgent(
                 : null,
             getSecret(character, "FAL_API_KEY") ||
             getSecret(character, "OPENAI_API_KEY") ||
+            getSecret(character, "VENICE_API_KEY") ||
             getSecret(character, "HEURIST_API_KEY")
                 ? imageGenerationPlugin
                 : null,

diff --git a/packages/core/src/generation.ts b/packages/core/src/generation.ts
@@ -954,7 +954,8 @@ export const generateImage = async (
             : (runtime.getSetting("HEURIST_API_KEY") ??
               runtime.getSetting("TOGETHER_API_KEY") ??
               runtime.getSetting("FAL_API_KEY") ??
-              runtime.getSetting("OPENAI_API_KEY"));
+              runtime.getSetting("OPENAI_API_KEY") ??
+              runtime.getSetting("VENICE_API_KEY"));
 
     try {
         if (runtime.imageModelProvider === ModelProviderName.HEURIST) {
@@ -1100,6 +1101,40 @@ export const generateImage = async (
             });
 
             const base64s = await Promise.all(base64Promises);
+            return { success: true, data: base64s };
+        } else if (runtime.imageModelProvider === ModelProviderName.VENICE) {
+            const response = await fetch(
+                "https://api.venice.ai/api/v1/image/generate",
+                {
+                    method: "POST",
+                    headers: {
+                        Authorization: `Bearer ${apiKey}`,
+                        "Content-Type": "application/json",
+                    },
+                    body: JSON.stringify({
+                        model: data.modelId || "fluently-xl",
+                        prompt: data.prompt,
+                        negative_prompt: data.negativePrompt,
+                        width: data.width || 1024,
+                        height: data.height || 1024,
+                        steps: data.numIterations || 20,
+                    }),
+                }
+            );
+
+            const result = await response.json();
+
+            if (!result.images || !Array.isArray(result.images)) {
+                throw new Error("Invalid response format from Venice AI");
+            }
+
+            const base64s = result.images.map((base64String) => {
+                if (!base64String) {
+                    throw new Error("Empty base64 string in Venice AI response");
+                }
+                return `data:image/png;base64,${base64String}`;
+            });
+
             return { success: true, data: base64s };
         } else {
             let targetSize = `${data.width}x${data.height}`;

diff --git a/packages/core/src/models.ts b/packages/core/src/models.ts
@@ -463,6 +463,7 @@ export const models: Models = {
             [ModelClass.SMALL]: settings.SMALL_VENICE_MODEL || "llama-3.3-70b",
             [ModelClass.MEDIUM]: settings.MEDIUM_VENICE_MODEL || "llama-3.3-70b",
             [ModelClass.LARGE]: settings.LARGE_VENICE_MODEL || "llama-3.1-405b",
+            [ModelClass.IMAGE]: settings.IMAGE_VENICE_MODEL || "fluently-xl",
         },
     },
 };

diff --git a/packages/plugin-image-generation/src/environment.ts b/packages/plugin-image-generation/src/environment.ts
@@ -8,6 +8,7 @@ export const imageGenEnvSchema = z
         HEURIST_API_KEY: z.string().optional(),
         FAL_API_KEY: z.string().optional(),
         OPENAI_API_KEY: z.string().optional(),
+        VENICE_API_KEY: z.string().optional(),
     })
     .refine(
         (data) => {
@@ -16,12 +17,13 @@ export const imageGenEnvSchema = z
                 data.TOGETHER_API_KEY ||
                 data.HEURIST_API_KEY ||
                 data.FAL_API_KEY ||
-                data.OPENAI_API_KEY
+                data.OPENAI_API_KEY ||
+                data.VENICE_API_KEY
             );
         },
         {
             message:
-                "At least one of ANTHROPIC_API_KEY, TOGETHER_API_KEY, HEURIST_API_KEY, FAL_API_KEY or OPENAI_API_KEY is required",
+                "At least one of ANTHROPIC_API_KEY, TOGETHER_API_KEY, HEURIST_API_KEY, FAL_API_KEY, OPENAI_API_KEY or VENICE_API_KEY is required",
         }
     );
 
@@ -46,6 +48,9 @@ export async function validateImageGenConfig(
             OPENAI_API_KEY:
                 runtime.getSetting("OPENAI_API_KEY") ||
                 process.env.OPENAI_API_KEY,
+            VENICE_API_KEY:
+                runtime.getSetting("VENICE_API_KEY") ||
+                process.env.VENICE_API_KEY,
         };
 
         return imageGenEnvSchema.parse(config);

diff --git a/packages/plugin-image-generation/src/index.ts b/packages/plugin-image-generation/src/index.ts
@@ -84,13 +84,15 @@ const imageGeneration: Action = {
         const heuristApiKeyOk = !!runtime.getSetting("HEURIST_API_KEY");
         const falApiKeyOk = !!runtime.getSetting("FAL_API_KEY");
         const openAiApiKeyOk = !!runtime.getSetting("OPENAI_API_KEY");
+        const veniceApiKeyOk = !!runtime.getSetting("VENICE_API_KEY");
 
         return (
             anthropicApiKeyOk ||
             togetherApiKeyOk ||
             heuristApiKeyOk ||
             falApiKeyOk ||
-            openAiApiKeyOk
+            openAiApiKeyOk ||
+            veniceApiKeyOk
         );
     },
     handler: async (