Add TogetherAI for image generation

public/scripts/extensions/stable-diffusion/index.js

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

public/scripts/extensions/stable-diffusion/settings.html

@@ -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>

src/endpoints/stable-diffusion.js

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

src/util.js

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