LibreChat Stable Diffusion RunPod

StableDiffusion.js

// Generates image using stable diffusion webui's api (automatic1111)
const fs = require('fs');
const { Tool } = require('langchain/tools');
const path = require('path');
const axios = require('axios');
const sharp = require('sharp');
// const { run } = require('jest');

class StableDiffusionAPI extends Tool {
  constructor(fields) {
    super();
    this.name = 'stable-diffusion';
    this.url = fields.SD_WEBUI_URL || this.getServerURL();
    this.runPodApiKey = fields.RUNPOD_API_KEY || this.getRunPodApiKey();
    this.headers = this.getHeader(this.runPodApiKey);
    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.
- It's best to follow this format for image creation:
"detailed keywords to describe the subject, separated by comma | keywords we want to exclude from the final image"
- Here's an example prompt for generating a realistic portrait photo of a man:
"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 | 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`;
  }

  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;
  }

  getRunPodApiKey() {
    const key = process.env.RUNPOD_API_KEY || '';
    if (!key) {
      throw new Error('Missing RUNPOD_API_KEY environment variable.');
    }
    return key;
  }

  getTxt2ImgUrl(url) {
    return `${url}`;
  }

  getTxt2ImgPayHeaders() {
    return this.headers;
  }

  getTxt2ImgPayload(input) {
    const payload = {
      input: {
        api_name: 'txt2img',
        prompt: input.split('|')[0],
        restore_faces: true,
        negative_prompt: input.split('|')[1],
        num_outputs: 1,
        override_settings: {
          sd_model_checkpoint: 'model.safetensors [62adc0f3ce]'
        },
        seed: 3302206224,
        cfg_scale: 5,
        sampler_index: 'DDIM',
        num_inference_steps: 20
      }
    };
    return payload;
  }

  getPngInfoUrl(url) {
    return `${url}`;
  }

  getPngInfoHeaders() {
    return this.headers;
  }
  
  getPngInfoPayload(image) {
    return {
      input: {
        api_name: 'png-info',
        image: `data:image/png;base64,${image}`
      }
    };
  }

  getHeader(runPodApiKey) {
    const headers = {
      'Content-Type': 'application/json',
      Authorization: 'Bearer ' + runPodApiKey
    };
    return headers;
  }

  async _call(input) {
    const url = this.url;
    const payload = this.getTxt2ImgPayload(input);
    // add header Bearer token to axios.post
    const response = await axios.post(this.getTxt2ImgUrl(url), payload, { headers: this.getTxt2ImgPayHeaders() });
    console.log('response', response.data);
    const image = response.data.output.images[0];

    const pngPayload = this.getPngInfoPayload(image);
    const response2 = await axios.post(this.getPngInfoUrl(url), pngPayload, {
      headers: this.getPngInfoHeaders()
    });
    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;