ai-image
Version:
A unified wrapper for image generation inference APIs (OpenAI and Replicate) with CLI support, useful for MCP and sub-processes
168 lines • 7.47 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.Replicate = exports.OpenAI = exports.ImageGenerator = void 0;
const openai_1 = __importDefault(require("openai"));
exports.OpenAI = openai_1.default;
const replicate_1 = __importDefault(require("replicate"));
exports.Replicate = replicate_1.default;
const promises_1 = __importDefault(require("fs/promises"));
const path_1 = __importDefault(require("path"));
const dotenv_1 = __importDefault(require("dotenv"));
// Load environment variables
dotenv_1.default.config();
class ImageGenerator {
provider;
openai;
replicate;
outputDir;
outputFilename;
constructor(options) {
this.provider = options.provider;
this.outputDir = options.outputDir || process.cwd();
this.outputFilename = options.outputFilename;
const apiKey = options.apiKey || this.getApiKeyFromEnv(options.provider);
if (!apiKey) {
throw new Error(`[AI-IMAGE] API key for ${options.provider} not found. Please provide it via parameter or environment variable.`);
}
switch (options.provider) {
case 'openai':
this.openai = new openai_1.default({ apiKey });
break;
case 'replicate':
this.replicate = new replicate_1.default({ auth: apiKey });
break;
default:
throw new Error(`Unsupported provider: ${options.provider}`);
}
}
getApiKeyFromEnv(provider) {
switch (provider) {
case 'openai':
return process.env.OPENAI_API_KEY;
case 'replicate':
return process.env.REPLICATE_API_TOKEN;
default:
return undefined;
}
}
sanitizeFilename(filename) {
// Remove or replace invalid characters for filenames
return filename
.replace(/[<>:"/\\|?*\x00-\x1F]/g, '')
.replace(/\s+/g, '_')
.substring(0, 200); // Limit length
}
async getOutputPath(prompt, index = 0, extension = 'png') {
let filename;
if (this.outputFilename) {
// If multiple images, append index
const ext = path_1.default.extname(this.outputFilename);
const name = path_1.default.basename(this.outputFilename, ext);
filename = index > 0 ? `${name}_${index}${ext || `.${extension}`}` : `${name}${ext || `.${extension}`}`;
}
else {
// Use sanitized prompt as filename
const sanitized = this.sanitizeFilename(prompt);
filename = index > 0 ? `${sanitized}_${index}.${extension}` : `${sanitized}.${extension}`;
}
let outputPath = path_1.default.join(this.outputDir, filename);
// Handle filename collisions
let counter = 1;
while (await this.fileExists(outputPath)) {
const ext = path_1.default.extname(filename);
const name = path_1.default.basename(filename, ext);
const newFilename = `${name} ${counter}${ext}`;
outputPath = path_1.default.join(this.outputDir, newFilename);
counter++;
}
return outputPath;
}
async fileExists(filePath) {
try {
await promises_1.default.access(filePath);
return true;
}
catch {
return false;
}
}
async generate(options) {
const savedPaths = [];
try {
if (this.provider === 'openai' && this.openai) {
const effectiveModel = options.model || 'gpt-image-1';
const generateParams = {
model: effectiveModel,
prompt: options.prompt,
size: options.size || '1024x1024',
quality: options.quality || 'auto',
n: options.n || 1,
// response_format: 'b64_json'
};
// Add optional parameters
if (options.format && options.format !== 'png') {
generateParams.output_format = options.format;
}
if (options.compression && (options.format === 'jpeg' || options.format === 'webp')) {
generateParams.output_compression = options.compression;
}
if (options.background) {
generateParams.background = options.background;
}
if (options.debug) {
console.log('🐛 Debug - Full request parameters:');
console.log(JSON.stringify(generateParams, null, 2));
}
const result = await this.openai.images.generate(generateParams);
// Save each generated image
for (let i = 0; i < (result.data?.length || 0); i++) {
const imageData = result.data?.[i];
if (imageData?.b64_json) {
const imageBytes = Buffer.from(imageData.b64_json, 'base64');
const fileExtension = options.format || 'png';
const outputPath = await this.getOutputPath(options.prompt, i, fileExtension);
await promises_1.default.writeFile(outputPath, imageBytes);
savedPaths.push(outputPath);
}
}
}
else if (this.provider === 'replicate' && this.replicate) {
// Default to SDXL model if not specified
const model = options.model || 'stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b';
const replicateInput = {
prompt: options.prompt,
width: parseInt(options.size?.split('x')[0] || '1024'),
height: parseInt(options.size?.split('x')[1] || '1024'),
num_outputs: options.n || 1
};
if (options.debug) {
console.log('🐛 Debug - Replicate model input:');
console.log(JSON.stringify(replicateInput, null, 2));
}
const output = await this.replicate.run(model, {
input: replicateInput
});
// Save each generated image
for (let i = 0; i < output.length; i++) {
const imageUrl = output[i];
const response = await fetch(imageUrl);
const arrayBuffer = await response.arrayBuffer();
const buffer = Buffer.from(arrayBuffer);
const outputPath = await this.getOutputPath(options.prompt, i);
await promises_1.default.writeFile(outputPath, buffer);
savedPaths.push(outputPath);
}
}
return savedPaths;
}
catch (error) {
console.error(`❌ Image generation failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
throw new Error(`Failed to generate image: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
}
exports.ImageGenerator = ImageGenerator;
//# sourceMappingURL=index.js.map