replicate-flux-mcp
Version:
MCP for Replicate Flux Model
106 lines (105 loc) • 4.46 kB
JavaScript
import { replicate } from "../services/replicate.js";
import { handleError } from "../utils/error.js";
import { outputToBase64 } from "../utils/image.js";
import { CONFIG } from "../config/index.js";
export const registerGenerateImageVariantsTool = async (input) => {
const { prompt, num_variants, seed, support_image_mcp_response_type, prompt_variations, variation_mode, ...commonParams } = input;
try {
let effectiveVariants = num_variants;
let usingPromptVariations = false;
// Decide if we're using prompt variations
if (prompt_variations && prompt_variations.length > 0) {
usingPromptVariations = true;
// If using prompt variations, number of variants is limited by available variations
effectiveVariants = Math.min(num_variants, prompt_variations.length);
}
// Process all variants in parallel
const generationPromises = Array.from({ length: effectiveVariants }, (_, index) => {
// If seed is provided, create deterministic variants by adding the index
const variantSeed = seed !== undefined ? seed + index : undefined;
// Determine which prompt to use for this variant
let variantPrompt = prompt;
if (usingPromptVariations) {
const variation = prompt_variations[index];
if (variation_mode === "append") {
variantPrompt = `${prompt} ${variation}`;
}
else {
// 'replace' mode
variantPrompt = variation;
}
}
return replicate
.run(CONFIG.imageModelId, {
input: {
prompt: variantPrompt,
seed: variantSeed,
...commonParams,
},
})
.then((outputs) => {
const [output] = outputs;
const imageUrl = output.url();
if (support_image_mcp_response_type) {
return outputToBase64(output).then((imageBase64) => ({
variantIndex: index + 1,
imageUrl,
imageBase64,
usedPrompt: variantPrompt,
}));
}
return {
variantIndex: index + 1,
imageUrl,
usedPrompt: variantPrompt,
};
});
});
// Wait for all variant generation to complete
const results = (await Promise.all(generationPromises));
// Build response content
const responseContent = [];
// Add intro text - different based on whether we're using prompt variations
if (usingPromptVariations) {
responseContent.push({
type: "text",
text: `Generated ${results.length} variants of "${prompt}" using custom prompt variations (${variation_mode} mode)`,
});
}
else {
responseContent.push({
type: "text",
text: `Generated ${results.length} variants of: "${prompt}" using seed variations`,
});
}
// Add each variant with its index and prompt info
for (const result of results) {
// Build an appropriate description based on variant type
let variantDescription = `Variant #${result.variantIndex}`;
if (usingPromptVariations) {
variantDescription += `\nPrompt: "${result.usedPrompt}"`;
}
else if (seed !== undefined) {
variantDescription += ` (seed: ${seed + (result.variantIndex - 1)})`;
}
variantDescription += `\nImage URL: ${result.imageUrl}`;
responseContent.push({
type: "text",
text: `\n\n${variantDescription}`,
});
if (support_image_mcp_response_type && result.imageBase64) {
responseContent.push({
type: "image",
data: result.imageBase64,
mimeType: `image/${input.output_format === "jpg" ? "jpeg" : input.output_format}`,
});
}
}
return {
content: responseContent,
};
}
catch (error) {
handleError(error);
}
};