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)
|
|||
|
|
})
|
|||
|
|
}
|