Flux Image MCP Server

by ckz
Verified
  • src
#!/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'); } class FluxImageServer { private server: Server; private axiosInstance; constructor() { this.server = new Server( { name: 'flux-img-mcp', version: '1.0.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 on Replicate', inputSchema: { type: 'object', properties: { prompt: { type: 'string', description: 'Text prompt describing the desired image', } }, required: ['prompt'], }, }, ], })); this.server.setRequestHandler(CallToolRequestSchema, async (request) => { if (request.params.name !== 'generate_image') { throw new McpError( ErrorCode.MethodNotFound, `Unknown tool: ${request.params.name}` ); } const args = request.params.arguments; if (typeof args !== 'object' || !args || typeof args.prompt !== 'string') { throw new McpError( ErrorCode.InvalidParams, 'Invalid parameters. Expected prompt string.' ); } try { const response = await this.axiosInstance.post( '/models/black-forest-labs/flux-schnell/predictions', { input: { prompt: args.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 Image MCP server running on stdio'); } } const server = new FluxImageServer(); server.run().catch(console.error);