#!/usr/bin/env node
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { z } from "zod";
import Replicate from "replicate";
import { writeFile } from "fs/promises";
import * as fs from 'fs';
import * as path from 'path';
import * as https from 'https';
import * as http from 'http';
// Check for required environment variable
const REPLICATE_API_TOKEN = process.env.REPLICATE_API_TOKEN;
let replicateClient: Replicate | null = null;
if (!REPLICATE_API_TOKEN) {
console.error('REPLICATE_API_TOKEN environment variable is required');
console.error('Please set your Replicate API token: export REPLICATE_API_TOKEN=r8_your_token_here');
// Server continues running, no process.exit()
} else {
// Configure Replicate client
replicateClient = new Replicate({
auth: REPLICATE_API_TOKEN
});
}
// Define types based on FLUX.1 Kontext [Max] API documentation
interface FluxImageResult {
images?: Array<{
url: string;
content_type?: string;
width?: number;
height?: number;
}>;
output?: string | string[]; // Replicate returns output as URL or array of URLs
seed?: number;
has_nsfw_concepts?: boolean[];
prompt?: string;
}
// Download image function
async function downloadImage(url: string, filename: string): Promise<string> {
return new Promise((resolve, reject) => {
try {
const parsedUrl = new URL(url);
const client = parsedUrl.protocol === 'https:' ? https : http;
// Create images directory if it doesn't exist
const imagesDir = path.join(process.cwd(), 'images');
if (!fs.existsSync(imagesDir)) {
fs.mkdirSync(imagesDir, { recursive: true });
}
const filePath = path.join(imagesDir, filename);
const file = fs.createWriteStream(filePath);
client.get(url, (response) => {
if (response.statusCode !== 200) {
reject(new Error(`Failed to download image: HTTP ${response.statusCode}`));
return;
}
response.pipe(file);
file.on('finish', () => {
file.close();
resolve(filePath);
});
file.on('error', (err) => {
fs.unlink(filePath, () => {}); // Delete partial file
reject(err);
});
}).on('error', (err) => {
reject(err);
});
} catch (error) {
reject(error);
}
});
}
// Generate safe filename for images
function generateImageFilename(prompt: string, index: number, seed?: number): string {
const safePrompt = prompt
.toLowerCase()
.replace(/[^a-z0-9\s]/g, '')
.replace(/\s+/g, '_')
.substring(0, 50);
const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
const seedStr = seed ? `_${seed}` : '';
return `flux_kontext_max_${safePrompt}${seedStr}_${index}_${timestamp}.jpg`;
}
// Create MCP server
const server = new McpServer({
name: "replicate-flux-kontext-max-server",
version: "1.2.0",
});
// Tool: Generate images with FLUX.1 Kontext [Max]
server.tool(
"flux_kontext_max_generate",
{
description: "Generate high-quality images using FLUX.1 Kontext [Max] via Replicate - Frontier image generation model with advanced text rendering and contextual understanding",
inputSchema: {
type: "object",
properties: {
prompt: {
type: "string",
description: "Text description of what you want to generate, or the instruction on how to edit the given image"
},
input_image: {
type: "string",
description: "Image to use as reference. Must be jpeg, png, gif, or webp. Can be a URL or base64 data URI"
},
seed: {
type: "integer",
description: "Random seed. Set for reproducible generation"
},
aspect_ratio: {
type: "string",
enum: ["match_input_image", "1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "4:5", "5:4", "21:9", "9:21", "2:1", "1:2"],
description: "Aspect ratio of the generated image. Use 'match_input_image' to match the aspect ratio of the input image",
default: "match_input_image"
},
output_format: {
type: "string",
enum: ["jpg", "png"],
description: "Output format for the generated image",
default: "jpg"
},
safety_tolerance: {
type: "integer",
minimum: 0,
maximum: 6,
description: "Safety tolerance, 0 is most strict and 6 is most permissive. 2 is currently the maximum allowed when input images are used",
default: 2
}
},
required: ["prompt"]
}
},
async (args: any) => {
// Check if Replicate client is configured
if (!replicateClient) {
return {
content: [{
type: "text",
text: "Error: REPLICATE_API_TOKEN environment variable is not set. Please configure your Replicate API token."
}],
isError: true
};
}
const {
prompt,
input_image,
seed,
aspect_ratio = "match_input_image",
output_format = "jpg",
safety_tolerance = 2
} = args;
try {
// Prepare input for Replicate API
const input: any = {
prompt,
aspect_ratio,
output_format,
safety_tolerance
};
// Add optional parameters if provided
if (input_image) {
input.input_image = input_image;
}
if (seed !== undefined) {
input.seed = seed;
}
console.error(`Generating image with FLUX.1 Kontext [Max] via Replicate - prompt: "${prompt}"`);
// Call Replicate FLUX.1 Kontext [Max] API
const output = await replicateClient.run("black-forest-labs/flux-kontext-max", { input });
// Handle different output formats from Replicate
let imageUrls: string[] = [];
if (typeof output === 'string') {
imageUrls = [output];
} else if (Array.isArray(output)) {
imageUrls = output;
} else {
throw new Error('Unexpected output format from Replicate API');
}
// Download images locally
console.error("Downloading images locally...");
const downloadedImages = [];
for (let i = 0; i < imageUrls.length; i++) {
const url = imageUrls[i];
const filename = generateImageFilename(prompt, i + 1, seed);
try {
const localPath = await downloadImage(url, filename);
downloadedImages.push({
url,
localPath,
index: i + 1,
content_type: `image/${output_format}`,
filename
});
console.error(`Downloaded: ${filename}`);
} catch (downloadError) {
console.error(`Failed to download image ${i + 1}:`, downloadError);
// Still add the image info without local path
downloadedImages.push({
url,
localPath: null,
index: i + 1,
content_type: `image/${output_format}`,
filename
});
}
}
// Format response with download information
const imageDetails = downloadedImages.map(img => {
let details = `Image ${img.index}:`;
if (img.localPath) {
details += `\n Local Path: ${img.localPath}`;
}
details += `\n Original URL: ${img.url}`;
details += `\n Filename: ${img.filename}`;
return details;
}).join('\n\n');
const responseText = `Successfully generated ${downloadedImages.length} image(s) using FLUX.1 Kontext [Max] via Replicate:
Prompt: "${prompt}"
Aspect Ratio: ${aspect_ratio}
Output Format: ${output_format}
Safety Tolerance: ${safety_tolerance}
${seed ? `Seed: ${seed}` : 'Seed: Auto-generated'}
${input_image ? `Input Image: ${input_image}` : ''}
Generated Images:
${imageDetails}
${downloadedImages.some(img => img.localPath) ? 'Images have been downloaded to the local \'images\' directory.' : 'Note: Local download failed, but original URLs are available.'}`;
return {
content: [
{
type: "text",
text: responseText
}
]
};
} catch (error) {
console.error('Error generating image:', error);
let errorMessage = "Failed to generate image with FLUX.1 Kontext [Max] via Replicate.";
if (error instanceof Error) {
errorMessage += ` Error: ${error.message}`;
}
return {
content: [
{
type: "text",
text: errorMessage
}
],
isError: true
};
}
}
);
// Tool: Generate images using async method with prediction tracking
server.tool(
"flux_kontext_max_generate_async",
{
description: "Generate images using FLUX.1 Kontext [Max] via Replicate with async prediction tracking for monitoring progress",
inputSchema: {
type: "object",
properties: {
prompt: {
type: "string",
description: "Text description of what you want to generate, or the instruction on how to edit the given image"
},
input_image: {
type: "string",
description: "Image to use as reference. Must be jpeg, png, gif, or webp. Can be a URL or base64 data URI"
},
seed: {
type: "integer",
description: "Random seed. Set for reproducible generation"
},
aspect_ratio: {
type: "string",
enum: ["match_input_image", "1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "4:5", "5:4", "21:9", "9:21", "2:1", "1:2"],
description: "Aspect ratio of the generated image. Use 'match_input_image' to match the aspect ratio of the input image",
default: "match_input_image"
},
output_format: {
type: "string",
enum: ["jpg", "png"],
description: "Output format for the generated image",
default: "jpg"
},
safety_tolerance: {
type: "integer",
minimum: 0,
maximum: 6,
description: "Safety tolerance, 0 is most strict and 6 is most permissive. 2 is currently the maximum allowed when input images are used",
default: 2
}
},
required: ["prompt"]
}
},
async (args: any) => {
// Check if Replicate client is configured
if (!replicateClient) {
return {
content: [{
type: "text",
text: "Error: REPLICATE_API_TOKEN environment variable is not set. Please configure your Replicate API token."
}],
isError: true
};
}
const {
prompt,
input_image,
seed,
aspect_ratio = "match_input_image",
output_format = "jpg",
safety_tolerance = 2
} = args;
try {
// Prepare input for Replicate API
const input: any = {
prompt,
aspect_ratio,
output_format,
safety_tolerance
};
// Add optional parameters if provided
if (input_image) {
input.input_image = input_image;
}
if (seed !== undefined) {
input.seed = seed;
}
console.error(`Creating async prediction for FLUX.1 Kontext [Max] via Replicate - prompt: "${prompt}"`);
// Create prediction
const prediction = await replicateClient.predictions.create({
model: "black-forest-labs/flux-kontext-max",
input
});
console.error(`Prediction created with ID: ${prediction.id}`);
// Poll for completion
let result = prediction;
let attempts = 0;
const maxAttempts = 60; // 5 minutes with 5-second intervals
while (attempts < maxAttempts && !["succeeded", "failed", "canceled"].includes(result.status)) {
await new Promise(resolve => setTimeout(resolve, 5000));
result = await replicateClient.predictions.get(prediction.id);
attempts++;
console.error(`Status check ${attempts}: ${result.status}`);
if (result.logs && result.logs.length > 0) {
const logsStr = Array.isArray(result.logs) ? result.logs.join('\n') : result.logs;
console.error(`Logs: ${logsStr}`);
}
}
if (result.status === "failed") {
throw new Error(`Prediction failed: ${result.error || 'Unknown error'}`);
}
if (result.status === "canceled") {
throw new Error("Prediction was canceled");
}
if (result.status !== "succeeded") {
throw new Error("Prediction timed out after 5 minutes");
}
// Handle different output formats from Replicate
let imageUrls: string[] = [];
if (typeof result.output === 'string') {
imageUrls = [result.output];
} else if (Array.isArray(result.output)) {
imageUrls = result.output;
} else {
throw new Error('Unexpected output format from Replicate API');
}
// Download images locally
console.error("Downloading images locally...");
const downloadedImages = [];
for (let i = 0; i < imageUrls.length; i++) {
const url = imageUrls[i];
const filename = generateImageFilename(prompt, i + 1, seed);
try {
const localPath = await downloadImage(url, filename);
downloadedImages.push({
url,
localPath,
index: i + 1,
content_type: `image/${output_format}`,
filename
});
console.error(`Downloaded: ${filename}`);
} catch (downloadError) {
console.error(`Failed to download image ${i + 1}:`, downloadError);
// Still add the image info without local path
downloadedImages.push({
url,
localPath: null,
index: i + 1,
content_type: `image/${output_format}`,
filename
});
}
}
// Format response with download information
const imageDetails = downloadedImages.map(img => {
let details = `Image ${img.index}:`;
if (img.localPath) {
details += `\n Local Path: ${img.localPath}`;
}
details += `\n Original URL: ${img.url}`;
details += `\n Filename: ${img.filename}`;
return details;
}).join('\n\n');
const responseText = `Successfully generated ${downloadedImages.length} image(s) using FLUX.1 Kontext [Max] via Replicate (Async):
Prompt: "${prompt}"
Aspect Ratio: ${aspect_ratio}
Output Format: ${output_format}
Safety Tolerance: ${safety_tolerance}
${seed ? `Seed: ${seed}` : 'Seed: Auto-generated'}
${input_image ? `Input Image: ${input_image}` : ''}
Prediction ID: ${prediction.id}
Generated Images:
${imageDetails}
${downloadedImages.some(img => img.localPath) ? 'Images have been downloaded to the local \'images\' directory.' : 'Note: Local download failed, but original URLs are available.'}`;
return {
content: [
{
type: "text",
text: responseText
}
]
};
} catch (error) {
console.error('Error generating image:', error);
let errorMessage = "Failed to generate image with FLUX.1 Kontext [Max] via Replicate (Async).";
if (error instanceof Error) {
errorMessage += ` Error: ${error.message}`;
}
return {
content: [
{
type: "text",
text: errorMessage
}
],
isError: true
};
}
}
);
// Tool: Cancel a prediction
server.tool(
"flux_kontext_max_cancel_prediction",
{
description: "Cancel a running prediction by its ID",
inputSchema: {
type: "object",
properties: {
prediction_id: {
type: "string",
description: "The ID of the prediction to cancel"
}
},
required: ["prediction_id"]
}
},
async (args: any) => {
// Check if Replicate client is configured
if (!replicateClient) {
return {
content: [{
type: "text",
text: "Error: REPLICATE_API_TOKEN environment variable is not set. Please configure your Replicate API token."
}],
isError: true
};
}
const { prediction_id } = args;
try {
console.error(`Canceling prediction: ${prediction_id}`);
const result = await replicateClient.predictions.cancel(prediction_id);
return {
content: [
{
type: "text",
text: `Successfully canceled prediction ${prediction_id}. Status: ${result.status}`
}
]
};
} catch (error) {
console.error('Error canceling prediction:', error);
let errorMessage = "Failed to cancel prediction.";
if (error instanceof Error) {
errorMessage += ` Error: ${error.message}`;
}
return {
content: [
{
type: "text",
text: errorMessage
}
],
isError: true
};
}
}
);
// Tool: Get prediction status
server.tool(
"flux_kontext_max_get_prediction",
{
description: "Get the status and details of a prediction by its ID",
inputSchema: {
type: "object",
properties: {
prediction_id: {
type: "string",
description: "The ID of the prediction to check"
}
},
required: ["prediction_id"]
}
},
async (args: any) => {
// Check if Replicate client is configured
if (!replicateClient) {
return {
content: [{
type: "text",
text: "Error: REPLICATE_API_TOKEN environment variable is not set. Please configure your Replicate API token."
}],
isError: true
};
}
const { prediction_id } = args;
try {
console.error(`Getting prediction status: ${prediction_id}`);
const prediction = await replicateClient.predictions.get(prediction_id);
let responseText = `Prediction ${prediction_id}:
Status: ${prediction.status}
Created: ${prediction.created_at}
Model: ${prediction.model}`;
if (prediction.started_at) {
responseText += `\nStarted: ${prediction.started_at}`;
}
if (prediction.completed_at) {
responseText += `\nCompleted: ${prediction.completed_at}`;
}
if (prediction.input) {
responseText += `\nInput: ${JSON.stringify(prediction.input, null, 2)}`;
}
if (prediction.output) {
responseText += `\nOutput: ${JSON.stringify(prediction.output, null, 2)}`;
}
if (prediction.error) {
responseText += `\nError: ${prediction.error}`;
}
if (prediction.logs && prediction.logs.length > 0) {
const logsStr = Array.isArray(prediction.logs) ? prediction.logs.join('\n') : prediction.logs;
responseText += `\nLogs:\n${logsStr}`;
}
return {
content: [
{
type: "text",
text: responseText
}
]
};
} catch (error) {
console.error('Error getting prediction:', error);
let errorMessage = "Failed to get prediction status.";
if (error instanceof Error) {
errorMessage += ` Error: ${error.message}`;
}
return {
content: [
{
type: "text",
text: errorMessage
}
],
isError: true
};
}
}
);
// Graceful shutdown handlers
process.on('SIGINT', () => {
console.error('Received SIGINT, shutting down gracefully...');
process.exit(0);
});
process.on('SIGTERM', () => {
console.error('Received SIGTERM, shutting down gracefully...');
process.exit(0);
});
async function main() {
const transport = new StdioServerTransport();
await server.connect(transport);
console.error('Replicate FLUX.1 Kontext [Max] MCP server running on stdio');
}
main().catch((error) => {
console.error('Fatal error in main():', error);
process.exit(1);
});
