🎨 fix: Optimize StableDiffusion API Tool and Fix for Assistants Usage (#2253)

* chore: update docs

* fix(StableDiffusion): optimize API responses and file handling, return expected metadata for Assistants endpoint
This commit is contained in:
Danny Avila 2024-03-30 20:09:59 -04:00 committed by GitHub
parent 56ea0f9ae7
commit bb8a40dd98
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 113 additions and 52 deletions

View file

@ -12,14 +12,15 @@ const { logger } = require('~/config');
class DALLE3 extends Tool { class DALLE3 extends Tool {
constructor(fields = {}) { constructor(fields = {}) {
super(); super();
/* Used to initialize the Tool without necessary variables. */ /** @type {boolean} Used to initialize the Tool without necessary variables. */
this.override = fields.override ?? false; this.override = fields.override ?? false;
/* Necessary for output to contain all image metadata. */ /** @type {boolean} Necessary for output to contain all image metadata. */
this.returnMetadata = fields.returnMetadata ?? false; this.returnMetadata = fields.returnMetadata ?? false;
this.userId = fields.userId; this.userId = fields.userId;
this.fileStrategy = fields.fileStrategy; this.fileStrategy = fields.fileStrategy;
if (fields.processFileURL) { if (fields.processFileURL) {
/** @type {processFileURL} Necessary for output to contain all image metadata. */
this.processFileURL = fields.processFileURL.bind(this); this.processFileURL = fields.processFileURL.bind(this);
} }
@ -165,13 +166,7 @@ Error Message: ${error.message}`;
}); });
if (this.returnMetadata) { if (this.returnMetadata) {
this.result = { this.result = result;
file_id: result.file_id,
filename: result.filename,
filepath: result.filepath,
height: result.height,
width: result.width,
};
} else { } else {
this.result = this.wrapInMarkdown(result.filepath); this.result = this.wrapInMarkdown(result.filepath);
} }

View file

@ -4,14 +4,27 @@ const { z } = require('zod');
const path = require('path'); const path = require('path');
const axios = require('axios'); const axios = require('axios');
const sharp = require('sharp'); const sharp = require('sharp');
const { v4: uuidv4 } = require('uuid');
const { StructuredTool } = require('langchain/tools'); const { StructuredTool } = require('langchain/tools');
const { FileContext } = require('librechat-data-provider');
const paths = require('~/config/paths');
const { logger } = require('~/config'); const { logger } = require('~/config');
class StableDiffusionAPI extends StructuredTool { class StableDiffusionAPI extends StructuredTool {
constructor(fields) { constructor(fields) {
super(); super();
/* Used to initialize the Tool without necessary variables. */ /** @type {string} User ID */
this.userId = fields.userId;
/** @type {Express.Request | undefined} Express Request object, only provided by ToolService */
this.req = fields.req;
/** @type {boolean} Used to initialize the Tool without necessary variables. */
this.override = fields.override ?? false; this.override = fields.override ?? false;
/** @type {boolean} Necessary for output to contain all image metadata. */
this.returnMetadata = fields.returnMetadata ?? false;
if (fields.uploadImageBuffer) {
/** @type {uploadImageBuffer} Necessary for output to contain all image metadata. */
this.uploadImageBuffer = fields.uploadImageBuffer.bind(this);
}
this.name = 'stable-diffusion'; this.name = 'stable-diffusion';
this.url = fields.SD_WEBUI_URL || this.getServerURL(); this.url = fields.SD_WEBUI_URL || this.getServerURL();
@ -47,7 +60,7 @@ class StableDiffusionAPI extends StructuredTool {
getMarkdownImageUrl(imageName) { getMarkdownImageUrl(imageName) {
const imageUrl = path const imageUrl = path
.join(this.relativeImageUrl, imageName) .join(this.relativePath, this.userId, imageName)
.replace(/\\/g, '/') .replace(/\\/g, '/')
.replace('public/', ''); .replace('public/', '');
return `![generated image](/${imageUrl})`; return `![generated image](/${imageUrl})`;
@ -73,46 +86,67 @@ class StableDiffusionAPI extends StructuredTool {
width: 1024, width: 1024,
height: 1024, height: 1024,
}; };
const response = await axios.post(`${url}/sdapi/v1/txt2img`, payload); const generationResponse = await axios.post(`${url}/sdapi/v1/txt2img`, payload);
const image = response.data.images[0]; const image = generationResponse.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 /** @type {{ height: number, width: number, seed: number, infotexts: string[] }} */
const imageName = `${Date.now()}.png`; let info = {};
this.outputPath = path.resolve( try {
__dirname, info = JSON.parse(generationResponse.data.info);
'..', } catch (error) {
'..', logger.error('[StableDiffusion] Error while getting image metadata:', error);
'..', }
'..',
'..',
'client',
'public',
'images',
);
const appRoot = path.resolve(__dirname, '..', '..', '..', '..', '..', 'client');
this.relativeImageUrl = path.relative(appRoot, this.outputPath);
// Check if directory exists, if not create it const file_id = uuidv4();
if (!fs.existsSync(this.outputPath)) { const imageName = `${file_id}.png`;
fs.mkdirSync(this.outputPath, { recursive: true }); const { imageOutput: imageOutputPath, clientPath } = paths;
const filepath = path.join(imageOutputPath, this.userId, imageName);
this.relativePath = path.relative(clientPath, imageOutputPath);
if (!fs.existsSync(imageOutputPath)) {
fs.mkdirSync(imageOutputPath, { recursive: true });
} }
try { try {
const buffer = Buffer.from(image.split(',', 1)[0], 'base64'); const buffer = Buffer.from(image.split(',', 1)[0], 'base64');
if (this.returnMetadata && this.uploadImageBuffer && this.req) {
const file = await this.uploadImageBuffer({
req: this.req,
context: FileContext.image_generation,
resize: false,
metadata: {
buffer,
height: info.height,
width: info.width,
bytes: Buffer.byteLength(buffer),
filename: imageName,
type: 'image/png',
file_id,
},
});
const generationInfo = info.infotexts[0].split('\n').pop();
return {
...file,
prompt,
metadata: {
negative_prompt,
seed: info.seed,
info: generationInfo,
},
};
}
await sharp(buffer) await sharp(buffer)
.withMetadata({ .withMetadata({
iptcpng: { iptcpng: {
parameters: info, parameters: info.infotexts[0],
}, },
}) })
.toFile(this.outputPath + '/' + imageName); .toFile(filepath);
this.result = this.getMarkdownImageUrl(imageName); this.result = this.getMarkdownImageUrl(imageName);
} catch (error) { } catch (error) {
logger.error('[StableDiffusion] Error while saving the image:', error); logger.error('[StableDiffusion] Error while saving the image:', error);
// this.result = theImageUrl;
} }
return this.result; return this.result;

View file

@ -237,9 +237,11 @@ const loadTools = async ({
} }
const imageGenOptions = { const imageGenOptions = {
req: options.req,
fileStrategy: options.fileStrategy, fileStrategy: options.fileStrategy,
processFileURL: options.processFileURL, processFileURL: options.processFileURL,
returnMetadata: options.returnMetadata, returnMetadata: options.returnMetadata,
uploadImageBuffer: options.uploadImageBuffer,
}; };
const toolOptions = { const toolOptions = {

View file

@ -1,7 +1,9 @@
const path = require('path'); const path = require('path');
module.exports = { module.exports = {
root: path.resolve(__dirname, '..', '..'),
uploads: path.resolve(__dirname, '..', '..', 'uploads'), uploads: path.resolve(__dirname, '..', '..', 'uploads'),
clientPath: path.resolve(__dirname, '..', '..', 'client'),
dist: path.resolve(__dirname, '..', '..', 'client', 'dist'), dist: path.resolve(__dirname, '..', '..', 'client', 'dist'),
publicPath: path.resolve(__dirname, '..', '..', 'client', 'public'), publicPath: path.resolve(__dirname, '..', '..', 'client', 'public'),
imageOutput: path.resolve(__dirname, '..', '..', 'client', 'public', 'images'), imageOutput: path.resolve(__dirname, '..', '..', 'client', 'public', 'images'),

View file

@ -62,14 +62,14 @@ async function resizeImageBuffer(inputBuffer, resolution, endpoint) {
} }
/** /**
* Resizes an image buffer to webp format as well as reduces 150 px width. * Resizes an image buffer to webp format as well as reduces by specified or default 150 px width.
* *
* @param {Buffer} inputBuffer - The buffer of the image to be resized. * @param {Buffer} inputBuffer - The buffer of the image to be resized.
* @returns {Promise<{ buffer: Buffer, width: number, height: number, bytes: number }>} An object containing the resized image buffer, its size and dimensions. * @returns {Promise<{ buffer: Buffer, width: number, height: number, bytes: number }>} An object containing the resized image buffer, its size and dimensions.
* @throws Will throw an error if the resolution parameter is invalid. * @throws Will throw an error if the resolution parameter is invalid.
*/ */
async function resizeAndConvert(inputBuffer) { async function resizeAndConvert(inputBuffer, width = 150) {
const resizedBuffer = await sharp(inputBuffer).resize({ width: 150 }).toFormat('webp').toBuffer(); const resizedBuffer = await sharp(inputBuffer).resize({ width }).toFormat('webp').toBuffer();
const resizedMetadata = await sharp(resizedBuffer).metadata(); const resizedMetadata = await sharp(resizedBuffer).metadata();
return { return {
buffer: resizedBuffer, buffer: resizedBuffer,

View file

@ -223,26 +223,32 @@ const processImageFile = async ({ req, res, file, metadata }) => {
* @param {Object} params - The parameters object. * @param {Object} params - The parameters object.
* @param {Express.Request} params.req - The Express request object. * @param {Express.Request} params.req - The Express request object.
* @param {FileContext} params.context - The context of the file (e.g., 'avatar', 'image_generation', etc.) * @param {FileContext} params.context - The context of the file (e.g., 'avatar', 'image_generation', etc.)
* @param {boolean} [params.resize=true] - Whether to resize and convert the image to WebP. Default is `true`.
* @param {{ buffer: Buffer, width: number, height: number, bytes: number, filename: string, type: string, file_id: string }} [params.metadata] - Required metadata for the file if resize is false.
* @returns {Promise<{ filepath: string, filename: string, source: string, type: 'image/webp'}>} * @returns {Promise<{ filepath: string, filename: string, source: string, type: 'image/webp'}>}
*/ */
const uploadImageBuffer = async ({ req, context }) => { const uploadImageBuffer = async ({ req, context, metadata = {}, resize = true }) => {
const source = req.app.locals.fileStrategy; const source = req.app.locals.fileStrategy;
const { saveBuffer } = getStrategyFunctions(source); const { saveBuffer } = getStrategyFunctions(source);
const { buffer, width, height, bytes } = await resizeAndConvert(req.file.buffer); let { buffer, width, height, bytes, filename, file_id, type } = metadata;
const file_id = v4(); if (resize) {
const fileName = `img-${file_id}.webp`; file_id = v4();
type = 'image/webp';
({ buffer, width, height, bytes } = await resizeAndConvert(req.file.buffer));
filename = path.basename(req.file.originalname, path.extname(req.file.originalname)) + '.webp';
}
const filepath = await saveBuffer({ userId: req.user.id, fileName, buffer }); const filepath = await saveBuffer({ userId: req.user.id, fileName: filename, buffer });
return await createFile( return await createFile(
{ {
user: req.user.id, user: req.user.id,
file_id, file_id,
bytes, bytes,
filepath, filepath,
filename: req.file.originalname, filename,
context, context,
source, source,
type: 'image/webp', type,
width, width,
height, height,
}, },

View file

@ -12,8 +12,8 @@ const {
openapiToFunction, openapiToFunction,
validateAndParseOpenAPISpec, validateAndParseOpenAPISpec,
} = require('librechat-data-provider'); } = require('librechat-data-provider');
const { processFileURL, uploadImageBuffer } = require('~/server/services/Files/process');
const { loadActionSets, createActionTool, domainParser } = require('./ActionService'); const { loadActionSets, createActionTool, domainParser } = require('./ActionService');
const { processFileURL } = require('~/server/services/Files/process');
const { recordUsage } = require('~/server/services/Threads'); const { recordUsage } = require('~/server/services/Threads');
const { loadTools } = require('~/app/clients/tools/util'); const { loadTools } = require('~/app/clients/tools/util');
const { redactMessage } = require('~/config/parsers'); const { redactMessage } = require('~/config/parsers');
@ -147,7 +147,7 @@ const processVisionRequest = async (client, currentAction) => {
/** /**
* Processes return required actions from run. * Processes return required actions from run.
* @param {OpenAIClient} client - OpenAI or StreamRunManager Client. * @param {OpenAIClient | StreamRunManager} client - OpenAI (legacy) or StreamRunManager Client.
* @param {RequiredAction[]} requiredActions - The required actions to submit outputs for. * @param {RequiredAction[]} requiredActions - The required actions to submit outputs for.
* @returns {Promise<ToolOutputs>} The outputs of the tools. * @returns {Promise<ToolOutputs>} The outputs of the tools.
*/ */
@ -164,6 +164,8 @@ async function processRequiredActions(client, requiredActions) {
functions: true, functions: true,
options: { options: {
processFileURL, processFileURL,
req: client.req,
uploadImageBuffer,
openAIApiKey: client.apiKey, openAIApiKey: client.apiKey,
fileStrategy: client.req.app.locals.fileStrategy, fileStrategy: client.req.app.locals.fileStrategy,
returnMetadata: true, returnMetadata: true,

View file

@ -386,6 +386,18 @@
* @memberof typedefs * @memberof typedefs
*/ */
/**
* @exports uploadImageBuffer
* @typedef {import('~/server/services/Files/process').uploadImageBuffer} uploadImageBuffer
* @memberof typedefs
*/
/**
* @exports processFileURL
* @typedef {import('~/server/services/Files/process').processFileURL} processFileURL
* @memberof typedefs
*/
/** /**
* @exports AssistantCreateParams * @exports AssistantCreateParams
* @typedef {import('librechat-data-provider').AssistantCreateParams} AssistantCreateParams * @typedef {import('librechat-data-provider').AssistantCreateParams} AssistantCreateParams

View file

@ -51,7 +51,12 @@ Once you have completed all the setup, you can start the LibreChat application b
That's it! If you need more detailed information on configuring your compose file, see my notes below. That's it! If you need more detailed information on configuring your compose file, see my notes below.
## Updating LibreChat ## Updating LibreChat
The following commands will fetch the latest code of LibreChat and build a new docker image.
As of v0.7.0+, Docker installations transitioned from building images locally to using prebuilt images [hosted on Github Container registry](https://github.com/danny-avila?tab=packages&repo_name=LibreChat).
You can still build the image locally, as shown in the commented commands below. More info on building the image locally in the [Docker Compose Override Section](../configuration/docker_override.md).
The following commands will fetch the latest LibreChat project changes, including any necessary changes to the docker compose files, as well as the latest prebuilt images.
```bash ```bash
# Stop the running container(s) # Stop the running container(s)
@ -63,13 +68,16 @@ git pull
# Pull the latest LibreChat image (default setup) # Pull the latest LibreChat image (default setup)
docker compose pull docker compose pull
# If building the LibreChat image Locally, build without cache (legacy setup)
# docker compose build --no-cache
# Start LibreChat # Start LibreChat
docker compose up docker compose up
``` ```
If you're having issues running the above commands, you can try a comprehensive approach: If you're having issues running the above commands, you can try a comprehensive approach instead:
Prefix commands with `sudo` according to your environment permissions. Note: you may need to prefix commands with `sudo` according to your environment permissions.
```bash ```bash
# Stop the container (if running) # Stop the container (if running)