439 lines
13 KiB
JavaScript
439 lines
13 KiB
JavaScript
#!/usr/bin/env node
|
||
|
||
/**
|
||
* MJ Image Generator - Midjourney 图片生成工具
|
||
*
|
||
* 功能：
|
||
* - 提交 imagine 任务
|
||
* - 轮询直到完成
|
||
* - 下载 4 合 1 结果图
|
||
* - 自动拆分为 4 张独立图片
|
||
* - 支持参考图（图生图）
|
||
*
|
||
* 用法：
|
||
* node mj-image-generator.js "a cute cat" -o ./output
|
||
* node mj-image-generator.js "cyberpunk city" -o ./output -r http://example.com/ref.jpg
|
||
* node mj-image-generator.js batch ./prompts.txt -o ./output
|
||
*/
|
||
|
||
const fs = require('fs')
|
||
const path = require('path')
|
||
const https = require('https')
|
||
const http = require('http')
|
||
const sharp = require('sharp')
|
||
|
||
// ============================================================================
|
||
// 配置
|
||
// ============================================================================
|
||
|
||
function loadConfig() {
|
||
const configPath = path.join(__dirname, '..', '..', 'config.json')
|
||
if (fs.existsSync(configPath)) {
|
||
return JSON.parse(fs.readFileSync(configPath, 'utf-8'))
|
||
}
|
||
return {}
|
||
}
|
||
|
||
const cfg = loadConfig()
|
||
|
||
const Config = {
|
||
baseUrl: cfg.mjApiBaseUrl,
|
||
apiKey: cfg.mjApiKey || '',
|
||
pollInterval: 5000, // 轮询间隔 5 秒
|
||
maxPollTime: 300000, // 最大等待 5 分钟
|
||
}
|
||
|
||
// ============================================================================
|
||
// API 调用
|
||
// ============================================================================
|
||
|
||
const MJApi = {
|
||
/**
|
||
* 提交 imagine 任务
|
||
*/
|
||
async submit(prompt, options = {}) {
|
||
const { referenceImages = [], botType = 'mj', aspectRatio = '', mjParams = '', styleWeight = 100 } = options
|
||
|
||
let finalPrompt = prompt
|
||
if (referenceImages.length > 0) {
|
||
const srefSection = `--sref ${referenceImages.join(' ')} --sw ${styleWeight}`
|
||
finalPrompt = `${prompt} ${srefSection}`.trim()
|
||
}
|
||
if (aspectRatio) {
|
||
finalPrompt = `${finalPrompt} --ar ${aspectRatio}`
|
||
}
|
||
if (mjParams) {
|
||
finalPrompt = `${finalPrompt} ${mjParams}`
|
||
}
|
||
|
||
const body = {
|
||
prompt: finalPrompt,
|
||
base64Array: [],
|
||
botType,
|
||
}
|
||
|
||
console.log(`\n📡 提交 MJ 任务`)
|
||
console.log(` 提示词: ${finalPrompt.substring(0, 80)}...`)
|
||
console.log(` 参考图: ${referenceImages.length} 张`)
|
||
|
||
const res = await fetch(`${Config.baseUrl}/mj/submit/imagine`, {
|
||
method: 'POST',
|
||
headers: {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': `Bearer ${Config.apiKey}`,
|
||
},
|
||
body: JSON.stringify(body),
|
||
})
|
||
|
||
const result = await res.json()
|
||
|
||
if (result.code !== 1 && result.code !== '1') {
|
||
throw new Error(`MJ 提交失败: code=${result.code}, ${result.description || JSON.stringify(result)}`)
|
||
}
|
||
|
||
const taskId = result.result
|
||
console.log(` 任务 ID: ${taskId}`)
|
||
return taskId
|
||
},
|
||
|
||
/**
|
||
* 轮询任务状态
|
||
*/
|
||
async poll(taskId) {
|
||
const startTime = Date.now()
|
||
let lastProgress = ''
|
||
|
||
console.log(`\n⏳ 等待 MJ 生成...`)
|
||
|
||
while (Date.now() - startTime < Config.maxPollTime) {
|
||
const res = await fetch(`${Config.baseUrl}/mj/task/${taskId}/fetch`, {
|
||
headers: { 'Authorization': `Bearer ${Config.apiKey}` },
|
||
})
|
||
|
||
const task = await res.json()
|
||
const status = task.status
|
||
|
||
if (status === 'SUCCESS') {
|
||
console.log(`\n✅ 生成完成!`)
|
||
console.log(` 图片 URL: ${task.imageUrl}`)
|
||
return {
|
||
success: true,
|
||
imageUrl: task.imageUrl,
|
||
prompt: task.prompt || task.promptEn,
|
||
}
|
||
}
|
||
|
||
if (status === 'FAILURE') {
|
||
const errMsg = task.failReason || '未知原因'
|
||
throw new Error(`MJ 生成失败: ${errMsg}`)
|
||
}
|
||
|
||
// 显示进度
|
||
const progress = task.progress || ''
|
||
if (progress !== lastProgress) {
|
||
lastProgress = progress
|
||
process.stdout.write(` 进度: ${progress}% 状态: ${status}\r`)
|
||
}
|
||
|
||
await new Promise(r => setTimeout(r, Config.pollInterval))
|
||
}
|
||
|
||
throw new Error(`MJ 生成超时 (${Config.maxPollTime / 1000}s)`)
|
||
},
|
||
}
|
||
|
||
// ============================================================================
|
||
// 图片处理
|
||
// ============================================================================
|
||
|
||
const ImageUtils = {
|
||
/**
|
||
* 下载图片到本地
|
||
*/
|
||
async download(url, outputPath) {
|
||
const protocol = url.startsWith('https') ? https : http
|
||
|
||
return new Promise((resolve, reject) => {
|
||
const file = fs.createWriteStream(outputPath)
|
||
protocol.get(url, (response) => {
|
||
// 处理重定向
|
||
if (response.statusCode >= 300 && response.statusCode < 400 && response.headers.location) {
|
||
file.close()
|
||
fs.unlinkSync(outputPath)
|
||
return ImageUtils.download(response.headers.location, outputPath).then(resolve).catch(reject)
|
||
}
|
||
response.pipe(file)
|
||
file.on('finish', () => {
|
||
file.close()
|
||
resolve(outputPath)
|
||
})
|
||
}).on('error', (err) => {
|
||
file.close()
|
||
if (fs.existsSync(outputPath)) fs.unlinkSync(outputPath)
|
||
reject(err)
|
||
})
|
||
})
|
||
},
|
||
|
||
/**
|
||
* 将 4 合 1 图片拆分为 4 张独立图片
|
||
*/
|
||
async split4(gridImagePath, outputDir, prefix = 'mj') {
|
||
const image = sharp(gridImagePath)
|
||
const metadata = await image.metadata()
|
||
const { width, height } = metadata
|
||
|
||
// MJ 4 合 1 是 2x2 网格，每格约一半
|
||
const halfW = Math.floor(width / 2)
|
||
const halfH = Math.floor(height / 2)
|
||
|
||
const positions = [
|
||
{ name: `${prefix}_1`, x: 0, y: 0 }, // 左上
|
||
{ name: `${prefix}_2`, x: halfW, y: 0 }, // 右上
|
||
{ name: `${prefix}_3`, x: 0, y: halfH }, // 左下
|
||
{ name: `${prefix}_4`, x: halfW, y: halfH }, // 右下
|
||
]
|
||
|
||
const files = []
|
||
for (const pos of positions) {
|
||
const outputPath = path.join(outputDir, `${pos.name}.png`)
|
||
await sharp(gridImagePath)
|
||
.extract({ left: pos.x, top: pos.y, width: halfW, height: halfH })
|
||
.toFile(outputPath)
|
||
files.push(outputPath)
|
||
}
|
||
|
||
return files
|
||
},
|
||
}
|
||
|
||
// ============================================================================
|
||
// 核心流程
|
||
// ============================================================================
|
||
|
||
async function generate(prompt, options = {}) {
|
||
const {
|
||
outputDir = './output',
|
||
referenceImages = [],
|
||
aspectRatio = '',
|
||
mjParams = '',
|
||
split = true,
|
||
keepGrid = false,
|
||
} = options
|
||
|
||
if (!Config.apiKey) {
|
||
throw new Error('未配置 mjApiKey，请在 config.json 中添加')
|
||
}
|
||
|
||
fs.mkdirSync(outputDir, { recursive: true })
|
||
|
||
// 1. 提交任务
|
||
const taskId = await MJApi.submit(prompt, { referenceImages, aspectRatio, mjParams, styleWeight: options.styleWeight })
|
||
|
||
// 2. 轮询等待
|
||
const result = await MJApi.poll(taskId)
|
||
|
||
// 3. 下载
|
||
const timestamp = new Date().toISOString().replace(/[:.]/g, '-')
|
||
const gridFile = path.join(outputDir, `${timestamp}_grid.png`)
|
||
await ImageUtils.download(result.imageUrl, gridFile)
|
||
console.log(` 下载完成: ${gridFile}`)
|
||
|
||
// 4. 拆分
|
||
const allFiles = [gridFile]
|
||
if (split) {
|
||
const prefix = timestamp
|
||
const splitFiles = await ImageUtils.split4(gridFile, outputDir, prefix)
|
||
allFiles.push(...splitFiles)
|
||
console.log(` 拆分完成: ${splitFiles.length} 张图片`)
|
||
|
||
if (!keepGrid) {
|
||
fs.unlinkSync(gridFile)
|
||
allFiles.shift()
|
||
console.log(` 已删除网格图`)
|
||
}
|
||
}
|
||
|
||
return {
|
||
success: true,
|
||
taskId,
|
||
imageUrl: result.imageUrl,
|
||
files: allFiles,
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 批量生成（并发提交 + 并行轮询）
|
||
*/
|
||
async function batchGenerate(prompts, options = {}) {
|
||
const concurrency = options.concurrency || prompts.length
|
||
const results = new Array(prompts.length).fill(null)
|
||
|
||
// 分批提交，每批 concurrency 个并行
|
||
for (let batchStart = 0; batchStart < prompts.length; batchStart += concurrency) {
|
||
const batchEnd = Math.min(batchStart + concurrency, prompts.length)
|
||
const batchIndices = []
|
||
for (let i = batchStart; i < batchEnd; i++) batchIndices.push(i)
|
||
|
||
// 并行提交所有任务
|
||
console.log(`\n📡 批量提交 [${batchStart + 1}-${batchEnd}/${prompts.length}]...`)
|
||
const taskIds = await Promise.all(batchIndices.map(async (i) => {
|
||
try {
|
||
const taskId = await MJApi.submit(prompts[i], {
|
||
referenceImages: options.referenceImages,
|
||
aspectRatio: options.aspectRatio,
|
||
mjParams: options.mjParams,
|
||
styleWeight: options.styleWeight,
|
||
})
|
||
return { i, taskId }
|
||
} catch (err) {
|
||
console.error(` [${i + 1}] ❌ 提交失败: ${err.message}`)
|
||
results[i] = { success: false, prompt: prompts[i], error: err.message }
|
||
return { i, taskId: null }
|
||
}
|
||
}))
|
||
|
||
// 过滤出成功的任务，并行轮询
|
||
const activeTasks = taskIds.filter(t => t.taskId)
|
||
console.log(` 已提交 ${activeTasks.length} 个任务，并行等待生成...\n`)
|
||
|
||
const pollResults = await Promise.all(activeTasks.map(async ({ i, taskId }) => {
|
||
try {
|
||
const result = await MJApi.poll(taskId)
|
||
// 下载 + 拆分
|
||
const timestamp = new Date().toISOString().replace(/[:.]/g, '-')
|
||
const gridFile = path.join(options.outputDir || './output', `${timestamp}_grid.png`)
|
||
fs.mkdirSync(path.dirname(gridFile), { recursive: true })
|
||
await ImageUtils.download(result.imageUrl, gridFile)
|
||
|
||
const allFiles = [gridFile]
|
||
if (options.split !== false) {
|
||
const prefix = timestamp
|
||
const splitFiles = await ImageUtils.split4(gridFile, options.outputDir || './output', prefix)
|
||
allFiles.push(...splitFiles)
|
||
if (!options.keepGrid) {
|
||
fs.unlinkSync(gridFile)
|
||
allFiles.shift()
|
||
}
|
||
}
|
||
|
||
console.log(` [${i + 1}/${prompts.length}] ✅ 完成`)
|
||
return { i, success: true, prompt: prompts[i], taskId, imageUrl: result.imageUrl, files: allFiles }
|
||
} catch (err) {
|
||
console.error(` [${i + 1}/${prompts.length}] ❌ 失败: ${err.message}`)
|
||
return { i, success: false, prompt: prompts[i], error: err.message }
|
||
}
|
||
}))
|
||
|
||
for (const r of pollResults) results[r.i] = r
|
||
}
|
||
|
||
const ok = results.filter(r => r && r.success).length
|
||
console.log(`\n✨ 批量完成: ${ok}/${prompts.length} 成功`)
|
||
|
||
return results
|
||
}
|
||
|
||
// ============================================================================
|
||
// CLI
|
||
// ============================================================================
|
||
|
||
function showHelp() {
|
||
console.log(`
|
||
🎨 MJ Image Generator - Midjourney 图片生成工具
|
||
|
||
用法:
|
||
node mj-image-generator.js <prompt> [options]
|
||
node mj-image-generator.js batch <file> [options]
|
||
|
||
选项:
|
||
-o, --output <dir> 输出目录 (默认: ./output)
|
||
-r, --ref <urls> 参考图 URL，逗号分隔
|
||
-a, --ar <ratio> 宽高比 (1:1, 16:9, 9:16, 3:4, 4:3 等)
|
||
-c, --concurrency <n> 并发数 (默认: 全部并行)
|
||
--no-split 不拆分 4 合 1
|
||
--keep-grid 保留原始网格图
|
||
-h, --help 帮助
|
||
|
||
示例:
|
||
node mj-image-generator.js "a cute cat" -o ./cats
|
||
node mj-image-generator.js "cyberpunk city" -a 16:9
|
||
node mj-image-generator.js "portrait" -r http://img.com/ref.jpg -a 9:16
|
||
node mj-image-generator.js batch ./prompts.txt -o ./batch
|
||
`)
|
||
}
|
||
|
||
async function main() {
|
||
const args = process.argv.slice(2)
|
||
|
||
if (args.includes('-h') || args.includes('--help') || args.length === 0) {
|
||
showHelp()
|
||
return
|
||
}
|
||
|
||
let command = 'generate'
|
||
let params = []
|
||
const options = { outputDir: './output', split: true, keepGrid: false, referenceImages: [], aspectRatio: '', concurrency: 0, mjParams: '', styleWeight: 100 }
|
||
|
||
let i = 0
|
||
if (args[0] === 'batch') {
|
||
command = 'batch'
|
||
i = 1
|
||
}
|
||
|
||
while (i < args.length) {
|
||
const arg = args[i]
|
||
if (arg === '-o' || arg === '--output') {
|
||
options.outputDir = args[++i]
|
||
} else if (arg === '-a' || arg === '--ar') {
|
||
options.aspectRatio = args[++i]
|
||
} else if (arg === '-r' || arg === '--ref') {
|
||
options.referenceImages = args[++i].split(',').map(s => s.trim()).filter(Boolean)
|
||
} else if (arg === '--no-split') {
|
||
options.split = false
|
||
} else if (arg === '--keep-grid') {
|
||
options.keepGrid = true
|
||
} else if (arg === '-c' || arg === '--concurrency') {
|
||
options.concurrency = parseInt(args[++i], 10) || 0
|
||
} else if (arg === '--mj-params') {
|
||
options.mjParams = args[++i]
|
||
} else if (arg === '--sw') {
|
||
options.styleWeight = parseInt(args[++i], 10) || 100
|
||
} else {
|
||
params.push(arg)
|
||
}
|
||
i++
|
||
}
|
||
|
||
if (command === 'batch') {
|
||
const filePath = params[0]
|
||
if (!filePath || !fs.existsSync(filePath)) {
|
||
console.error('请提供提示词文件路径')
|
||
process.exit(1)
|
||
}
|
||
const prompts = fs.readFileSync(filePath, 'utf-8')
|
||
.split('\n').filter(l => l.trim()).map(l => l.trim())
|
||
await batchGenerate(prompts, options)
|
||
} else {
|
||
const prompt = params.join(' ')
|
||
if (!prompt) {
|
||
console.error('请提供提示词')
|
||
process.exit(1)
|
||
}
|
||
await generate(prompt, options)
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// 导出
|
||
// ============================================================================
|
||
|
||
module.exports = { generate, batchGenerate, MJApi, ImageUtils }
|
||
|
||
if (require.main === module) {
|
||
main().catch(err => {
|
||
console.error(`\n❌ 错误: ${err.message}`)
|
||
process.exit(1)
|
||
})
|
||
}
|