rixty
/
SillyTavern
mirror of https://github.com/SillyTavern/SillyTavern.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add TogetherAI for image generation

This commit is contained in:
Cohee 2023-12-18 03:33:05 +02:00
parent 4473532151
commit ac70a0a592
4 changed files with 184 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

78
public/scripts/extensions/stable-diffusion/index.js
View File

@ -46,6 +46,7 @@ const sources = {
vlad: 'vlad',
openai: 'openai',
comfy: 'comfy',
togetherai: 'togetherai',
};
const generationMode = {
@ -917,7 +918,7 @@ async function onModelChange() {
extension_settings.sd.model = $('#sd_model').find(':selected').val();
saveSettingsDebounced();
const cloudSources = [sources.horde, sources.novel, sources.openai];
const cloudSources = [sources.horde, sources.novel, sources.openai, sources.togetherai];
if (cloudSources.includes(extension_settings.sd.source)) {
return;
@ -1050,11 +1051,14 @@ async function loadSamplers() {
samplers = await loadVladSamplers();
break;
case sources.openai:
samplers = await loadOpenAiSamplers();
samplers = ['N/A'];
break;
case sources.comfy:
samplers = await loadComfySamplers();
break;
case sources.togetherai:
samplers = ['N/A'];
break;
}
for (const sampler of samplers) {
@ -1064,6 +1068,11 @@ async function loadSamplers() {
option.selected = sampler === extension_settings.sd.sampler;
$('#sd_sampler').append(option);
}
if (!extension_settings.sd.sampler && samplers.length > 0) {
extension_settings.sd.sampler = samplers[0];
$('#sd_sampler').val(extension_settings.sd.sampler).trigger('change');
}
}
async function loadHordeSamplers() {
@ -1120,10 +1129,6 @@ async function loadAutoSamplers() {
}
}
async function loadOpenAiSamplers() {
return ['N/A'];
}
async function loadVladSamplers() {
if (!extension_settings.sd.vlad_url) {
return [];
@ -1212,6 +1217,9 @@ async function loadModels() {
case sources.comfy:
models = await loadComfyModels();
break;
case sources.togetherai:
models = await loadTogetherAIModels();
break;
}
for (const model of models) {
@ -1221,6 +1229,30 @@ async function loadModels() {
option.selected = model.value === extension_settings.sd.model;
$('#sd_model').append(option);
}
if (!extension_settings.sd.model && models.length > 0) {
extension_settings.sd.model = models[0].value;
$('#sd_model').val(extension_settings.sd.model).trigger('change');
}
}
async function loadTogetherAIModels() {
if (!secret_state[SECRET_KEYS.TOGETHERAI]) {
console.debug('TogetherAI API key is not set.');
return [];
}
const result = await fetch('/api/sd/together/models', {
method: 'POST',
headers: getRequestHeaders(),
});
if (result.ok) {
const data = await result.json();
return data;
}
return [];
}
async function loadHordeModels() {
@ -1434,6 +1466,9 @@ async function loadSchedulers() {
case sources.openai:
schedulers = ['N/A'];
break;
case sources.togetherai:
schedulers = ['N/A'];
break;
case sources.comfy:
schedulers = await loadComfySchedulers();
break;
@ -1493,6 +1528,9 @@ async function loadVaes() {
case sources.openai:
vaes = ['N/A'];
break;
case sources.togetherai:
vaes = ['N/A'];
break;
case sources.comfy:
vaes = await loadComfyVaes();
break;
@ -1873,6 +1911,9 @@ async function sendGenerationRequest(generationType, prompt, characterName = nul
case sources.comfy:
result = await generateComfyImage(prefixedPrompt);
break;
case sources.togetherai:
result = await generateTogetherAIImage(prefixedPrompt);
break;
}
if (!result.data) {
@ -1895,6 +1936,29 @@ async function sendGenerationRequest(generationType, prompt, characterName = nul
callback ? callback(prompt, base64Image, generationType) : sendMessage(prompt, base64Image, generationType);
}
async function generateTogetherAIImage(prompt) {
const result = await fetch('/api/sd/together/generate', {
method: 'POST',
headers: getRequestHeaders(),
body: JSON.stringify({
prompt: prompt,
negative_prompt: extension_settings.sd.negative_prompt,
model: extension_settings.sd.model,
steps: extension_settings.sd.steps,
width: extension_settings.sd.width,
height: extension_settings.sd.height,
}),
});
if (result.ok) {
const data = await result.json();
return { format: 'jpg', data: data?.output?.choices?.[0]?.image_base64 };
} else {
const text = await result.text();
throw new Error(text);
}
}
/**
* Generates an "extras" image using a provided prompt and other settings.
*
@ -2435,6 +2499,8 @@ function isValidState() {
return secret_state[SECRET_KEYS.OPENAI];
case sources.comfy:
return true;
case sources.togetherai:
return secret_state[SECRET_KEYS.TOGETHERAI];
}
}

1
public/scripts/extensions/stable-diffusion/settings.html
View File

@ -35,6 +35,7 @@
<option value="novel">NovelAI Diffusion</option>
<option value="openai">OpenAI (DALL-E)</option>
<option value="comfy">ComfyUI</option>
<option value="togetherai">TogetherAI</option>
</select>
<div data-sd-source="auto">
<label for="sd_auto_url">SD Web UI URL</label>

96
src/endpoints/stable-diffusion.js
View File

@ -1,11 +1,12 @@
const express = require('express');
const fetch = require('node-fetch').default;
const sanitize = require('sanitize-filename');
const { getBasicAuthHeader, delay } = require('../util.js');
const { getBasicAuthHeader, delay, getHexString } = require('../util.js');
const fs = require('fs');
const { DIRECTORIES } = require('../constants.js');
const writeFileAtomicSync = require('write-file-atomic').sync;
const { jsonParser } = require('../express-common');
const { readSecret, SECRET_KEYS } = require('./secrets.js');
/**
* Sanitizes a string.
@ -545,6 +546,99 @@ comfy.post('/generate', jsonParser, async (request, response) => {
}
});
const together = express.Router();
together.post('/models', jsonParser, async (_, response) => {
try {
const key = readSecret(SECRET_KEYS.TOGETHERAI);
if (!key) {
console.log('TogetherAI key not found.');
return response.sendStatus(400);
}
const modelsResponse = await fetch('https://api.together.xyz/api/models', {
method: 'GET',
headers: {
'Authorization': `Bearer ${key}`,
},
});
if (!modelsResponse.ok) {
console.log('TogetherAI returned an error.');
return response.sendStatus(500);
}
const data = await modelsResponse.json();
if (!Array.isArray(data)) {
console.log('TogetherAI returned invalid data.');
return response.sendStatus(500);
}
const models = data
.filter(x => x.display_type === 'image')
.map(x => ({ value: x.name, text: x.display_name }));
return response.send(models);
} catch (error) {
console.log(error);
return response.sendStatus(500);
}
});
together.post('/generate', jsonParser, async (request, response) => {
try {
const key = readSecret(SECRET_KEYS.TOGETHERAI);
if (!key) {
console.log('TogetherAI key not found.');
return response.sendStatus(400);
}
console.log('TogetherAI request:', request.body);
const result = await fetch('https://api.together.xyz/api/inference', {
method: 'POST',
body: JSON.stringify({
request_type: 'image-model-inference',
prompt: request.body.prompt,
negative_prompt: request.body.negative_prompt,
height: request.body.height,
width: request.body.width,
model: request.body.model,
steps: request.body.steps,
n: 1,
seed: Math.floor(Math.random() * 10_000_000), // Limited to 10000 on playground, works fine with more.
sessionKey: getHexString(40), // Don't know if that's supposed to be random or not. It works either way.
}),
headers: {
'Content-Type': 'application/json',
'Authorization': `Bearer ${key}`,
},
});
if (!result.ok) {
console.log('TogetherAI returned an error.');
return response.sendStatus(500);
}
const data = await result.json();
console.log('TogetherAI response:', data);
if (data.status !== 'finished') {
console.log('TogetherAI job failed.');
return response.sendStatus(500);
}
return response.send(data);
} catch (error) {
console.log(error);
return response.sendStatus(500);
}
});
router.use('/comfy', comfy);
router.use('/together', together);
module.exports = { router };

16
src/util.js
View File

@ -105,6 +105,21 @@ function delay(ms) {
return new Promise(resolve => setTimeout(resolve, ms));
}
/**
* Generates a random hex string of the given length.
* @param {number} length String length
* @returns {string} Random hex string
* @example getHexString(8) // 'a1b2c3d4'
*/
function getHexString(length) {
const chars = '0123456789abcdef';
let result = '';
for (let i = 0; i < length; i++) {
result += chars[Math.floor(Math.random() * chars.length)];
}
return result;
}
/**
* Extracts a file with given extension from an ArrayBuffer containing a ZIP archive.
* @param {ArrayBuffer} archiveBuffer Buffer containing a ZIP archive
@ -404,4 +419,5 @@ module.exports = {
removeOldBackups,
getImages,
forwardFetchResponse,
getHexString,
};