feat: support DALLE2 image edit

apps/draw.js

@@ -1,6 +1,6 @@
 import plugin from '../../../lib/plugins/plugin.js'
 import { segment } from 'oicq'
-import { createImage, imageVariation } from '../utils/dalle.js'
+import { createImage, editImage, imageVariation } from '../utils/dalle.js'
 import { makeForwardMsg } from '../utils/common.js'
 import _ from 'lodash'
 
@@ -13,16 +13,20 @@ export class dalle extends plugin {
       priority: 500,
       rule: [
         {
-          reg: '#(chatgpt|ChatGPT|dalle|Dalle)(绘图|画图)',
+          reg: '^#(chatgpt|ChatGPT|dalle|Dalle)(绘图|画图)',
           fnc: 'draw'
         },
         {
-          reg: '#(chatgpt|ChatGPT|dalle|Dalle)(修图|图片变形|改图)',
+          reg: '^#(chatgpt|ChatGPT|dalle|Dalle)(修图|图片变形|改图)$',
           fnc: 'variation'
         },
         {
-          reg: '#(搞|改)(她|他)头像',
+          reg: '^#(搞|改)(她|他)头像',
           fnc: 'avatarVariation'
+        },
+        {
+          reg: '^#(chatgpt|dalle)编辑图片',
+          fnc: 'edit'
         }
       ]
     })
@@ -139,4 +143,67 @@ export class dalle extends plugin {
       }
     }
   }
+
+  async edit (e) {
+    let ttl = await redis.ttl(`CHATGPT:EDIT:${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) {
+      console.log(e.img)
+      imgUrl = e.img[0]
+    }
+    if (!imgUrl) {
+      this.reply('图呢？')
+      return false
+    }
+    await redis.set(`CHATGPT:EDIT:${e.sender.user_id}`, 'c', { EX: 30 })
+    await this.reply('正在为您编辑图片，请稍候……')
+
+    let command = _.trimStart(e.msg, '#chatgpt编辑图片')
+    command = _.trimStart(command, '#dalle编辑图片')
+    // command = 'A bird on it/100,100,300,200/2/512x512'
+    let args = command.split('/')
+    let [prompt = '', position = '', num = '1', size = '512x512'] = args.slice(0, 4)
+    if (!prompt || !position) {
+      this.reply('编辑图片必须填写prompt和涂抹位置.参考格式：A bird on it/100,100,300,200/2/512x512')
+      return false
+    }
+    num = parseInt(num, 10)
+    if (num > 5) {
+      this.reply('太多啦！你要花光我的余额吗！')
+      return false
+    }
+    try {
+      let images = (await editImage(imgUrl, position.split(',').map(p => parseInt(p, 10)), prompt, num, size))
+        .map(image => segment.image(`base64://${image}`))
+      if (images.length > 1) {
+        this.reply(await makeForwardMsg(e, images, prompt))
+      } else {
+        this.reply(images[0], true)
+      }
+    } catch (err) {
+      logger.error(err)
+      this.reply(`图片编辑失败: ${err}`, true)
+      await redis.del(`CHATGPT:EDIT:${e.sender.user_id}`)
+    }
+  }
 }

package.json

@@ -17,6 +17,7 @@
     "uuid": "^9.0.0"
   },
   "optionalDependencies": {
-    "sharp": "^0.31.3"
+    "sharp": "^0.31.3",
+    "jimp": "^0.22.7"
   }
 }

utils/dalle.js

@@ -83,3 +83,74 @@ async function resizeAndCropImage (inputFilePath, outputFilePath, size = 512) {
 
   console.log('Image resized and cropped successfully!')
 }
+
+export async function editImage (originalImage, mask = [], prompt, num = 1, size = '512x512') {
+  const configuration = new Configuration({
+    apiKey: Config.apiKey
+  })
+  const openai = new OpenAIApi(configuration)
+  if (Config.debug) {
+    logger.info({ originalImage, mask, num, size })
+  }
+  const imageResponse = await fetch(originalImage)
+  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)
+  let maskFileLoc = await createMask(croppedFileLoc, mask)
+  let response = await openai.createImageEdit(
+    fs.createReadStream(croppedFileLoc),
+    fs.createReadStream(maskFileLoc),
+    prompt,
+    num,
+    size,
+    'b64_json'
+  )
+  if (response.status !== 200) {
+    console.log(response.data.error)
+  }
+  await fs.unlinkSync(fileLoc)
+  await fs.unlinkSync(croppedFileLoc)
+  await fs.unlinkSync(maskFileLoc)
+  return response.data.data?.map(pic => pic.b64_json)
+}
+
+async function createMask (inputFilePath, mask = []) {
+  let sharp, Jimp
+  try {
+    sharp = (await import('sharp')).default
+  } catch (e) {
+    logger.error('sharp未安装，请执行 pnpm install sharp@0.31.3')
+    throw new Error('sharp未安装，请执行 pnpm install sharp@0.31.3')
+  }
+  try {
+    Jimp = (await import('jimp')).default
+  } catch (e) {
+    logger.error('jimp未安装，请执行 pnpm install jimp')
+    throw new Error('jimp未安装，请执行 pnpm install jimp')
+  }
+  let image = await sharp(inputFilePath)
+    .png()
+    .ensureAlpha()
+    .toBuffer()
+    .then(inputData => {
+      // Load the PNG input data with Jimp
+      return Jimp.read(inputData)
+    })
+  let [x, y, width, height] = mask
+  // Set the transparency for a specified rectangular area
+  image.scan(x, y, width, height, function (x, y, idx) {
+    this.bitmap.data[idx + 3] = 0 // set alpha to 0 to make transparent
+  })
+
+  // Write the modified PNG data to a new file
+  const outputFilePath = `data/chatgpt/imagesAccept/${Date.now()}_masked.png`
+  await image.writeAsync(outputFilePath)
+  return outputFilePath
+}