update Heurist image gen settings

ai16z · Nov 17, 2024 · 00eb3ef · 00eb3ef
1 parent 04de0f4
commit 00eb3ef
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 7 deletions.
diff --git a/packages/core/src/generation.ts b/packages/core/src/generation.ts
@@ -767,11 +767,13 @@ export const generateImage = async (
                                 num_iterations: data.numIterations || 20,
                                 width: data.width || 512,
                                 height: data.height || 512,
-                                guidance_scale: data.guidanceScale,
+                                guidance_scale: data.guidanceScale || 3,
                                 seed: data.seed || -1,
                             },
                         },
                         model_id: data.modelId || "FLUX.1-dev",
+                        deadline: 60,
+                        priority: 1,
                     }),
                 }
             );
@@ -782,8 +784,8 @@ export const generateImage = async (
                 );
             }
 
-            const result = await response.json();
-            return { success: true, data: [result.url] };
+            const imageURL = await response.json();
+            return { success: true, data: [imageURL] };
         } else if (
             runtime.character.modelProvider === ModelProviderName.LLAMACLOUD
         ) {

diff --git a/packages/plugin-image-generation/src/index.ts b/packages/plugin-image-generation/src/index.ts
@@ -34,17 +34,52 @@ export function saveBase64Image(base64Data: string, filename: string): string {
     return filepath;
 }
 
+export async function saveHeuristImage(imageUrl: string, filename: string): Promise<string> {
+    const imageDir = path.join(process.cwd(), "generatedImages");
+    if (!fs.existsSync(imageDir)) {
+        fs.mkdirSync(imageDir, { recursive: true });
+    }
+
+    // Fetch image from URL
+    const response = await fetch(imageUrl);
+    if (!response.ok) {
+        throw new Error(`Failed to fetch image: ${response.statusText}`);
+    }
+
+    const arrayBuffer = await response.arrayBuffer();
+    const imageBuffer = Buffer.from(arrayBuffer);
+
+    // Create full file path
+    const filepath = path.join(imageDir, `${filename}.png`);
+
+    // Save the file
+    fs.writeFileSync(filepath, imageBuffer);
+
+    return filepath;
+}
+
 const imageGeneration: Action = {
     name: "GENERATE_IMAGE",
-    similes: ["IMAGE_GENERATION", "IMAGE_GEN", "CREATE_IMAGE", "MAKE_PICTURE"],
+    similes: [
+        "IMAGE_GENERATION", 
+        "IMAGE_GEN", 
+        "CREATE_IMAGE", 
+        "MAKE_PICTURE",
+        "GENERATE_IMAGE",
+        "GENERATE_A",
+        "DRAW",
+        "DRAW_A",
+        "MAKE_A"
+    ],
     description: "Generate an image to go along with the message.",
     validate: async (runtime: IAgentRuntime, message: Memory) => {
         const anthropicApiKeyOk = !!runtime.getSetting("ANTHROPIC_API_KEY");
         const togetherApiKeyOk = !!runtime.getSetting("TOGETHER_API_KEY");
+        const heuristApiKeyOk = !!runtime.getSetting("HEURIST_API_KEY");
 
         // TODO: Add openai DALL-E generation as well
 
-        return anthropicApiKeyOk && togetherApiKeyOk;
+        return anthropicApiKeyOk || togetherApiKeyOk || heuristApiKeyOk;
     },
     handler: async (
         runtime: IAgentRuntime,
@@ -84,10 +119,14 @@ const imageGeneration: Action = {
             for (let i = 0; i < images.data.length; i++) {
                 const image = images.data[i];
 
-                const base64Image = images.data[i];
                 // Save the image and get filepath
                 const filename = `generated_${Date.now()}_${i}`;
-                const filepath = saveBase64Image(base64Image, filename);
+
+                // Choose save function based on image data format
+                const filepath = image.startsWith('http') 
+                    ? await saveHeuristImage(image, filename)
+                    : saveBase64Image(image, filename);
+
                 elizaLogger.log(`Processing image ${i + 1}:`, filename);
 
                 //just dont even add a caption or a description just have it generate & send