Move Stability generation to backend

This commit is contained in:
Cohee
2024-07-04 22:36:17 +03:00
parent 6608e530c5
commit e32b0cc223
5 changed files with 198 additions and 161 deletions

View File

@ -7,6 +7,7 @@ const path = require('path');
const writeFileAtomicSync = require('write-file-atomic').sync;
const { jsonParser } = require('../express-common');
const { readSecret, SECRET_KEYS } = require('./secrets.js');
const FormData = require('form-data');
/**
* Sanitizes a string.
@ -793,9 +794,68 @@ pollinations.post('/generate', jsonParser, async (request, response) => {
}
});
const stability = express.Router();
stability.post('/generate', jsonParser, async (request, response) => {
try {
const key = readSecret(request.user.directories, SECRET_KEYS.STABILITY);
if (!key) {
console.log('Stability AI key not found.');
return response.sendStatus(400);
}
const { payload, model } = request.body;
const formData = new FormData();
for (const [key, value] of Object.entries(payload)) {
if (value !== undefined) {
formData.append(key, String(value));
}
}
let apiUrl;
switch (model) {
case 'stable-image-ultra':
apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/ultra';
break;
case 'stable-image-core':
apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/core';
break;
case 'stable-diffusion-3':
apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/sd3';
break;
default:
throw new Error('Invalid Stability AI model selected');
}
const result = await fetch(apiUrl, {
method: 'POST',
headers: {
'Authorization': `Bearer ${key}`,
'Accept': 'image/*',
},
body: formData,
});
if (!result.ok) {
const text = await result.text();
console.log('Stability AI returned an error.', result.status, result.statusText, text);
return response.sendStatus(500);
}
const buffer = await result.buffer();
return response.send(buffer);
} catch (error) {
console.log(error);
return response.sendStatus(500);
}
});
router.use('/comfy', comfy);
router.use('/together', together);
router.use('/drawthings', drawthings);
router.use('/pollinations', pollinations);
router.use('/stability', stability);
module.exports = { router };