From a372a08b6efb2ebba0f0713062c9c4f6a72cc096 Mon Sep 17 00:00:00 2001
From: derpystuff <3515180-derpystuff@users.noreply.gitlab.com>
Date: Sat, 10 Dec 2022 23:07:33 +0100
Subject: [PATCH] support img2img

---
 commands/message/image/anime.js | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/commands/message/image/anime.js b/commands/message/image/anime.js
index 2c98e9c..37653d7 100644
--- a/commands/message/image/anime.js
+++ b/commands/message/image/anime.js
@@ -2,12 +2,13 @@ const { createEmbed } = require("../../../labscore/utils/embed");
 const { editOrReply } = require("../../../labscore/utils/message");
 
 const superagent = require('superagent');
+const { getRecentImage } = require("../../../labscore/utils/attachment");
 
 module.exports = {
   name: 'anime',
   label: 'query',
   metadata: {
-    description: 'Uses Stable Diffusion to generate an image using anime styling from a text prompt.',
+    description: 'Uses Stable Diffusion to generate an image using anime styling from a text prompt.\n\n**Replying** to a message with an image will use said image as a base for the AI (img2img), this does __not__ support inpainting/adding elements to the image.',
     description_short: 'AI Anime image generation',
     explicit: true,
     examples: ['anime otter'],
@@ -22,6 +23,15 @@ module.exports = {
   run: async (context, args) => {
     if(!context.channel.nsfw) return editOrReply(context, {embeds:[createEmbed("nsfw", context)]})
     if(args.query.length == 0) return editOrReply(context, {embeds:[createEmbed("warning", context, "Missing prompt")]})
+
+    let hasUrl = false;
+    let image;
+    if(context.message.messageReference){
+      image = await getRecentImage(context, 50)
+      hasUrl = true;
+      if(!image) return editOrReply(context, createEmbed("warning", context, "No images found."))
+    }
+
     let response = await editOrReply(context, { embeds: [createEmbed("loading", context, `Generating image...`)] })
 
     let noticeTimer = setTimeout(()=>{
@@ -35,11 +45,19 @@ module.exports = {
     try{
       let t = Date.now();
 
-      let img = await superagent.get(`${process.env.AI_SERVER}/anime`)
-        .query({
-          prompt: args.query
-        })
-
+      let img;
+      if(hasUrl){
+        img = await superagent.get(`${process.env.AI_SERVER}/anime/img2img`)
+          .query({
+            prompt: args.query,
+            url: image
+          })
+      } else {
+        img = await superagent.get(`${process.env.AI_SERVER}/anime`)
+          .query({
+            prompt: args.query
+          })
+      }
       clearTimeout(noticeTimer)
       
       if(img.body.message) return await response.edit({embeds:[createEmbed("warning", context, img.body.message)]})