@bratcliffe909/mcp-server-segmind
Version:
Model Context Protocol server for Segmind API - Generate images and videos using AI models
440 lines • 18 kB
JavaScript
import { z } from 'zod';
import { logger } from '../utils/logger.js';
export var ModelCategory;
(function (ModelCategory) {
ModelCategory["TEXT_TO_IMAGE"] = "text2img";
ModelCategory["IMAGE_TO_IMAGE"] = "img2img";
ModelCategory["VIDEO_GENERATION"] = "video";
ModelCategory["IMAGE_ENHANCEMENT"] = "enhancement";
ModelCategory["SPECIALIZED"] = "specialized";
})(ModelCategory || (ModelCategory = {}));
export var OutputType;
(function (OutputType) {
OutputType["IMAGE"] = "image";
OutputType["VIDEO"] = "video";
OutputType["AUDIO"] = "audio";
OutputType["TEXT"] = "text";
})(OutputType || (OutputType = {}));
export class ModelRegistry {
models = new Map();
modelsByCategory = new Map();
constructor() {
this.initializeModels();
}
initializeModels() {
Object.values(ModelCategory).forEach(category => {
this.modelsByCategory.set(category, []);
});
this.registerTextToImageModels();
this.registerImageToImageModels();
this.registerVideoModels();
this.registerEnhancementModels();
this.registerSpecializedModels();
logger.info('Model registry initialized', {
totalModels: this.models.size,
categories: Object.values(ModelCategory).map(cat => ({
category: cat,
count: this.modelsByCategory.get(cat)?.length || 0,
})),
});
}
registerModel(config) {
this.models.set(config.id, config);
const categoryModels = this.modelsByCategory.get(config.category) || [];
categoryModels.push(config.id);
this.modelsByCategory.set(config.category, categoryModels);
}
registerTextToImageModels() {
this.registerModel({
id: 'flux-1-pro',
name: 'FLUX.1 Pro',
description: 'State-of-the-art image generation with superior quality and prompt adherence',
category: ModelCategory.TEXT_TO_IMAGE,
endpoint: '/flux-1-pro',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 10,
creditsPerUse: 2,
parameters: z.object({
prompt: z.string().min(1).max(2000),
negative_prompt: z.string().optional(),
num_inference_steps: z.number().int().min(1).max(100).default(50),
guidance_scale: z.number().min(0).max(20).default(7.5),
width: z.number().int().multipleOf(8).min(512).max(2048).default(1024),
height: z.number().int().multipleOf(8).min(512).max(2048).default(1024),
seed: z.number().int().optional(),
scheduler: z.enum(['ddim', 'euler', 'euler_a', 'pndm']).default('euler_a'),
}),
defaultParams: {
num_inference_steps: 50,
guidance_scale: 7.5,
width: 1024,
height: 1024,
scheduler: 'euler_a',
},
supportedFormats: ['png', 'jpeg', 'webp'],
maxDimensions: { width: 2048, height: 2048 },
});
this.registerModel({
id: 'sdxl',
name: 'Stable Diffusion XL',
description: 'Industry standard model with high customizability and quality',
category: ModelCategory.TEXT_TO_IMAGE,
endpoint: '/sdxl',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 8,
creditsPerUse: 1,
parameters: z.object({
prompt: z.string().min(1).max(2000),
negative_prompt: z.string().optional(),
num_inference_steps: z.number().int().min(1).max(150).default(30),
guidance_scale: z.number().min(0).max(30).default(7.5),
width: z.number().int().multipleOf(64).min(512).max(2048).default(1024),
height: z.number().int().multipleOf(64).min(512).max(2048).default(1024),
seed: z.number().int().optional(),
sampler: z.enum(['ddim', 'k_euler', 'k_euler_a', 'k_dpm_2', 'k_dpm_2_a']).default('k_euler_a'),
base_64: z.boolean().default(true),
}),
defaultParams: {
num_inference_steps: 30,
guidance_scale: 7.5,
width: 1024,
height: 1024,
sampler: 'k_euler_a',
base_64: true,
},
supportedFormats: ['png', 'jpeg'],
maxDimensions: { width: 2048, height: 2048 },
});
this.registerModel({
id: 'ideogram-3',
name: 'Ideogram 3.0',
description: 'Excellent for text rendering and photorealistic images',
category: ModelCategory.TEXT_TO_IMAGE,
endpoint: '/ideogram-3',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 12,
creditsPerUse: 2,
parameters: z.object({
prompt: z.string().min(1).max(1000),
negative_prompt: z.string().optional(),
aspect_ratio: z.enum(['1:1', '16:9', '9:16', '4:3', '3:4']).default('1:1'),
style_type: z.enum(['auto', 'realistic', 'design', 'render_3d', 'anime']).default('auto'),
magic_prompt: z.boolean().default(true),
seed: z.number().int().optional(),
}),
defaultParams: {
aspect_ratio: '1:1',
style_type: 'auto',
magic_prompt: true,
},
supportedFormats: ['png', 'jpeg'],
});
this.registerModel({
id: 'gpt-image-1',
name: 'GPT Image 1',
description: 'High-quality general purpose image generation',
category: ModelCategory.TEXT_TO_IMAGE,
endpoint: '/gpt-image-1',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 15,
creditsPerUse: 3,
parameters: z.object({
prompt: z.string().min(1).max(4000),
size: z.enum(['256x256', '512x512', '1024x1024', '1792x1024', '1024x1792']).default('1024x1024'),
quality: z.enum(['standard', 'hd']).default('standard'),
style: z.enum(['vivid', 'natural']).default('vivid'),
n: z.number().int().min(1).max(4).default(1),
}),
defaultParams: {
size: '1024x1024',
quality: 'standard',
style: 'vivid',
n: 1,
},
supportedFormats: ['png'],
});
this.registerModel({
id: 'seedream-v3',
name: 'Seedream V3',
description: 'Bilingual high-resolution image generation with fast processing',
category: ModelCategory.TEXT_TO_IMAGE,
endpoint: '/seedream-v3',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 6,
creditsPerUse: 1,
parameters: z.object({
prompt: z.string().min(1).max(2000),
negative_prompt: z.string().optional(),
width: z.number().int().multipleOf(8).min(256).max(2048).default(1024),
height: z.number().int().multipleOf(8).min(256).max(2048).default(1024),
num_inference_steps: z.number().int().min(1).max(50).default(25),
guidance_scale: z.number().min(1).max(20).default(7),
seed: z.number().int().optional(),
language: z.enum(['en', 'zh', 'auto']).default('auto'),
}),
defaultParams: {
width: 1024,
height: 1024,
num_inference_steps: 25,
guidance_scale: 7,
language: 'auto',
},
supportedFormats: ['png', 'jpeg'],
maxDimensions: { width: 2048, height: 2048 },
});
}
registerImageToImageModels() {
this.registerModel({
id: 'flux-kontext-pro',
name: 'FLUX Kontext Pro',
description: 'Advanced image transformation with context understanding',
category: ModelCategory.IMAGE_TO_IMAGE,
endpoint: '/flux-kontext-pro',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 12,
creditsPerUse: 2,
parameters: z.object({
prompt: z.string().min(1).max(2000),
image: z.string(),
strength: z.number().min(0).max(1).default(0.75),
num_inference_steps: z.number().int().min(1).max(100).default(50),
guidance_scale: z.number().min(0).max(20).default(7.5),
seed: z.number().int().optional(),
}),
defaultParams: {
strength: 0.75,
num_inference_steps: 50,
guidance_scale: 7.5,
},
supportedFormats: ['png', 'jpeg', 'webp'],
});
this.registerModel({
id: 'controlnet',
name: 'ControlNet',
description: 'Precise control over image generation using reference images',
category: ModelCategory.IMAGE_TO_IMAGE,
endpoint: '/controlnet',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 10,
creditsPerUse: 2,
parameters: z.object({
prompt: z.string().min(1).max(2000),
negative_prompt: z.string().optional(),
image: z.string(),
control_type: z.enum(['canny', 'depth', 'pose', 'scribble', 'segmentation']).default('canny'),
control_strength: z.number().min(0).max(2).default(1),
num_inference_steps: z.number().int().min(1).max(100).default(30),
guidance_scale: z.number().min(0).max(20).default(7.5),
width: z.number().int().multipleOf(64).min(512).max(2048).default(1024),
height: z.number().int().multipleOf(64).min(512).max(2048).default(1024),
seed: z.number().int().optional(),
}),
defaultParams: {
control_type: 'canny',
control_strength: 1,
num_inference_steps: 30,
guidance_scale: 7.5,
width: 1024,
height: 1024,
},
supportedFormats: ['png', 'jpeg'],
maxDimensions: { width: 2048, height: 2048 },
});
}
registerVideoModels() {
this.registerModel({
id: 'veo-2',
name: 'Google Veo 2',
description: 'Advanced text-to-video generation with cinematic effects',
category: ModelCategory.VIDEO_GENERATION,
endpoint: '/veo-2',
apiVersion: 'v2',
outputType: OutputType.VIDEO,
estimatedTime: 120,
creditsPerUse: 10,
parameters: z.object({
prompt: z.string().min(1).max(2000),
duration: z.number().min(1).max(30).default(5),
fps: z.number().int().min(12).max(60).default(24),
aspect_ratio: z.enum(['16:9', '9:16', '1:1', '4:3']).default('16:9'),
quality: z.enum(['standard', 'high', 'ultra']).default('high'),
seed: z.number().int().optional(),
}),
defaultParams: {
duration: 5,
fps: 24,
aspect_ratio: '16:9',
quality: 'high',
},
supportedFormats: ['mp4', 'webm'],
});
this.registerModel({
id: 'kling-video',
name: 'Kling AI Video Generator',
description: 'Transform static images into animated videos',
category: ModelCategory.VIDEO_GENERATION,
endpoint: '/kling-video',
apiVersion: 'v2',
outputType: OutputType.VIDEO,
estimatedTime: 90,
creditsPerUse: 8,
parameters: z.object({
image: z.string(),
motion_prompt: z.string().min(1).max(500),
duration: z.number().min(1).max(10).default(3),
motion_strength: z.number().min(0).max(1).default(0.5),
fps: z.number().int().min(12).max(30).default(24),
seed: z.number().int().optional(),
}),
defaultParams: {
duration: 3,
motion_strength: 0.5,
fps: 24,
},
supportedFormats: ['mp4', 'gif'],
});
}
registerEnhancementModels() {
this.registerModel({
id: 'esrgan',
name: 'ESRGAN Upscaler',
description: 'AI-powered image upscaling (2x, 4x, 8x)',
category: ModelCategory.IMAGE_ENHANCEMENT,
endpoint: '/esrgan',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 15,
creditsPerUse: 2,
parameters: z.object({
image: z.string(),
scale: z.enum(['2', '4', '8']).default('4'),
face_enhance: z.boolean().default(false),
tile_size: z.number().int().min(0).max(512).default(0),
}),
defaultParams: {
scale: '4',
face_enhance: false,
tile_size: 0,
},
supportedFormats: ['png', 'jpeg', 'webp'],
});
this.registerModel({
id: 'bg-removal',
name: 'Background Removal',
description: 'Remove background from images with high accuracy',
category: ModelCategory.IMAGE_ENHANCEMENT,
endpoint: '/background-removal',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 5,
creditsPerUse: 1,
parameters: z.object({
image: z.string(),
return_mask: z.boolean().default(false),
alpha_matting: z.boolean().default(true),
alpha_matting_erode_size: z.number().int().min(0).max(40).default(10),
}),
defaultParams: {
return_mask: false,
alpha_matting: true,
alpha_matting_erode_size: 10,
},
supportedFormats: ['png'],
});
}
registerSpecializedModels() {
this.registerModel({
id: 'qr-generator',
name: 'Artistic QR Code Generator',
description: 'Generate QR codes with artistic styles',
category: ModelCategory.SPECIALIZED,
endpoint: '/qr-generator',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 8,
creditsPerUse: 1,
parameters: z.object({
qr_text: z.string().min(1).max(2000),
prompt: z.string().min(1).max(1000),
negative_prompt: z.string().optional(),
qr_code_strength: z.number().min(0.5).max(1.5).default(0.9),
guidance_scale: z.number().min(1).max(20).default(7.5),
num_inference_steps: z.number().int().min(10).max(100).default(30),
seed: z.number().int().optional(),
}),
defaultParams: {
qr_code_strength: 0.9,
guidance_scale: 7.5,
num_inference_steps: 30,
},
supportedFormats: ['png'],
});
this.registerModel({
id: 'face-to-sticker',
name: 'Face to Sticker',
description: 'Convert portrait photos into sticker-style art',
category: ModelCategory.SPECIALIZED,
endpoint: '/face-to-sticker',
apiVersion: 'v1',
outputType: OutputType.IMAGE,
estimatedTime: 6,
creditsPerUse: 1,
parameters: z.object({
image: z.string(),
style: z.enum(['cartoon', 'anime', 'pixel', 'sketch']).default('cartoon'),
background_remove: z.boolean().default(true),
enhance_face: z.boolean().default(true),
}),
defaultParams: {
style: 'cartoon',
background_remove: true,
enhance_face: true,
},
supportedFormats: ['png'],
});
}
getModel(id) {
return this.models.get(id);
}
getModelsByCategory(category) {
const modelIds = this.modelsByCategory.get(category) || [];
return modelIds.map(id => this.models.get(id)).filter(Boolean);
}
getAllModels() {
return Array.from(this.models.values());
}
validateParameters(modelId, params) {
const model = this.getModel(modelId);
if (!model) {
return { success: false, error: `Model ${modelId} not found` };
}
try {
const validated = model.parameters.parse(params);
return { success: true, data: validated };
}
catch (error) {
if (error instanceof z.ZodError) {
const issues = error.issues.map(i => `${i.path.join('.')}: ${i.message}`).join(', ');
return { success: false, error: `Invalid parameters: ${issues}` };
}
return { success: false, error: 'Parameter validation failed' };
}
}
recommendModel(category) {
const models = category
? this.getModelsByCategory(category)
: this.getAllModels();
if (models.length === 0)
return null;
return models[0] || null;
}
}
export const modelRegistry = new ModelRegistry();
//# sourceMappingURL=registry-old.js.map