Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Enhance Heurist Image Generation Settings and Image Handling #375

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions packages/core/src/generation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}),
}
);
Expand All @@ -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
) {
Expand Down
47 changes: 43 additions & 4 deletions packages/plugin-image-generation/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
Loading