├── .env.example ├── tsconfig.json ├── .gitignore ├── package.json ├── src ├── types.ts ├── server.ts └── imageService.ts └── README.md /.env.example: -------------------------------------------------------------------------------- 1 | # Replicate API token (required) 2 | REPLICATE_API_TOKEN=your_replicate_api_token_here -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "ES2022", 4 | "module": "ES2022", 5 | "lib": ["ES2022"], 6 | "outDir": "./dist", 7 | "rootDir": "./src", 8 | "strict": true, 9 | "esModuleInterop": true, 10 | "skipLibCheck": true, 11 | "forceConsistentCasingInFileNames": true, 12 | "resolveJsonModule": true, 13 | "moduleResolution": "node" 14 | }, 15 | "include": ["src/**/*"], 16 | "exclude": ["node_modules", "dist"] 17 | } -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Environment variables 2 | .env 3 | 4 | # Generated images 5 | images/ 6 | 7 | # Dependencies 8 | node_modules/ 9 | 10 | # Build output 11 | dist/ 12 | 13 | # Logs 14 | logs/ 15 | *.log 16 | npm-debug.log* 17 | yarn-debug.log* 18 | yarn-error.log* 19 | 20 | # IDE and editor files 21 | .idea/ 22 | .vscode/ 23 | *.swp 24 | *.swo 25 | .DS_Store 26 | 27 | # Coverage directory 28 | coverage/ 29 | 30 | # Temporary files 31 | tmp/ 32 | temp/ 33 | 34 | # Local development 35 | *.local 36 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "image-gen-mcp", 3 | "version": "1.0.0", 4 | "description": "MCP Server for image generation using Replicate's flux-schnell model", 5 | "type": "module", 6 | "main": "dist/server.js", 7 | "scripts": { 8 | "build": "tsc", 9 | "start": "node dist/server.js", 10 | "dev": "node --loader ts-node/esm src/server.ts", 11 | "test": "echo \"Error: no test specified\" && exit 1" 12 | }, 13 | "keywords": [ 14 | "mcp", 15 | "image-generation", 16 | "replicate", 17 | "flux-schnell" 18 | ], 19 | "author": "", 20 | "license": "ISC", 21 | "dependencies": { 22 | "@modelcontextprotocol/sdk": "^1.4.1", 23 | "@types/node": "^20.0.0", 24 | "dotenv": "^16.0.0", 25 | "replicate": "^0.25.0", 26 | "typescript": "^5.0.0", 27 | "zod": "^3.24.1" 28 | }, 29 | "devDependencies": { 30 | "ts-node": "^10.9.0" 31 | } 32 | } 33 | -------------------------------------------------------------------------------- /src/types.ts: -------------------------------------------------------------------------------- 1 | export interface ImageGenerationParams { 2 | prompt: string; 3 | output_dir: string; 4 | filename?: string; 5 | go_fast?: boolean; 6 | megapixels?: "1" | "2" | "4"; 7 | num_outputs?: number; 8 | aspect_ratio?: "1:1" | "4:3" | "16:9"; 9 | output_format?: "webp" | "png" | "jpeg"; 10 | output_quality?: number; 11 | num_inference_steps?: number; 12 | } 13 | 14 | export interface ImageGenerationResponse { 15 | image_paths: string[]; 16 | metadata: { 17 | model: string; 18 | inference_time_ms: number; 19 | cache_hit?: boolean; 20 | }; 21 | } 22 | 23 | export interface ValidationError extends Error { 24 | code: 'VALIDATION_ERROR'; 25 | details: Record; 26 | } 27 | 28 | export interface APIError extends Error { 29 | code: 'API_ERROR'; 30 | details: { 31 | message: string; 32 | status?: number; 33 | }; 34 | } 35 | 36 | export interface ServerError extends Error { 37 | code: 'SERVER_ERROR'; 38 | details: { 39 | message: string; 40 | system_error?: string; 41 | }; 42 | } -------------------------------------------------------------------------------- /src/server.ts: -------------------------------------------------------------------------------- 1 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; 2 | import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; 3 | import { z } from "zod"; 4 | import { ImageGenerationService } from './imageService.js'; 5 | import dotenv from 'dotenv'; 6 | import { promises as fs } from 'fs'; 7 | 8 | // Load environment variables 9 | dotenv.config(); 10 | 11 | const imageService = new ImageGenerationService(); 12 | 13 | const server = new McpServer({ 14 | name: "image-generator", 15 | version: "1.0.0" 16 | }); 17 | 18 | // Register the image generation tool 19 | server.tool( 20 | "generate-image", 21 | "Generate an image based on a prompt", 22 | { 23 | prompt: z.string(), 24 | output_dir: z.string().describe("Full absolute path to output directory. For Windows, use double backslashes like 'C:\\\\Users\\\\name\\\\path'. For Unix/Mac use '/path/to/dir'. Always use the proper path otherwise you will get an error."), 25 | filename: z.string().optional().describe("Base filename to save the image(s) with"), 26 | go_fast: z.boolean().optional(), 27 | megapixels: z.enum(["1", "2", "4"]).optional(), 28 | num_outputs: z.number().min(1).max(4).optional(), 29 | aspect_ratio: z.enum(["1:1", "4:3", "16:9"]).optional(), 30 | output_format: z.enum(["webp", "png", "jpeg"]).optional(), 31 | output_quality: z.number().min(1).max(100).optional(), 32 | num_inference_steps: z.number().min(4).max(20).optional() 33 | }, 34 | async (params) => { 35 | try { 36 | const result = await imageService.generateImages(params); 37 | return { 38 | content: [{ 39 | type: "text", 40 | text: JSON.stringify(result, null, 2) 41 | }] 42 | }; 43 | } catch (error: any) { 44 | return { 45 | content: [{ 46 | type: "text", 47 | text: `Error: ${error.message}` 48 | }] 49 | }; 50 | } 51 | } 52 | ); 53 | 54 | 55 | async function main() { 56 | const transport = new StdioServerTransport(); 57 | await server.connect(transport); 58 | console.error("Image Generation MCP Server running on stdio"); 59 | } 60 | 61 | main().catch(async (error) => { 62 | const fs = require('fs').promises; 63 | await fs.appendFile('server.log', `${new Date().toISOString()} - ${error.stack || error}\n`); 64 | console.error(error); 65 | process.exit(1); 66 | }); -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Image Generation MCP Server 2 | 3 | An [MCP (Model Context Protocol)](https://modelcontextprotocol.io/) server implementation for generating images using Replicate's [`black-forest-labs/flux-schnell`](https://replicate.com/black-forest-labs/flux-schnell) model. 4 | 5 | Ideally to be used with Cursor's MCP feature, but can be used with any MCP client. 6 | 7 | ## Features 8 | 9 | - Generate images from text prompts 10 | - Configurable image parameters (resolution, aspect ratio, quality) 11 | - Save generated images to specified directory 12 | - Full MCP protocol compliance 13 | - Error handling and validation 14 | 15 | ## Prerequisites 16 | 17 | - Node.js 16+ 18 | - Replicate API token 19 | - TypeScript SDK for MCP 20 | 21 | ## Setup 22 | 23 | 1. Clone the repository 24 | 2. Install dependencies: 25 | ```bash 26 | npm install 27 | ``` 28 | 3. Add your Replicate API token directly in the code at `src/imageService.ts` by updating the `apiToken` constant: 29 | ```bash 30 | // No environment variables are used since they can't be easily set in cursor 31 | const apiToken = "your-replicate-api-token-here"; 32 | ``` 33 | 34 | > **Note:** If using with Claude, you can create a `.env` file in the root directory and set your API token there: 35 | ```bash 36 | REPLICATE_API_TOKEN=your-replicate-api-token-here 37 | ``` 38 | 39 | Then build the project: 40 | ```bash 41 | npm run build 42 | ``` 43 | 44 | ## Usage 45 | 46 | To use with cursor: 47 | 1. Go to Settings 48 | 2. Select Features 49 | 3. Scroll down to "MCP Servers" 50 | 4. Click "Add new MCP Server" 51 | 5. Set Type to "Command" 52 | 6. Set Command to: `node ./path/to/dist/server.js` 53 | 54 | ## API Parameters 55 | 56 | | Parameter | Type | Required | Default | Description | 57 | |--------------------|---------|----------|---------|------------------------------------------------| 58 | | `prompt` | string | Yes | - | Text prompt for image generation | 59 | | `output_dir` | string | Yes | - | Server directory path to save generated images | 60 | | `go_fast` | boolean | No | false | Enable faster generation mode | 61 | | `megapixels` | string | No | "1" | Resolution quality ("1", "2", "4") | 62 | | `num_outputs` | number | No | 1 | Number of images to generate (1-4) | 63 | | `aspect_ratio` | string | No | "1:1" | Aspect ratio ("1:1", "4:3", "16:9") | 64 | | `output_format` | string | No | "webp" | Image format ("webp", "png", "jpeg") | 65 | | `output_quality` | number | No | 80 | Compression quality (1-100) | 66 | | `num_inference_steps`| number| No | 4 | Number of denoising steps (4-20) | 67 | 68 | ## Example Request 69 | 70 | ```json 71 | { 72 | "prompt": "black forest gateau cake spelling out 'FLUX SCHNELL'", 73 | "output_dir": "/var/output/images", 74 | "filename": "black_forest_cake", 75 | "output_format": "webp" 76 | "go_fast": true, 77 | "megapixels": "1", 78 | "num_outputs": 2, 79 | "aspect_ratio": "1:1" 80 | } 81 | ``` 82 | 83 | ## Example Response 84 | 85 | ```json 86 | { 87 | "image_paths": [ 88 | "/var/output/images/output_0.webp", 89 | "/var/output/images/output_1.webp" 90 | ], 91 | "metadata": { 92 | "model": "black-forest-labs/flux-schnell", 93 | "inference_time_ms": 2847 94 | } 95 | } 96 | ``` 97 | 98 | ## Error Handling 99 | 100 | The server handles the following error types: 101 | 102 | - Validation errors (invalid parameters) 103 | - API errors (Replicate API issues) 104 | - Server errors (filesystem, permissions) 105 | - Unknown errors (unexpected issues) 106 | 107 | Each error response includes: 108 | - Error code 109 | - Human-readable message 110 | - Detailed error information 111 | 112 | ## License 113 | 114 | ISC -------------------------------------------------------------------------------- /src/imageService.ts: -------------------------------------------------------------------------------- 1 | import Replicate from 'replicate'; 2 | import * as fs from 'fs'; 3 | import * as path from 'path'; 4 | import { ImageGenerationParams, ImageGenerationResponse, APIError, ServerError } from './types.js'; 5 | import { createHash } from 'crypto'; 6 | import { promises as fsPromises } from 'fs'; 7 | 8 | 9 | const apiToken = process.env.REPLICATE_API_TOKEN || "YOUR API TOKEN HERE"; 10 | 11 | export class ImageGenerationService { 12 | private replicate: Replicate; 13 | 14 | private readonly MODEL = 'black-forest-labs/flux-schnell'; 15 | private readonly MAX_RETRIES = 3; 16 | private readonly RETRY_DELAY = 1000; // ms 17 | private readonly cache = new Map(); 21 | private readonly CACHE_TTL = 1000 * 60 * 60; // 1 hour 22 | 23 | constructor() { 24 | 25 | if (!apiToken) { 26 | throw new Error('REPLICATE_API_TOKEN environment variable is required'); 27 | } 28 | 29 | this.replicate = new Replicate({ auth: apiToken }); 30 | 31 | // Start cache cleanup interval 32 | setInterval(() => this.cleanupCache(), this.CACHE_TTL); 33 | } 34 | 35 | private generateCacheKey(params: ImageGenerationParams): string { 36 | const relevantParams = { 37 | prompt: params.prompt, 38 | go_fast: params.go_fast ?? false, 39 | megapixels: params.megapixels ?? "1", 40 | num_outputs: params.num_outputs ?? 1, 41 | aspect_ratio: params.aspect_ratio ?? "1:1", 42 | num_inference_steps: params.num_inference_steps ?? 1, 43 | output_dir: params.output_dir 44 | }; 45 | return createHash('sha256') 46 | .update(JSON.stringify(relevantParams)) 47 | .digest('hex'); 48 | } 49 | 50 | private cleanupCache(): void { 51 | const now = Date.now(); 52 | for (const [key, value] of this.cache.entries()) { 53 | if (now - value.timestamp > this.CACHE_TTL) { 54 | this.cache.delete(key); 55 | } 56 | } 57 | } 58 | 59 | async generateImages(params: ImageGenerationParams): Promise { 60 | const startTime = Date.now(); 61 | 62 | try { 63 | // Check cache first 64 | const cacheKey = this.generateCacheKey(params); 65 | const cached = this.cache.get(cacheKey); 66 | 67 | if (cached) { 68 | // Verify files still exist 69 | const allFilesExist = cached.response.image_paths.every(path => fs.existsSync(path)); 70 | if (allFilesExist) { 71 | return { 72 | ...cached.response, 73 | metadata: { 74 | ...cached.response.metadata, 75 | cache_hit: true 76 | } 77 | }; 78 | } 79 | // If files don't exist, remove from cache 80 | this.cache.delete(cacheKey); 81 | } 82 | 83 | // Prepare model input 84 | const modelInput = { 85 | prompt: params.prompt, 86 | go_fast: params.go_fast ?? false, 87 | megapixels: params.megapixels ?? "1", 88 | num_outputs: params.num_outputs ?? 1, 89 | aspect_ratio: params.aspect_ratio ?? "1:1", 90 | num_inference_steps: params.num_inference_steps ?? 4 91 | }; 92 | 93 | // Call Replicate API 94 | const output = await this.replicate.run( 95 | this.MODEL, 96 | { input: modelInput } 97 | ) as string[]; 98 | 99 | // Download and save images 100 | const imagePaths = await this.saveImages( 101 | output, 102 | params.output_dir, 103 | params.output_format ?? 'webp', 104 | params.output_quality ?? 80, 105 | params.filename 106 | ); 107 | 108 | const endTime = Date.now(); 109 | 110 | const response: ImageGenerationResponse = { 111 | image_paths: imagePaths, 112 | metadata: { 113 | model: this.MODEL, 114 | inference_time_ms: endTime - startTime, 115 | cache_hit: false 116 | } 117 | }; 118 | 119 | // Cache the result 120 | this.cache.set(cacheKey, { 121 | response, 122 | timestamp: Date.now() 123 | }); 124 | 125 | return response; 126 | 127 | } catch (error: any) { 128 | 129 | if (error.response) { 130 | const apiError = new Error(error.message) as APIError; 131 | apiError.code = 'API_ERROR'; 132 | apiError.details = { 133 | message: error.message, 134 | status: error.response.status 135 | }; 136 | throw apiError; 137 | } 138 | 139 | const serverError = new Error('Server error occurred') as ServerError; 140 | serverError.code = 'SERVER_ERROR'; 141 | serverError.details = { 142 | message: 'Failed to generate or save images', 143 | system_error: error.message 144 | }; 145 | throw serverError; 146 | } 147 | } 148 | 149 | private async downloadWithRetry( 150 | url: string, 151 | retries = this.MAX_RETRIES 152 | ): Promise { 153 | for (let i = 0; i < retries; i++) { 154 | try { 155 | const response = await fetch(url); 156 | if (!response.ok) { 157 | throw new Error(`Failed to download image: ${response.statusText}`); 158 | } 159 | return Buffer.from(await response.arrayBuffer()); 160 | } catch (error) { 161 | if (i === retries - 1) throw error; 162 | await new Promise(resolve => setTimeout(resolve, this.RETRY_DELAY * (i + 1))); 163 | } 164 | } 165 | throw new Error('Failed to download after retries'); 166 | } 167 | 168 | private async saveImages( 169 | imageUrls: string[], 170 | outputDir: string, 171 | format: 'webp' | 'png' | 'jpeg', 172 | quality: number, 173 | baseFilename?: string 174 | ): Promise { 175 | // Create output directory if it doesn't exist 176 | if (!fs.existsSync(outputDir)) { 177 | fs.mkdirSync(outputDir, { recursive: true }); 178 | } 179 | 180 | // Prepare download tasks 181 | const downloadTasks = imageUrls.map(async (imageUrl, i) => { 182 | const filename = baseFilename 183 | ? (imageUrls.length > 1 ? `${baseFilename}_${i + 1}.${format}` : `${baseFilename}.${format}`) 184 | : `output_${i}.${format}`; 185 | const filePath = path.join(outputDir, filename); 186 | 187 | try { 188 | // Download image with retry mechanism 189 | const buffer = await this.downloadWithRetry(imageUrl); 190 | 191 | // Save image atomically using temporary file 192 | const tempPath = `${filePath}.tmp`; 193 | fs.writeFileSync(tempPath, buffer); 194 | fs.renameSync(tempPath, filePath); 195 | 196 | return filePath; 197 | } catch (error: any) { 198 | const serverError = new Error('Failed to save image') as ServerError; 199 | serverError.code = 'SERVER_ERROR'; 200 | serverError.details = { 201 | message: `Failed to save image ${i}`, 202 | system_error: error.message 203 | }; 204 | throw serverError; 205 | } 206 | }); 207 | 208 | // Execute all downloads in parallel with a concurrency limit 209 | const CONCURRENCY_LIMIT = 3; 210 | const imagePaths: string[] = []; 211 | 212 | for (let i = 0; i < downloadTasks.length; i += CONCURRENCY_LIMIT) { 213 | const batch = downloadTasks.slice(i, i + CONCURRENCY_LIMIT); 214 | const results = await Promise.all(batch); 215 | imagePaths.push(...results); 216 | } 217 | 218 | return imagePaths; 219 | } 220 | } --------------------------------------------------------------------------------