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: add Ideogram.ai image generation support (CONFLICTED) #3059

Closed
wants to merge 1 commit into from
Closed
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: 8 additions & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -349,6 +349,14 @@ ALI_BAILIAN_API_KEY= # Ali Bailian API Key
NANOGPT_API_KEY= # NanoGPT API Key
TOGETHER_API_KEY= # Together API Key

# Ideogram.ai Configuration
IDEOGRAM_API_KEY= # Ideogram.ai API key. Get itfrom https://developer.ideogram.ai/ideogram-api/api-setup
IDEOGRAM_MODEL= # Which model to use. Defaults to V_2
IDEOGRAM_MAGIC_PROMPT= # If it should use magic prompt or not. Defaults to AUTO
IDEOGRAM_STYLE_TYPE= # Which style to use. Defaults to AUTO. Check https://developer.ideogram.ai/api-reference/api-reference/generate for possible values
IDEOGRAM_COLOR_PALETTE= # Name of the color palette to use. Defaults to auto. Check https://developer.ideogram.ai/api-reference/api-reference/generate for possible values


######################################
#### Crypto Plugin Configurations ####
######################################
Expand Down
220 changes: 220 additions & 0 deletions packages/core/__tests__/imageGeneration.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,220 @@
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import { generateImage } from "../src/generation";
import { ModelProviderName } from "../src/types";
import type { AgentRuntime } from "../src/runtime";

describe("Image Generation", () => {
let mockRuntime: AgentRuntime;

beforeEach(() => {
global.fetch = vi.fn();
mockRuntime = {
imageModelProvider: ModelProviderName.IDEOGRAM,
getSetting: vi.fn((key: string) => {
switch (key) {
case "IDEOGRAM_API_KEY":
return "test-api-key";
case "IDEOGRAM_MAGIC_PROMPT":
return "auto";
case "IDEOGRAM_STYLE_TYPE":
return "auto";
case "IDEOGRAM_COLOR_PALETTE":
return "vibrant";
case "IDEOGRAM_MODEL":
return "V_2";
default:
return undefined;
}
}),
} as unknown as AgentRuntime;
});

afterEach(() => {
vi.resetAllMocks();
});

describe("Ideogram.ai Integration", () => {
it("should generate an image successfully", async () => {
const mockImageUrl = "https://example.com/image.jpg";
const mockBase64 = "base64-encoded-image";

// Mock the initial API call
(global.fetch as ReturnType<typeof vi.fn>).mockImplementationOnce(() =>
Promise.resolve({
ok: true,
json: () =>
Promise.resolve({
data: [{ url: mockImageUrl }],
}),
})
);

// Mock the image fetch call
(global.fetch as ReturnType<typeof vi.fn>).mockImplementationOnce(() =>
Promise.resolve({
ok: true,
blob: () =>
Promise.resolve(
new Blob(["mock-image-data"], {
type: "image/jpeg",
})
),
})
);

const result = await generateImage(
{
prompt: "A beautiful sunset",
width: 1024,
height: 1024,
count: 1,
negativePrompt: "blur, dark",
seed: 12345,
},
mockRuntime
);

expect(result.success).toBe(true);
expect(result.data).toBeDefined();
expect(result.data?.length).toBe(1);
expect(result.data?.[0]).toContain("data:image/jpeg;base64,");

// Verify the API call
expect(fetch).toHaveBeenCalledTimes(2);
expect(fetch).toHaveBeenNthCalledWith(
1,
"https://api.ideogram.ai/generate",
{
method: "POST",
headers: {
"Api-Key": "test-api-key",
"Content-Type": "application/json",
},
body: JSON.stringify({
image_request: {
prompt: "A beautiful sunset",
model: "V_2",
magic_prompt_option: "AUTO",
style_type: "AUTO",
resolution: "RESOLUTION_1024_1024",
num_images: 1,
negative_prompt: "blur, dark",
color_palette: {
name: "VIBRANT"
},
seed: 12345,
},
}),
}
);
});

it("should handle API errors gracefully", async () => {
(global.fetch as ReturnType<typeof vi.fn>).mockImplementationOnce(() =>
Promise.resolve({
ok: false,
statusText: "Bad Request",
json: () =>
Promise.resolve({
error: "Invalid request",
}),
})
);

const result = await generateImage(
{
prompt: "A beautiful sunset",
width: 1024,
height: 1024,
},
mockRuntime
);

expect(result.success).toBe(false);
expect(result.error).toBeDefined();
expect(result.error?.message).toMatch(/Failed to generate image/);
});

it("should handle empty response data", async () => {
(global.fetch as ReturnType<typeof vi.fn>).mockImplementationOnce(() =>
Promise.resolve({
ok: true,
json: () =>
Promise.resolve({
data: [],
}),
})
);

const result = await generateImage(
{
prompt: "A beautiful sunset",
width: 1024,
height: 1024,
},
mockRuntime
);

expect(result.success).toBe(false);
expect(result.error).toBeDefined();
expect(result.error?.message).toBe("No images generated");
});

it("should handle missing image URL", async () => {
(global.fetch as ReturnType<typeof vi.fn>).mockImplementationOnce(() =>
Promise.resolve({
ok: true,
json: () =>
Promise.resolve({
data: [{ }],
}),
})
);

const result = await generateImage(
{
prompt: "A beautiful sunset",
width: 1024,
height: 1024,
},
mockRuntime
);

expect(result.success).toBe(false);
expect(result.error).toBeDefined();
expect(result.error?.message).toBe("Empty base64 string in Ideogram AI response");
});

it("should handle image fetch errors", async () => {
(global.fetch as ReturnType<typeof vi.fn>).mockImplementationOnce(() =>
Promise.resolve({
ok: true,
json: () =>
Promise.resolve({
data: [{ url: "https://example.com/image.jpg" }],
}),
})
);

(global.fetch as ReturnType<typeof vi.fn>).mockImplementationOnce(() =>
Promise.resolve({
ok: false,
statusText: "Not Found",
})
);

const result = await generateImage(
{
prompt: "A beautiful sunset",
width: 1024,
height: 1024,
},
mockRuntime
);

expect(result.success).toBe(false);
expect(result.error).toBeDefined();
expect(result.error?.message).toBe("Failed to fetch image: Not Found");
});
});
});
72 changes: 72 additions & 0 deletions packages/core/src/generation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1669,9 +1669,12 @@ export const generateImage = async (
return runtime.getSetting("VENICE_API_KEY");
case ModelProviderName.LIVEPEER:
return runtime.getSetting("LIVEPEER_GATEWAY_URL");
case ModelProviderName.IDEOGRAM:
return runtime.getSetting("IDEOGRAM_API_KEY");
default:
// If no specific match, try the fallback chain
return (
runtime.getSetting("IDEOGRAM_API_KEY") ??
runtime.getSetting("HEURIST_API_KEY") ??
runtime.getSetting("NINETEEN_AI_API_KEY") ??
runtime.getSetting("TOGETHER_API_KEY") ??
Expand Down Expand Up @@ -1911,6 +1914,75 @@ export const generateImage = async (
});

return { success: true, data: base64s };
} else if (runtime.imageModelProvider === ModelProviderName.IDEOGRAM) {
let body: Record<string, any> = {
image_request: {
prompt: data.prompt,
model: model,
magic_prompt_option: (runtime.getSetting("IDEOGRAM_MAGIC_PROMPT") || "auto").toUpperCase(),
style_type: (runtime.getSetting("IDEOGRAM_STYLE_TYPE") || "auto").toUpperCase(),
resolution: `RESOLUTION_${data.width}_${data.height}`,
num_images: data.count || 1,
},
}
if (data.negativePrompt) {
body.image_request.negative_prompt = data.negativePrompt;
}
if (runtime.getSetting("IDEOGRAM_COLOR_PALETTE")) {
body.image_request.color_palette = {
name: runtime.getSetting("IDEOGRAM_COLOR_PALETTE").toUpperCase()
}
}
if (data.seed) {
body.image_request.seed = data.seed;
}

let jsonBody = JSON.stringify(body);
const response = await fetch(
"https://api.ideogram.ai/generate",
{
method: "POST",
headers: {
"Api-Key": apiKey,
"Content-Type": "application/json",
},
body: jsonBody,
},
);
const result = await response.json();

if (!response.ok) {
throw new Error(`Failed to generate image with body ${jsonBody}: ${response.statusText} : ${result}`);
}

if (!result.data?.length) {
throw new Error("No images generated");
}

const base64Images = await Promise.all(
result.data.map(async (image) => {
if (!image.url) {
throw new Error(
"Empty base64 string in Nineteen AI response"
);
}
elizaLogger.debug(`Image URL: ${image.url}`);

const imageResponse = await fetch(image.url);
if (!imageResponse.ok) {
throw new Error(
`Failed to fetch image: ${imageResponse.statusText}`
);
}

const blob = await imageResponse.blob();
const arrayBuffer = await blob.arrayBuffer();
const base64 = Buffer.from(arrayBuffer).toString("base64");
return `data:image/jpeg;base64,${base64}`;
})
);

return { success: true, data: base64Images };
} else if (runtime.imageModelProvider === ModelProviderName.LIVEPEER) {
if (!apiKey) {
throw new Error("Livepeer Gateway is not defined");
Expand Down
6 changes: 6 additions & 0 deletions packages/core/src/models.ts
Original file line number Diff line number Diff line change
Expand Up @@ -669,6 +669,12 @@ export const models: Models = {
[ModelClass.IMAGE]: { name: "fal-ai/flux-lora", steps: 28 },
},
},
[ModelProviderName.IDEOGRAM]: {
endpoint: "https://api.ideogram.ai",
model: {
[ModelClass.IMAGE]: { name: settings.IDEOGRAM_MODEL || "V_2"},
},
},
[ModelProviderName.GAIANET]: {
endpoint: settings.GAIANET_SERVER_URL,
model: {
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,7 @@ export type Models = {
[ModelProviderName.INFERA]: Model;
[ModelProviderName.BEDROCK]: Model;
[ModelProviderName.ATOMA]: Model;
[ModelProviderName.IDEOGRAM]: Model;
};

/**
Expand Down Expand Up @@ -272,6 +273,7 @@ export enum ModelProviderName {
INFERA = "infera",
BEDROCK = "bedrock",
ATOMA = "atoma",
IDEOGRAM = "ideogram",
}

/**
Expand Down
7 changes: 6 additions & 1 deletion packages/plugin-image-generation/src/environment.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import { z } from "zod";

export const imageGenEnvSchema = z
.object({
IDEOGRAM_API_KEY: z.string().optional(),
ANTHROPIC_API_KEY: z.string().optional(),
NINETEEN_AI_API_KEY: z.string().optional(),
TOGETHER_API_KEY: z.string().optional(),
Expand All @@ -15,6 +16,7 @@ export const imageGenEnvSchema = z
.refine(
(data) => {
return !!(
data.IDEOGRAM_API_KEY ||
data.ANTHROPIC_API_KEY ||
data.NINETEEN_AI_API_KEY ||
data.TOGETHER_API_KEY ||
Expand All @@ -27,7 +29,7 @@ export const imageGenEnvSchema = z
},
{
message:
"At least one of ANTHROPIC_API_KEY, NINETEEN_AI_API_KEY, TOGETHER_API_KEY, HEURIST_API_KEY, FAL_API_KEY, OPENAI_API_KEY, VENICE_API_KEY or LIVEPEER_GATEWAY_URL is required",
"At least one of IDEOGRAM_API_KEY, ANTHROPIC_API_KEY, NINETEEN_AI_API_KEY, TOGETHER_API_KEY, HEURIST_API_KEY, FAL_API_KEY, OPENAI_API_KEY, VENICE_API_KEY or LIVEPEER_GATEWAY_URL is required",
}
);

Expand Down Expand Up @@ -62,6 +64,9 @@ export async function validateImageGenConfig(
LIVEPEER_GATEWAY_URL:
runtime.getSetting("LIVEPEER_GATEWAY_URL") ||
process.env.LIVEPEER_GATEWAY_URL,
IDEOGRAM_API_KEY:
runtime.getSetting("IDEOGRAM_API_KEY") ||
process.env.IDEOGRAM_API_KEY,
};

return imageGenEnvSchema.parse(config);
Expand Down
Loading
Loading