huggingface-mcp-server/dist/handler.js

MCP Server for HuggingFace inference endpoints with custom LoRA and story generation

"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.handleToolCalls = handleToolCalls;
exports.generateImage = generateImage;
exports.generateStory = generateStory;
const axios_1 = __importDefault(require("axios"));
const fs_1 = __importDefault(require("fs"));
const path_1 = __importDefault(require("path"));
async function handleToolCalls(toolCalls, apiKey) {
    const results = [];
    for (const toolCall of toolCalls) {
        const functionName = toolCall.function.name;
        const arguments_ = JSON.parse(toolCall.function.arguments);
        if (functionName === 'generate_image') {
            const result = await generateImage(arguments_, apiKey);
            results.push({ tool_call_id: toolCall.id, role: 'tool', content: result });
        }
        else if (functionName === 'generate_story') {
            const result = await generateStory(arguments_, apiKey);
            results.push({ tool_call_id: toolCall.id, role: 'tool', content: result });
        }
    }
    return results;
}
async function generateImage(input, apiKey) {
    if (!apiKey) {
        return 'Error: HuggingFace API key is not set';
    }
    // Use the FLUX.1 model
    const model = 'black-forest-labs/FLUX.1-dev';
    try {
        // Set default values for parameters if not provided
        const numInferenceSteps = input.num_inference_steps || 25;
        const height = input.height || 1024;
        const width = input.width || 1024;
        const guidanceScale = input.guidance_scale || 3.5;
        const seed = input.seed || Math.floor(Math.random() * 2147483647); // MAX_SEED in Python
        const numImagesPerPrompt = input.num_images_per_prompt || 1;
        // Since we can't use diffusers.js directly, we'll use the Hugging Face API
        // Create the payload with the parameters
        const payload = {
            inputs: input.prompt,
            parameters: {
                num_inference_steps: numInferenceSteps,
                height: height,
                width: width,
                guidance_scale: guidanceScale,
                seed: seed.toString(), // Convert to string as required by the API
                num_images_per_prompt: numImagesPerPrompt
            }
        };
        // If LoRA is specified, add it to the payload
        if (input.lora_name) {
            try {
                console.log("Adding LoRA to payload:", input.lora_name);
                // FLUX.1 uses a specific nested structure for LoRA parameters
                payload.parameters.pipeline_params = {
                    lora_weights: input.lora_name,
                    lora_scale: input.lora_scale || 0.7 // Default scale if not provided
                };
            }
            catch (error) {
                console.error("Error setting LoRA:", error);
                return `Error setting LoRA: ${error.message}`;
            }
        }
        // Call the Hugging Face API
        const response = await axios_1.default.post(`https://api-inference.huggingface.co/models/${model}`, payload, {
            headers: {
                Authorization: `Bearer ${apiKey}`,
                'Content-Type': 'application/json'
            },
            responseType: 'arraybuffer'
        });
        // For image generation, the response is binary data
        if (response.status === 200) {
            // In a real scenario you would return a URL or base64 of the image
            const imageBuffer = Buffer.from(response.data);
            const timestamp = Date.now();
            const filename = `image-${timestamp}.jpeg`;
            const imageDir = path_1.default.join(__dirname, 'images');
            // Create the image directory if it doesn't exist
            if (!fs_1.default.existsSync(imageDir)) {
                fs_1.default.mkdirSync(imageDir);
            }
            const filePath = path_1.default.join(imageDir, filename);
            fs_1.default.writeFileSync(filePath, imageBuffer);
            return filePath;
        }
        else {
            return `Error generating image: ${response.statusText}`;
        }
    }
    catch (error) {
        console.error("Error generating image:", error);
        return `Error: ${error.message}`;
    }
}
async function generateStory(input, apiKey) {
    if (!apiKey) {
        return 'Error: HuggingFace API key is not set';
    }
    const model = 'mistralai/Mistral-7B-Instruct-v0.2'; // Example model, change as needed
    try {
        const response = await axios_1.default.post(`https://api-inference.huggingface.co/models/${model}`, {
            inputs: `Generate a short story based on this prompt: ${input.prompt}`,
            parameters: {
                max_length: 1000,
                temperature: 0.7,
                top_p: 0.9
            }
        }, {
            headers: {
                Authorization: `Bearer ${apiKey}`,
                'Content-Type': 'application/json'
            }
        });
        if (response.status === 200) {
            return response.data[0].generated_text;
        }
        else {
            return `Error generating story: ${response.statusText}`;
        }
    }
    catch (error) {
        return `Error: ${error.message}`;
    }
}