support stable diffusion plugin (#49)

app/api/langchain-tools/stable_diffusion_image_generator.ts

@@ -0,0 +1,53 @@
+import { Tool } from "langchain/tools";
+import S3FileStorage from "../../utils/r2_file_storage";
+
+export class StableDiffusionWrapper extends Tool {
+  name = "stable_diffusion_image_generator";
+
+  constructor() {
+    super();
+  }
+
+  /** @ignore */
+  async _call(prompt: string) {
+    let url = process.env.STABLE_DIFFUSION_API_URL;
+    const data = {
+      prompt: prompt,
+      negative_prompt:
+        process.env.STABLE_DIFFUSION_NEGATIVE_PROMPT ??
+        "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
+      seed: -1,
+      subseed: -1,
+      subseed_strength: 0,
+      batch_size: 1,
+      n_iter: 1,
+      steps: process.env.STABLE_DIFFUSION_STEPS ?? 20,
+      cfg_scale: process.env.STABLE_DIFFUSION_CFG_SCALE ?? 7,
+      width: process.env.STABLE_DIFFUSION_WIDTH ?? 720,
+      height: process.env.STABLE_DIFFUSION_HEIGHT ?? 720,
+      restore_faces: process.env.STABLE_DIFFUSION_RESTORE_FACES ?? false,
+      eta: 0,
+      sampler_index: process.env.STABLE_DIFFUSION_SAMPLER_INDEX ?? "Euler a",
+    };
+    console.log(`[${this.name}]`, data);
+    const response = await fetch(`${url}/sdapi/v1/txt2img`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify(data),
+    });
+    const json = await response.json();
+    let imageBase64 = json.images[0];
+    if (!imageBase64) return "No image was generated";
+    const buffer = Buffer.from(imageBase64, "base64");
+    const filePath = await S3FileStorage.put(`${Date.now()}.png`, buffer);
+    console.log(`[${this.name}]`, filePath);
+    return filePath;
+  }
+
+  description = `stable diffusion is an ai art generation model similar to dalle-2.
+    input requires english.
+    output will be the image link url.
+    use markdown to display images. like: ![img](/api/file/xxx.png)`;
+}

app/api/langchain/tool/agent/route.ts

@@ -21,6 +21,7 @@ import { DynamicTool, Tool } from "langchain/tools";
 import { DallEAPIWrapper } from "@/app/api/langchain-tools/dalle_image_generator";
 import { BaiduSearch } from "@/app/api/langchain-tools/baidu_search";
 import { GoogleSearch } from "@/app/api/langchain-tools/google_search";
+import { StableDiffusionWrapper } from "@/app/api/langchain-tools/stable_diffusion_image_generator";
 
 const serverConfig = getServerSideConfig();
 
@@ -228,10 +229,13 @@ async function handle(req: NextRequest) {
     const webBrowserTool = new WebBrowser({ model, embeddings });
     const calculatorTool = new Calculator();
     const dallEAPITool = new DallEAPIWrapper(apiKey, baseUrl);
+    const stableDiffusionTool = new StableDiffusionWrapper();
     if (useTools.includes("web-search")) tools.push(searchTool);
     if (useTools.includes(webBrowserTool.name)) tools.push(webBrowserTool);
     if (useTools.includes(calculatorTool.name)) tools.push(calculatorTool);
     if (useTools.includes(dallEAPITool.name)) tools.push(dallEAPITool);
+    if (useTools.includes(stableDiffusionTool.name))
+      tools.push(stableDiffusionTool);
 
     useTools.forEach((toolName) => {
       if (toolName) {

app/masks/cn.ts

@@ -1,6 +1,35 @@
 import { BuiltinMask } from "./typing";
 
 export const CN_MASKS: BuiltinMask[] = [
+  {
+    avatar: "1f3a8",
+    name: "Stable Diffusion",
+    context: [
+      {
+        id: "SVx3ybvohJAKXDQ1KKQcs",
+        date: "",
+        role: "system",
+        content:
+          "Stable Diffusion is an AI art generation model similar to DALLE-2.\nHere are some prompts for generating art with Stable Diffusion.\n\nPrompt Example:\n\n- A ghostly apparition drifting through a haunted mansion's grand ballroom, illuminated by flickering candlelight. Eerie, ethereal, moody lighting.\n- portait of a homer simpson archer shooting arrow at forest monster, front game card, drark, marvel comics, dark, smooth\n- pirate, deep focus, fantasy, matte, sharp focus\n- red dead redemption 2, cinematic view, epic sky, detailed, low angle, high detail, warm lighting, volumetric, godrays, vivid, beautiful\n- a fantasy style portrait painting of rachel lane / alison brie hybrid in the style of francois boucher oil painting, rpg portrait\n- athena, greek goddess, claudia black, bronze greek armor, owl crown, d & d, fantasy, portrait, headshot, sharp focus\n- closeup portrait shot of a large strong female biomechanic woman in a scenic scifi environment, elegant, smooth, sharp focus, warframe\n- ultra realistic illustration of steve urkle as the hulk, elegant, smooth, sharp focus\n- portrait of beautiful happy young ana de armas, ethereal, realistic anime, clean lines, sharp lines, crisp lines, vibrant color scheme\n- A highly detailed and hyper realistic portrait of a gorgeous young ana de armas, lisa frank, butterflies, floral, sharp focus\n- lots of delicious tropical fruits with drops of moisture on table, floating colorful water, mysterious expression, in a modern and abstract setting, with bold and colorful abstract art, blurred background, bright lighting\n- 1girl, The most beautiful form of chaos, Fauvist design, Flowing colors, Vivid colors, dynamic angle, fantasy world\n- solo, sitting, close-up, girl in the hourglass, Sand is spilling out of the broken hourglass, flowing sand, huge hourglass art, hologram, particles, nebula, magic circle\n- geometric abstract background, 1girl, depth of field, zentangle, mandala, tangle, entangle, beautiful and aesthetic, dynamic angle, glowing skin, floating colorful sparkles the most beautiful form of chaos, elegant, a brutalist designed, vivid colours, romanticism\n\nFollow the structure of the example prompts. This means a very short description of the scene, followed by modifiers divided by commas to alter the mood, style, lighting, and more.\nIf the user input is in English, directly use the user input as a parameter to call the stable_diffusion_image_generator plugin. If the user input is not in English, generate an English prompt word based on the example and then call the stable_diffusion_image_generator plugin.",
+      },
+    ],
+    modelConfig: {
+      model: "gpt-3.5-turbo",
+      temperature: 1,
+      top_p: 1,
+      max_tokens: 2000,
+      presence_penalty: 0,
+      frequency_penalty: 0,
+      sendMemory: false,
+      historyMessageCount: 0,
+      compressMessageLengthThreshold: 1000,
+    },
+    lang: "cn",
+    builtin: false,
+    createdAt: 1697205441045,
+    usePlugins: true,
+    hideContext: true,
+  },
   {
     avatar: "1f5bc-fe0f",
     name: "以文搜图",

app/masks/en.ts

@@ -1,6 +1,35 @@
 import { BuiltinMask } from "./typing";
 
 export const EN_MASKS: BuiltinMask[] = [
+  {
+    avatar: "1f3a8",
+    name: "Stable Diffusion",
+    context: [
+      {
+        id: "SVx3ybvohJAKXDQ1KKQcs",
+        date: "",
+        role: "system",
+        content:
+          "Stable Diffusion is an AI art generation model similar to DALLE-2.\nHere are some prompts for generating art with Stable Diffusion.\n\nPrompt Example:\n\n- A ghostly apparition drifting through a haunted mansion's grand ballroom, illuminated by flickering candlelight. Eerie, ethereal, moody lighting.\n- portait of a homer simpson archer shooting arrow at forest monster, front game card, drark, marvel comics, dark, smooth\n- pirate, deep focus, fantasy, matte, sharp focus\n- red dead redemption 2, cinematic view, epic sky, detailed, low angle, high detail, warm lighting, volumetric, godrays, vivid, beautiful\n- a fantasy style portrait painting of rachel lane / alison brie hybrid in the style of francois boucher oil painting, rpg portrait\n- athena, greek goddess, claudia black, bronze greek armor, owl crown, d & d, fantasy, portrait, headshot, sharp focus\n- closeup portrait shot of a large strong female biomechanic woman in a scenic scifi environment, elegant, smooth, sharp focus, warframe\n- ultra realistic illustration of steve urkle as the hulk, elegant, smooth, sharp focus\n- portrait of beautiful happy young ana de armas, ethereal, realistic anime, clean lines, sharp lines, crisp lines, vibrant color scheme\n- A highly detailed and hyper realistic portrait of a gorgeous young ana de armas, lisa frank, butterflies, floral, sharp focus\n- lots of delicious tropical fruits with drops of moisture on table, floating colorful water, mysterious expression, in a modern and abstract setting, with bold and colorful abstract art, blurred background, bright lighting\n- 1girl, The most beautiful form of chaos, Fauvist design, Flowing colors, Vivid colors, dynamic angle, fantasy world\n- solo, sitting, close-up, girl in the hourglass, Sand is spilling out of the broken hourglass, flowing sand, huge hourglass art, hologram, particles, nebula, magic circle\n- geometric abstract background, 1girl, depth of field, zentangle, mandala, tangle, entangle, beautiful and aesthetic, dynamic angle, glowing skin, floating colorful sparkles the most beautiful form of chaos, elegant, a brutalist designed, vivid colours, romanticism\n\nFollow the structure of the example prompts. This means a very short description of the scene, followed by modifiers divided by commas to alter the mood, style, lighting, and more.\nIf the user input is in English, directly use the user input as a parameter to call the stable_diffusion_image_generator plugin. If the user input is not in English, generate an English prompt word based on the example and then call the stable_diffusion_image_generator plugin.",
+      },
+    ],
+    modelConfig: {
+      model: "gpt-3.5-turbo",
+      temperature: 1,
+      top_p: 1,
+      max_tokens: 2000,
+      presence_penalty: 0,
+      frequency_penalty: 0,
+      sendMemory: false,
+      historyMessageCount: 0,
+      compressMessageLengthThreshold: 1000,
+    },
+    lang: "en",
+    builtin: false,
+    createdAt: 1697205441045,
+    usePlugins: true,
+    hideContext: true,
+  },
   {
     avatar: "1f47e",
     name: "GitHub Copilot",

app/plugins/cn.ts

@@ -48,4 +48,14 @@ export const CN_PLUGINS: BuiltinPlugin[] = [
     createdAt: 1694703673000,
     enable: false,
   },
+  {
+    name: "Stable Diffusion",
+    toolName: "stable_diffusion_image_generator",
+    lang: "cn",
+    description:
+      "Stable Diffusion 图像生成模型。使用本插件需要配置 Cloudflare R2 对象存储服务以及 stable-diffusion-webui 接口。",
+    builtin: true,
+    createdAt: 1688899480510,
+    enable: false,
+  },
 ];

app/plugins/en.ts

@@ -50,4 +50,14 @@ export const EN_PLUGINS: BuiltinPlugin[] = [
     createdAt: 1694703673000,
     enable: false,
   },
+  {
+    name: "Stable Diffusion",
+    toolName: "stable_diffusion_image_generator",
+    lang: "en",
+    description:
+      "Stable Diffusion text-to-image model. Using this plugin requires configuring Cloudflare R2 object storage service and stable-diffusion-webui API.",
+    builtin: true,
+    createdAt: 1688899480510,
+    enable: false,
+  },
 ];

app/utils/r2_file_storage.ts

@@ -51,11 +51,16 @@ export default class S3FileStorage {
 
     console.log(signedUrl);
 
-    await fetch(signedUrl, {
-      method: "PUT",
-      body: data,
-    });
+    try {
+      await fetch(signedUrl, {
+        method: "PUT",
+        body: data,
+      });
 
-    return `/api/file/${fileName}`;
+      return `/api/file/${fileName}`;
+    } catch (e) {
+      console.error("[R2]", e);
+      throw e;
+    }
   }
 }