generateImage.ts•1.62 kB
import { ImageGenerationParams } from "../types/index.js";
import { replicate } from "../services/replicate.js";
import { handleError } from "../utils/error.js";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { FileOutput } from "replicate";
import { outputToBase64 } from "../utils/image.js";
import { CONFIG } from "../config/index.js";
export const registerGenerateImageTool = async (
input: ImageGenerationParams
): Promise<CallToolResult> => {
const { support_image_mcp_response_type, ...predictionInput } = input;
try {
const [output] = (await replicate.run(CONFIG.imageModelId, {
input: predictionInput,
})) as [FileOutput];
const imageUrl = output.url() as unknown as string;
if (support_image_mcp_response_type) {
const imageBase64 = await outputToBase64(output);
return {
content: [
{
type: "text",
text: `This is a generated image link: ${imageUrl}`,
},
{
type: "image",
data: imageBase64,
mimeType: "image/png",
},
{
type: "text",
text: `The image above is generated by the Flux model and prompt: ${input.prompt}`,
},
],
};
}
return {
content: [
{
type: "text",
text: `This is a generated image link: ${imageUrl}`,
},
{
type: "text",
text: `The image above is generated by the Flux model and prompt: ${input.prompt}`,
},
],
};
} catch (error) {
handleError(error);
}
};