UNPKG

@bratcliffe909/mcp-server-segmind

Version:

Model Context Protocol server for Segmind API - Generate images and videos using AI models

182 lines 7.83 kB
import { z } from 'zod'; import { modelRegistry, ModelCategory } from '../models/registry.js'; import { logger } from '../utils/logger.js'; import { BaseTool } from './base.js'; const GenerateVideoSchema = z.object({ prompt: z.string().min(1).max(2000).describe('Text prompt or motion description for video generation'), model: z.string().optional().describe('Model ID to use for video generation'), image: z.string().optional().describe('Input image for image-to-video generation (base64 or URL)'), duration: z.number().min(1).max(30).default(5).describe('Video duration in seconds'), fps: z.number().int().min(12).max(60).default(24).describe('Frames per second'), aspect_ratio: z.enum(['16:9', '9:16', '1:1', '4:3']).default('16:9').describe('Video aspect ratio'), quality: z.enum(['standard', 'high', 'ultra']).default('high').describe('Video quality preset'), motion_strength: z.number().min(0).max(1).optional().describe('Motion intensity for image-to-video'), seed: z.number().int().optional().describe('Seed for reproducible generation'), save_location: z.string().optional().describe('Directory path to save the video. Overrides default save location.'), }); export class GenerateVideoTool extends BaseTool { name = 'generate_video'; description = 'Generate videos from text prompts or animate static images'; async execute(params) { try { const validated = GenerateVideoSchema.parse(params); const isImageToVideo = !!validated.image; const model = this.selectModel(validated, isImageToVideo); if (!model) { return { content: [{ type: 'text', text: 'No suitable model found for video generation.', }], isError: true, }; } logger.info(`Selected model ${model.id} for video generation`, { type: isImageToVideo ? 'image-to-video' : 'text-to-video', }); if (validated.image) { const imageValidation = await this.validateImageInput(validated.image); if (!imageValidation.isValid) { return { content: [{ type: 'text', text: `Invalid image input: ${imageValidation.error}`, }], isError: true, }; } } const modelParams = await this.prepareModelParameters(validated, model); logger.info('Video generation parameters', { model: model.id, originalParams: validated, preparedParams: modelParams, }); const paramValidation = modelRegistry.validateModelParameters(model.id, modelParams); if (!paramValidation.success) { logger.error('Parameter validation failed', { model: model.id, params: modelParams, error: paramValidation.error, }); return { content: [{ type: 'text', text: `Invalid parameters for model ${model.id}: ${paramValidation.error}`, }], isError: true, }; } logger.info('Starting video generation', { model: model.id, duration: validated.duration, estimatedTime: model.estimatedTime, }); const result = await this.callModel(model, paramValidation.data, validated.save_location); return { content: result.content }; } catch (error) { logger.error('Video generation failed', { error }); return this.createErrorResponse(error); } } selectModel(params, isImageToVideo) { if (params.model) { const model = modelRegistry.getModel(params.model); if (model && model.category === ModelCategory.TEXT_TO_VIDEO) { return model; } logger.warn(`Model ${params.model} not found or not a video generation model`); } const videoModels = modelRegistry.getModelsByCategory(ModelCategory.TEXT_TO_VIDEO); if (isImageToVideo) { return videoModels[0]; } else { return videoModels.find(m => m.id === 'seedance-v1-lite') || videoModels[0]; } } async prepareModelParameters(params, model) { const baseParams = {}; if (model.id === 'veo-3') { baseParams.prompt = params.prompt; if (params.seed !== undefined) { baseParams.seed = params.seed; } if (params.aspect_ratio) { baseParams.aspect_ratio = params.aspect_ratio; } } else if (model.id === 'seedance-v1-lite') { baseParams.prompt = params.prompt; if (params.duration !== undefined) { baseParams.duration = Math.min(Math.max(params.duration, 5), 10); } const supportedAspectRatios = ['16:9', '4:3', '1:1', '3:4', '9:16']; if (params.aspect_ratio && supportedAspectRatios.includes(params.aspect_ratio)) { baseParams.aspect_ratio = params.aspect_ratio; } if (model.parameters.shape.resolution) { if (params.quality === 'standard') { baseParams.resolution = '480p'; } else { baseParams.resolution = '720p'; } } if (params.seed !== undefined && model.parameters.shape.seed) { baseParams.seed = params.seed; } } else { baseParams.prompt = params.prompt; if (params.image && model.parameters.shape.image) { baseParams.image = await this.processImageInput(params.image); } const modelShape = model.parameters.shape; if (params.duration !== undefined && modelShape.duration) { baseParams.duration = params.duration; } if (params.fps !== undefined && modelShape.fps) { baseParams.fps = params.fps; } if (params.aspect_ratio !== undefined && modelShape.aspect_ratio) { baseParams.aspect_ratio = params.aspect_ratio; } if (params.seed !== undefined && modelShape.seed) { baseParams.seed = params.seed; } } return this.mergeWithDefaults(baseParams, model); } async validateImageInput(input) { if (input.startsWith('http://') || input.startsWith('https://')) { return { isValid: true }; } const base64Regex = /^data:image\/(png|jpeg|jpg|webp);base64,/; if (input.match(base64Regex) || this.isValidBase64(input)) { return { isValid: true }; } return { isValid: false, error: 'Image must be a valid URL or base64 encoded string', }; } isValidBase64(str) { try { Buffer.from(str, 'base64'); return true; } catch { return false; } } async processImageInput(input) { if (input.startsWith('http') || input.startsWith('data:image/')) { return input; } return `data:image/png;base64,${input}`; } } export const generateVideoTool = new GenerateVideoTool(); //# sourceMappingURL=generate-video.js.map