ai-image-generate/dist/src/imageService.js

MCP Server for image generation using Replicate's flux-schnell model

import Replicate from 'replicate';
import * as fs from 'fs';
import * as path from 'path';
import { createHash } from 'crypto';
export class ImageGenerationService {
    replicate;
    model;
    MAX_RETRIES = 3;
    RETRY_DELAY = 1000; // ms
    cache = new Map();
    CACHE_TTL = 1000 * 60 * 60; // 1 hour
    constructor(apiToken, model) {
        if (!apiToken) {
            throw new Error('REPLICATE_API_TOKEN environment variable is required');
        }
        this.replicate = new Replicate({ auth: apiToken });
        this.model = model;
        // Start cache cleanup interval
        setInterval(() => this.cleanupCache(), this.CACHE_TTL);
    }
    generateCacheKey(params) {
        const relevantParams = {
            prompt: params.prompt,
            go_fast: params.go_fast ?? false,
            megapixels: params.megapixels ?? "1",
            num_outputs: params.num_outputs ?? 1,
            aspect_ratio: params.aspect_ratio ?? "1:1",
            num_inference_steps: params.num_inference_steps ?? 1,
            output_dir: params.output_dir
        };
        return createHash('sha256')
            .update(JSON.stringify(relevantParams))
            .digest('hex');
    }
    cleanupCache() {
        const now = Date.now();
        for (const [key, value] of this.cache.entries()) {
            if (now - value.timestamp > this.CACHE_TTL) {
                this.cache.delete(key);
            }
        }
    }
    async generateImages(params) {
        const startTime = Date.now();
        try {
            // Check cache first
            const cacheKey = this.generateCacheKey(params);
            const cached = this.cache.get(cacheKey);
            if (cached) {
                // Verify files still exist
                const allFilesExist = cached.response.image_paths.every(path => fs.existsSync(path));
                if (allFilesExist) {
                    return {
                        ...cached.response,
                        metadata: {
                            ...cached.response.metadata,
                            cache_hit: true
                        }
                    };
                }
                // If files don't exist, remove from cache
                this.cache.delete(cacheKey);
            }
            // Prepare model input
            const modelInput = {
                prompt: params.prompt,
                go_fast: params.go_fast ?? false,
                megapixels: params.megapixels ?? "1",
                num_outputs: params.num_outputs ?? 1,
                aspect_ratio: params.aspect_ratio ?? "1:1",
                num_inference_steps: params.num_inference_steps ?? 4
            };
            // Call Replicate API
            const output = await this.replicate.run(this.model, { input: modelInput });
            // Download and save images
            const imagePaths = await this.saveImages(output, params.output_dir, params.output_format ?? 'webp', params.output_quality ?? 80, params.filename);
            const endTime = Date.now();
            const response = {
                image_paths: imagePaths,
                metadata: {
                    model: this.model,
                    inference_time_ms: endTime - startTime,
                    cache_hit: false
                }
            };
            // Cache the result
            this.cache.set(cacheKey, {
                response,
                timestamp: Date.now()
            });
            return response;
        }
        catch (error) {
            if (error.response) {
                const apiError = new Error(error.message);
                apiError.code = 'API_ERROR';
                apiError.details = {
                    message: error.message,
                    status: error.response.status
                };
                throw apiError;
            }
            const serverError = new Error('Server error occurred');
            serverError.code = 'SERVER_ERROR';
            serverError.details = {
                message: 'Failed to generate or save images',
                system_error: error.message
            };
            throw serverError;
        }
    }
    async downloadWithRetry(url, retries = this.MAX_RETRIES) {
        for (let i = 0; i < retries; i++) {
            try {
                const response = await fetch(url);
                if (!response.ok) {
                    throw new Error(`Failed to download image: ${response.statusText}`);
                }
                return Buffer.from(await response.arrayBuffer());
            }
            catch (error) {
                if (i === retries - 1)
                    throw error;
                await new Promise(resolve => setTimeout(resolve, this.RETRY_DELAY * (i + 1)));
            }
        }
        throw new Error('Failed to download after retries');
    }
    async saveImages(imageUrls, outputDir, format, quality, baseFilename) {
        // Create output directory if it doesn't exist
        if (!fs.existsSync(outputDir)) {
            fs.mkdirSync(outputDir, { recursive: true });
        }
        // Prepare download tasks
        const downloadTasks = imageUrls.map(async (imageUrl, i) => {
            const filename = baseFilename
                ? (imageUrls.length > 1 ? `${baseFilename}_${i + 1}.${format}` : `${baseFilename}.${format}`)
                : `output_${i}.${format}`;
            const filePath = path.join(outputDir, filename);
            try {
                // Download image with retry mechanism
                const buffer = await this.downloadWithRetry(imageUrl);
                // Save image atomically using temporary file
                const tempPath = `${filePath}.tmp`;
                fs.writeFileSync(tempPath, buffer);
                fs.renameSync(tempPath, filePath);
                return filePath;
            }
            catch (error) {
                const serverError = new Error('Failed to save image');
                serverError.code = 'SERVER_ERROR';
                serverError.details = {
                    message: `Failed to save image ${i}`,
                    system_error: error.message
                };
                throw serverError;
            }
        });
        // Execute all downloads in parallel with a concurrency limit
        const CONCURRENCY_LIMIT = 3;
        const imagePaths = [];
        for (let i = 0; i < downloadTasks.length; i += CONCURRENCY_LIMIT) {
            const batch = downloadTasks.slice(i, i + CONCURRENCY_LIMIT);
            const results = await Promise.all(batch);
            imagePaths.push(...results);
        }
        return imagePaths;
    }
}