Skip to content

Commit

Permalink
Merge pull request elizaOS#786 from ProphetX10/main
Browse files Browse the repository at this point in the history
Update generation.ts to fix TOGETHER/LLAMACLOUD image generation
  • Loading branch information
lalalune authored Dec 2, 2024
2 parents 70030f0 + 98cd687 commit 46d6cfd
Showing 1 changed file with 43 additions and 17 deletions.
60 changes: 43 additions & 17 deletions packages/core/src/generation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -874,7 +874,6 @@ export const generateImage = async (
runtime.imageModelProvider === ModelProviderName.LLAMACLOUD
) {
const together = new Together({ apiKey: apiKey as string });
// Fix: steps 4 is for schnell; 28 is for dev.
const response = await together.images.create({
model: "black-forest-labs/FLUX.1-schnell",
prompt: data.prompt,
Expand All @@ -883,23 +882,41 @@ export const generateImage = async (
steps: modelSettings?.steps ?? 4,
n: data.count,
});
const urls: string[] = [];
for (let i = 0; i < response.data.length; i++) {
const json = response.data[i].b64_json;
// decode base64
const base64 = Buffer.from(json, "base64").toString("base64");
urls.push(base64);

// Add type assertion to handle the response properly
const togetherResponse = response as unknown as TogetherAIImageResponse;

if (!togetherResponse.data || !Array.isArray(togetherResponse.data)) {
throw new Error("Invalid response format from Together AI");
}
const base64s = await Promise.all(
urls.map(async (url) => {
const response = await fetch(url);
const blob = await response.blob();
const buffer = await blob.arrayBuffer();
let base64 = Buffer.from(buffer).toString("base64");
base64 = "data:image/jpeg;base64," + base64;
return base64;
})
);

// Rest of the code remains the same...
const base64s = await Promise.all(togetherResponse.data.map(async (image) => {
if (!image.url) {
elizaLogger.error("Missing URL in image data:", image);
throw new Error("Missing URL in Together AI response");
}

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

// Convert to blob and then to base64
const blob = await imageResponse.blob();
const arrayBuffer = await blob.arrayBuffer();
const base64 = Buffer.from(arrayBuffer).toString('base64');

// Return with proper MIME type
return `data:image/jpeg;base64,${base64}`;
}));

if (base64s.length === 0) {
throw new Error("No images generated by Together AI");
}

elizaLogger.debug(`Generated ${base64s.length} images`);
return { success: true, data: base64s };
} else if (runtime.imageModelProvider === ModelProviderName.FAL) {
fal.config({
Expand Down Expand Up @@ -1406,3 +1423,12 @@ async function handleOllama({
...modelOptions,
});
}

// Add type definition for Together AI response
interface TogetherAIImageResponse {
data: Array<{
url: string;
content_type?: string;
image_type?: string;
}>;
}

0 comments on commit 46d6cfd

Please sign in to comment.