#!/usr/bin/env node
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
CallToolRequestSchema,
ErrorCode,
ListToolsRequestSchema,
McpError,
} from '@modelcontextprotocol/sdk/types.js';
import axios from 'axios';
const REPLICATE_API_TOKEN = process.env.REPLICATE_API_TOKEN;
if (!REPLICATE_API_TOKEN) {
throw new Error('REPLICATE_API_TOKEN environment variable is required');
}
interface GenerateImageArgs {
prompt: string;
}
const isValidGenerateImageArgs = (args: any): args is GenerateImageArgs =>
typeof args === 'object' &&
args !== null &&
typeof args.prompt === 'string';
class FluxSchnellServer {
private server: Server;
private axiosInstance;
constructor() {
this.server = new Server(
{
name: 'flux-schnell-server',
version: '0.1.0',
},
{
capabilities: {
tools: {},
},
}
);
this.axiosInstance = axios.create({
baseURL: 'https://api.replicate.com/v1',
headers: {
'Authorization': `Bearer ${REPLICATE_API_TOKEN}`,
'Content-Type': 'application/json',
'Prefer': 'wait'
}
});
this.setupToolHandlers();
// Error handling
this.server.onerror = (error) => console.error('[MCP Error]', error);
process.on('SIGINT', async () => {
await this.server.close();
process.exit(0);
});
}
private setupToolHandlers() {
this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
tools: [
{
name: 'generate_image',
description: 'Generate an image using the Flux Schnell model',
inputSchema: {
type: 'object',
properties: {
prompt: {
type: 'string',
description: 'Text prompt describing the image to generate',
},
},
required: ['prompt'],
},
},
],
}));
this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
if (request.params.name !== 'generate_image') {
throw new McpError(
ErrorCode.MethodNotFound,
`Unknown tool: ${request.params.name}`
);
}
if (!isValidGenerateImageArgs(request.params.arguments)) {
throw new McpError(
ErrorCode.InvalidParams,
'Invalid generate_image arguments'
);
}
try {
const response = await this.axiosInstance.post(
'/models/black-forest-labs/flux-schnell/predictions',
{
input: {
prompt: request.params.arguments.prompt
}
}
);
return {
content: [
{
type: 'text',
text: JSON.stringify({
status: response.data.status,
output: response.data.output,
error: response.data.error
}, null, 2)
}
]
};
} catch (error) {
if (axios.isAxiosError(error)) {
return {
content: [
{
type: 'text',
text: `Replicate API error: ${
error.response?.data?.detail || error.message
}`
}
],
isError: true
};
}
throw error;
}
});
}
async run() {
const transport = new StdioServerTransport();
await this.server.connect(transport);
console.error('Flux Schnell MCP server running on stdio');
}
}
const server = new FluxSchnellServer();
server.run().catch(console.error);
