From a6a3eaaf0574b62b2eb33f7691ba92342f5cd401 Mon Sep 17 00:00:00 2001
From: derpystuff <3515180-derpystuff@users.noreply.gitlab.com>
Date: Sun, 21 Apr 2024 16:20:43 +0200
Subject: [PATCH] update wallpaper

---
 commands/message/genai/imagen.js    |  8 +--
 commands/message/genai/wallpaper.js | 75 +++++++++++++++++++++++++++++
 labscore/api/obelisk/endpoints.js   |  3 +-
 labscore/api/obelisk/index.js       | 11 ++++-
 4 files changed, 90 insertions(+), 7 deletions(-)
 create mode 100644 commands/message/genai/wallpaper.js

diff --git a/commands/message/genai/imagen.js b/commands/message/genai/imagen.js
index 04c1254..e713ed0 100644
--- a/commands/message/genai/imagen.js
+++ b/commands/message/genai/imagen.js
@@ -1,4 +1,4 @@
-const { geminiVision, aiWallpaper } = require("../../../labscore/api/obelisk");
+const { geminiVision, aiWallpaper, imagen } = require("../../../labscore/api/obelisk");
 const { getRecentImage } = require("../../../labscore/utils/attachment");
 const { createEmbed } = require("../../../labscore/utils/embed");
 const { editOrReply } = require("../../../labscore/utils/message");
@@ -16,14 +16,14 @@ module.exports = {
   aliases: ['aiimg'],
   metadata: {
     description: `${iconPill("generative_ai", "LIMITED TESTING")}\n\nGenerate images with Imagen 2`,
-    description_short: 'Create Wallpapers',
+    description_short: 'Create Images with Imagen',
     examples: ['imagen a painting of northern lights'],
     category: 'limited',
     usage: 'imagen <prompt>'
   },
   permissionsClient: [Permissions.EMBED_LINKS, Permissions.SEND_MESSAGES, Permissions.USE_EXTERNAL_EMOJIS, Permissions.ATTACH_FILES, Permissions.READ_MESSAGE_HISTORY],
   run: async (context, args) => {
-    if(!await hasFeature(context, "ai/wallpapers")) return;
+    if(!await hasFeature(context, "ai/imagen")) return;
     context.triggerTyping();
 
     if(!args.text) return editOrReply(context, createEmbed("warning", context, `Missing Parameter (prompt).`))
@@ -36,7 +36,7 @@ module.exports = {
         }
       }))
 
-      let res = await aiWallpaper(context, args.text);
+      let res = await imagen(context, args.text);
 
       // Construct Embeds
       let files = [];
diff --git a/commands/message/genai/wallpaper.js b/commands/message/genai/wallpaper.js
new file mode 100644
index 0000000..7a34a7a
--- /dev/null
+++ b/commands/message/genai/wallpaper.js
@@ -0,0 +1,75 @@
+const { geminiVision, aiWallpaper, imagen, wallpaper } = require("../../../labscore/api/obelisk");
+const { getRecentImage } = require("../../../labscore/utils/attachment");
+const { createEmbed } = require("../../../labscore/utils/embed");
+const { editOrReply } = require("../../../labscore/utils/message");
+const { getUser } = require("../../../labscore/utils/users");
+
+const { Permissions } = require("detritus-client/lib/constants");
+
+const { STATIC_ICONS } = require("../../../labscore/utils/statics");
+const { iconPill, stringwrap } = require("../../../labscore/utils/markdown");
+const { hasFeature } = require("../../../labscore/utils/testing");
+
+module.exports = {
+  name: 'wallpaper',
+  label: 'text',
+  aliases: ['aiwp'],
+  metadata: {
+    description: `${iconPill("generative_ai", "LIMITED TESTING")}\n\nGenerate AI Wallpapers`,
+    description_short: 'Create Wallpapers',
+    examples: ['wallpaper a painting of northern lights, in the bauhaus style -format square'],
+    category: 'limited',
+    usage: 'wallpaper <prompt> [-format <square|wide>]'
+  },
+  args: [
+    { name: 'format', default: 'wide', required: false, help: "Image style (wide, square)." }
+  ],
+  permissionsClient: [Permissions.EMBED_LINKS, Permissions.SEND_MESSAGES, Permissions.USE_EXTERNAL_EMOJIS, Permissions.ATTACH_FILES, Permissions.READ_MESSAGE_HISTORY],
+  run: async (context, args) => {
+    if(!await hasFeature(context, "ai/wallpapers")) return;
+    context.triggerTyping();
+
+    if(!args.text) return editOrReply(context, createEmbed("warning", context, `Missing Parameter (prompt).`))
+    if(!["wide","square"].includes(args.format.toLowerCase())) return editOrReply(context, createEmbed("warning", context, `Invalid Parameter (format).`))
+    try{
+      await editOrReply(context, createEmbed("defaultNoFooter", context, {
+        author: {
+          iconUrl: STATIC_ICONS.ai_image_processing,
+          name: "Generating images..."
+        }
+      }))
+
+      let res = await wallpaper(context, args.text, args.format.toLowerCase());
+
+      // Construct Embeds
+      let files = [];
+      let embeds = res.response.body.images.map((i)=>{
+        let imgName = `lcigen.${(Date.now() + Math.random()).toString(36)}.jpeg`;
+
+        files.push({
+          filename: imgName,
+          value: Buffer.from(i, 'base64')
+        })
+        return createEmbed("defaultNoFooter", context, {
+            url: "https://bignutty.gitlab.io",
+            author: {
+              iconUrl: STATIC_ICONS.ai_image,
+              name: stringwrap(args.text, 50, false),
+            },
+            image: {
+              url: `attachment://${imgName}`
+            },
+            footer: {
+              text: `Generative AI is experimental • Use AI Images responsibly.`
+            }
+          })
+        });
+      
+      return editOrReply(context, {embeds, files});
+    } catch(e){
+      console.log(e)
+      if(e.response?.body?.message) return editOrReply(context, createEmbed("error", context, e.response.body.message))
+      return editOrReply(context, createEmbed("error", context, `Unable to generate image.`))
+    }
+  }
+};
diff --git a/labscore/api/obelisk/endpoints.js b/labscore/api/obelisk/endpoints.js
index b9d4d4a..10534d2 100644
--- a/labscore/api/obelisk/endpoints.js
+++ b/labscore/api/obelisk/endpoints.js
@@ -17,7 +17,8 @@ const ObeliskApi = Object.freeze({
   WEB_ASK: "/flamingo/v1/web:ask",
   SUMMARIZE_WEBPAGES: "/flamingo/v1/web:summarize",
   
-  AI_WALLPAPER: "/robin/v1/wallpaper:generate",
+  GENERATE_IMAGEN: "/robin/v1/generate:imagen",
+  GENERATE_WALLPAPER: "/robin/v1/generate:wallpaper",
   
   WEBSHOT: "/peacock/v1/screenshot",
   TRANSCRIBE: "/peacock/v1/transcribe",
diff --git a/labscore/api/obelisk/index.js b/labscore/api/obelisk/index.js
index b481a9c..b2416fa 100644
--- a/labscore/api/obelisk/index.js
+++ b/labscore/api/obelisk/index.js
@@ -96,12 +96,19 @@ module.exports.summarizeWebpage = async function(context, url){
 }
 
 // ROBIN
-module.exports.aiWallpaper = async function(context, prompt){
-  return await request(ObeliskApi.AI_WALLPAPER, "POST", {}, {
+module.exports.imagen = async function(context, prompt){
+  return await request(ObeliskApi.GENERATE_IMAGEN, "POST", {}, {
     prompt
   })
 }
 
+module.exports.wallpaper = async function(context, prompt, model){
+  return await request(ObeliskApi.GENERATE_WALLPAPER, "POST", {}, {
+    prompt,
+    model
+  })
+}
+
 // PEACOCK
 module.exports.webshot = async function(context, url, allowAdultContent = false){
   return await request(ObeliskApi.WEBSHOT, "POST", {}, {