feat: support dalle image variation

This commit is contained in:
ikechan8370 2023-02-25 15:23:47 +08:00
parent 38de2ce714
commit 5fceef37ec
5 changed files with 127 additions and 6 deletions

View file

@ -279,7 +279,7 @@ export class chatgpt extends plugin {
}
if (reply) {
for (let val of reply) {
if (val.type == 'image') {
if (val.type === 'image') {
e.img = [val.url]
break
}

View file

@ -1,7 +1,7 @@
import plugin from '../../../lib/plugins/plugin.js'
import { segment } from 'oicq'
import { createImage } from '../utils/dalle.js'
import { makeForwardMsg} from "../utils/common.js";
import { createImage, imageVariation } from '../utils/dalle.js'
import { makeForwardMsg } from '../utils/common.js'
import _ from 'lodash'
export class dalle extends plugin {
@ -15,6 +15,10 @@ export class dalle extends plugin {
{
reg: '#(chatgpt|ChatGPT|dalle|Dalle)(绘图|画图)',
fnc: 'draw'
},
{
reg: '#(chatgpt|ChatGPT|dalle|Dalle)(修图|图片变形|改图)',
fnc: 'variation'
}
]
})
@ -37,7 +41,7 @@ export class dalle extends plugin {
this.reply('大小不符合要求必须是256x256/512x512/1024x1024中的一个')
return false
}
await redis.set(`CHATGPT:DRAW:${e.sender.user_id}`, 'c', {EX: 30})
await redis.set(`CHATGPT:DRAW:${e.sender.user_id}`, 'c', { EX: 30 })
let priceMap = {
'1024x1024': 0.02,
'512x512': 0.018,
@ -57,7 +61,51 @@ export class dalle extends plugin {
this.reply(`绘图失败: ${err}`, true)
await redis.del(`CHATGPT:DRAW:${e.sender.user_id}`)
}
}
async variation (e) {
let ttl = await redis.ttl(`CHATGPT:VARIATION:${e.sender.user_id}`)
if (ttl > 0 && !e.isMaster) {
this.reply(`冷却中,请${ttl}秒后再试`)
return false
}
let imgUrl
if (e.source) {
let reply
if (e.isGroup) {
reply = (await e.group.getChatHistory(e.source.seq, 1)).pop()?.message
} else {
reply = (await e.friend.getChatHistory(e.source.time, 1)).pop()?.message
}
if (reply) {
for (let val of reply) {
if (val.type === 'image') {
console.log(val)
imgUrl = val.url
break
}
}
}
} else if (e.img) {
imgUrl = e.img[0]?.url
}
if (!imgUrl) {
this.reply('图呢?')
return false
}
await redis.set(`CHATGPT:VARIATION:${e.sender.user_id}`, 'c', { EX: 30 })
await this.reply('正在为您生成图片变形,请稍候……')
try {
let images = (await imageVariation(imgUrl)).map(image => segment.image(`base64://${image}`))
if (images.length > 1) {
this.reply(await makeForwardMsg(e, images))
} else {
this.reply(images[0], true)
}
} catch (err) {
console.log(err)
this.reply(`绘图失败: ${err}`, true)
await redis.del(`CHATGPT:VARIATION:${e.sender.user_id}`)
}
}
}

View file

@ -13,6 +13,7 @@
"puppeteer-extra-plugin-recaptcha": "^3.6.6",
"puppeteer-extra-plugin-stealth": "^2.11.1",
"random": "^4.1.0",
"sharp": "^0.31.3",
"undici": "^5.19.1",
"uuid": "^9.0.0"
}

View file

@ -2,6 +2,8 @@
// import stripMarkdown from 'strip-markdown'
import { exec } from 'child_process'
import lodash from 'lodash'
import fs from 'node:fs'
import path from 'node:path'
// export function markdownToText (markdown) {
// return remark()
// .use(stripMarkdown)
@ -205,3 +207,14 @@ async function execSync (cmd) {
})
})
}
export function mkdirs (dirname) {
if (fs.existsSync(dirname)) {
return true
} else {
if (mkdirs(path.dirname(dirname))) {
fs.mkdirSync(dirname)
return true
}
}
}

View file

@ -1,6 +1,8 @@
import { Configuration, OpenAIApi } from 'openai'
import { Config } from './config.js'
import fs from 'fs'
import { mkdirs } from './common.js'
import sharp from 'sharp'
export async function createImage (prompt, n = 1, size = '512x512') {
const configuration = new Configuration({
apiKey: Config.apiKey
@ -17,3 +19,60 @@ export async function createImage (prompt, n = 1, size = '512x512') {
})
return response.data.data?.map(pic => pic.b64_json)
}
export async function imageVariation (imageUrl, n = 1, size = '512x512') {
const configuration = new Configuration({
apiKey: Config.apiKey
})
const openai = new OpenAIApi(configuration)
if (Config.debug) {
logger.info({ imageUrl, n, size })
}
const imageResponse = await fetch(imageUrl)
const fileType = imageResponse.headers.get('Content-Type').split('/')[1]
let fileLoc = `data/chatgpt/imagesAccept/${Date.now()}.${fileType}`
mkdirs('data/chatgpt/imagesAccept')
const blob = await imageResponse.blob()
const arrayBuffer = await blob.arrayBuffer()
const buffer = Buffer.from(arrayBuffer)
await fs.writeFileSync(fileLoc, buffer)
let croppedFileLoc = `data/chatgpt/imagesAccept/${Date.now()}_cropped.png`
await resizeAndCropImage(fileLoc, croppedFileLoc, 512)
const response = await openai.createImageVariation(
fs.createReadStream(croppedFileLoc),
n,
size,
'b64_json'
)
if (response.status !== 200) {
console.log(response.data.error)
}
await fs.unlinkSync(fileLoc)
await fs.unlinkSync(croppedFileLoc)
return response.data.data?.map(pic => pic.b64_json)
}
async function resizeAndCropImage (inputFilePath, outputFilePath, size = 512) {
// Determine the maximum dimension of the input image
const metadata = await sharp(inputFilePath).metadata()
const maxDimension = Math.max(metadata.width, metadata.height)
logger.mark(`original picture size is ${metadata.width} x ${metadata.height}`)
// Calculate the required dimensions for the output image
const outputWidth = size * metadata.width / maxDimension
const outputHeight = size * metadata.height / maxDimension
// Resize the image to the required dimensions
await sharp(inputFilePath)
.resize(outputWidth, outputHeight, {
fit: 'contain',
background: { r: 255, g: 255, b: 255, alpha: 1 }
})
.resize(size, size, { fit: 'cover', position: 'center' })
.png()
.toFile(outputFilePath)
console.log('Image resized successfully!')
console.log('Image resized and cropped successfully!')
}