From a42ef2944c7a847fb1f9df621e92293f342f8569 Mon Sep 17 00:00:00 2001
From: Daniel Avila <messagedaniel@protonmail.com>
Date: Wed, 14 Jun 2023 13:06:11 -0400
Subject: [PATCH] feat(StableDiffusion.js): add StableDiffusionAPI as a
 StructuredTool for Functions Agent

---
 .../tools/structured/StableDiffusion.js       | 88 +++++++++++++++++++
 1 file changed, 88 insertions(+)
 create mode 100644 api/app/langchain/tools/structured/StableDiffusion.js

diff --git a/api/app/langchain/tools/structured/StableDiffusion.js b/api/app/langchain/tools/structured/StableDiffusion.js
new file mode 100644
index 000000000..f2afbe0dc
--- /dev/null
+++ b/api/app/langchain/tools/structured/StableDiffusion.js
@@ -0,0 +1,88 @@
+// Generates image using stable diffusion webui's api (automatic1111)
+const fs = require('fs');
+const { StructuredTool } = require('langchain/tools');
+const { z } = require('zod');
+const path = require('path');
+const axios = require('axios');
+const sharp = require('sharp');
+
+class StableDiffusionAPI extends StructuredTool {
+  constructor(fields) {
+    super();
+    this.name = 'stable-diffusion';
+    this.url = fields.SD_WEBUI_URL || this.getServerURL();
+    this.description = `You can generate images with 'stable-diffusion'. This tool is exclusively for visual content.
+Guidelines:
+- Visually describe the moods, details, structures, styles, and/or proportions of the image. Remember, the focus is on visual attributes.
+- Craft your input by "showing" and not "telling" the imagery. Think in terms of what you'd want to see in a photograph or a painting.
+- Here's an example for generating a realistic portrait photo of a man:
+"prompt":"photo of a man in black clothes, half body, high detailed skin, coastline, overcast weather, wind, waves, 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3"
+"negative_prompt":"semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, out of frame, low quality, ugly, mutation, deformed"
+- Generate images only once per human query unless explicitly requested by the user`;
+    this.schema = z.object({
+      prompt: z.string().describe("Detailed keywords to describe the subject, separated by comma"),
+      negative_prompt: z.string().describe("Keywords we want to exclude from the final image")
+    });
+  }
+
+  replaceNewLinesWithSpaces(inputString) {
+    return inputString.replace(/\r\n|\r|\n/g, ' ');
+  }
+
+  getMarkdownImageUrl(imageName) {
+    const imageUrl = path.join(this.relativeImageUrl, imageName).replace(/\\/g, '/').replace('public/', '');
+    return `![generated image](/${imageUrl})`;
+  }
+
+  getServerURL() {
+    const url = process.env.SD_WEBUI_URL || '';
+    if (!url) {
+      throw new Error('Missing SD_WEBUI_URL environment variable.');
+    }
+    return url;
+  }
+
+  async _call(prompt, negative_prompt) {
+    const url = this.url;
+    const payload = {
+      prompt,
+      negative_prompt,
+      steps: 20
+    };
+    const response = await axios.post(`${url}/sdapi/v1/txt2img`, payload);
+    const image = response.data.images[0];
+    const pngPayload = { image: `data:image/png;base64,${image}` };
+    const response2 = await axios.post(`${url}/sdapi/v1/png-info`, pngPayload);
+    const info = response2.data.info;
+
+    // Generate unique name
+    const imageName = `${Date.now()}.png`;
+    this.outputPath = path.resolve(__dirname, '..', '..', '..', '..', 'client', 'public', 'images');
+    const appRoot = path.resolve(__dirname, '..', '..', '..', '..', 'client');
+    this.relativeImageUrl = path.relative(appRoot, this.outputPath);
+
+    // Check if directory exists, if not create it
+    if (!fs.existsSync(this.outputPath)) {
+      fs.mkdirSync(this.outputPath, { recursive: true });
+    }
+
+    try {
+      const buffer = Buffer.from(image.split(',', 1)[0], 'base64');
+      await sharp(buffer)
+        .withMetadata({
+          iptcpng: {
+            parameters: info
+          }
+        })
+        .toFile(this.outputPath + '/' + imageName);
+      this.result = this.getMarkdownImageUrl(imageName);
+    } catch (error) {
+      console.error('Error while saving the image:', error);
+      // this.result = theImageUrl;
+    }
+
+    return this.result;
+  }
+}
+
+module.exports = StableDiffusionAPI;