Switch branch
feat/multi-feature-pack maindefault
View all branches
Blame Raw
Cohee · e3f41666 · · 408 lines (14.1 KB)
1 contributor 
1import fetch from 'node-fetch';


2import express from 'express';


3import { AIHorde, ModelGenerationInputStableSamplers, ModelInterrogationFormTypes, HordeAsyncRequestStates } from '@zeldafan0225/ai_horde';


4import { getVersion, delay, Cache } from '../util.js';


5import { readSecret, SECRET_KEYS } from './secrets.js';


6


7const ANONYMOUS_KEY = '0000000000';


8const HORDE_TEXT_MODEL_METADATA_URL = 'https://raw.githubusercontent.com/db0/AI-Horde-text-model-reference/main/db.json';


9const cache = new Cache(60 * 1000);


10export const router = express.Router();


11


12/**


13 * Returns the AIHorde client agent.


14 * @returns {Promise<string>} AIHorde client agent


15 */


16async function getClientAgent() {


17    const version = await getVersion();


18    return version?.agent || 'SillyTavern:UNKNOWN:Cohee#1207';


19}


20


21/**


22 * Returns the AIHorde client.


23 * @returns {Promise<AIHorde>} AIHorde client


24 */


25async function getHordeClient() {


26    return new AIHorde({


27        client_agent: await getClientAgent(),


28    });


29}


30


31/**


32 * Removes dirty no-no words from the prompt.


33 * Taken verbatim from KAI Lite's implementation (AGPLv3).


34 * https://github.com/LostRuins/lite.koboldai.net/blob/main/index.html#L7786C2-L7811C1


35 * @param {string} prompt Prompt to sanitize


36 * @returns {string} Sanitized prompt


37 */


38function sanitizeHordeImagePrompt(prompt) {


39    if (!prompt) {


40        return '';


41    }


42


43    //to avoid flagging from some image models, always swap these words


44    prompt = prompt.replace(/\b(girl)\b/gmi, 'woman');


45    prompt = prompt.replace(/\b(boy)\b/gmi, 'man');


46    prompt = prompt.replace(/\b(girls)\b/gmi, 'women');


47    prompt = prompt.replace(/\b(boys)\b/gmi, 'men');


48    //always remove these high risk words from prompt, as they add little value to image gen while increasing the risk the prompt gets flagged


49    prompt = prompt.replace(/\b(under.age|under.aged|underage|underaged|loli|pedo|pedophile|(\w+).year.old|(\w+).years.old|minor|prepubescent|minors|shota)\b/gmi, '');


50    //replace risky subject nouns with person


51    prompt = prompt.replace(/\b(youngster|infant|baby|toddler|child|teen|kid|kiddie|kiddo|teenager|student|preteen|pre.teen)\b/gmi, 'person');


52    //remove risky adjectives and related words


53    prompt = prompt.replace(/\b(young|younger|youthful|youth|small|smaller|smallest|girly|boyish|lil|tiny|teenaged|lit[tl]le|school.aged|school|highschool|kindergarten|teens|children|kids)\b/gmi, '');


54


55    return prompt;


56}


57


58router.post('/text-workers', async (request, response) => {


59    try {


60        const cachedWorkers = cache.get('workers');


61


62        if (cachedWorkers && !request.body.force) {


63            return response.send(cachedWorkers);


64        }


65


66        const agent = await getClientAgent();


67        const fetchResult = await fetch('https://aihorde.net/api/v2/workers?type=text', {


68            headers: {


69                'Client-Agent': agent,


70            },


71        });


72        const data = await fetchResult.json();


73        cache.set('workers', data);


74        return response.send(data);


75    } catch (error) {


76        console.error(error);


77        response.sendStatus(500);


78    }


79});


80


81async function getHordeTextModelMetadata() {


82    const response = await fetch(HORDE_TEXT_MODEL_METADATA_URL);


83    return await response.json();


84}


85


86async function mergeModelsAndMetadata(models, metadata) {


87    return models.map(model => {


88        const metadataModel = metadata[model.name];


89        if (!metadataModel) {


90            return { ...model, is_whitelisted: false };


91        }


92        return { ...model, ...metadataModel, is_whitelisted: true };


93    });


94}


95


96router.post('/text-models', async (request, response) => {


97    try {


98        const cachedModels = cache.get('models');


99        if (cachedModels && !request.body.force) {


100            return response.send(cachedModels);


101        }


102


103        const agent = await getClientAgent();


104        const fetchResult = await fetch('https://aihorde.net/api/v2/status/models?type=text', {


105            headers: {


106                'Client-Agent': agent,


107            },


108        });


109


110        let data = await fetchResult.json();


111


112        // attempt to fetch and merge models metadata


113        try {


114            const metadata = await getHordeTextModelMetadata();


115            data = await mergeModelsAndMetadata(data, metadata);


116        } catch (error) {


117            console.error('Failed to fetch metadata:', error);


118        }


119


120        cache.set('models', data);


121        return response.send(data);


122    } catch (error) {


123        console.error(error);


124        response.sendStatus(500);


125    }


126});


127


128router.post('/status', async (_, response) => {


129    try {


130        const agent = await getClientAgent();


131        const fetchResult = await fetch('https://aihorde.net/api/v2/status/heartbeat', {


132            headers: {


133                'Client-Agent': agent,


134            },


135        });


136


137        return response.send({ ok: fetchResult.ok });


138    } catch (error) {


139        console.error(error);


140        response.sendStatus(500);


141    }


142});


143


144router.post('/cancel-task', async (request, response) => {


145    try {


146        const taskId = request.body.taskId;


147        const agent = await getClientAgent();


148        const fetchResult = await fetch(`https://aihorde.net/api/v2/generate/text/status/${taskId}`, {


149            method: 'DELETE',


150            headers: {


151                'Client-Agent': agent,


152            },


153        });


154


155        const data = await fetchResult.json();


156        console.info(`Cancelled Horde task ${taskId}`);


157        return response.send(data);


158    } catch (error) {


159        console.error(error);


160        response.sendStatus(500);


161    }


162});


163


164router.post('/task-status', async (request, response) => {


165    try {


166        const taskId = request.body.taskId;


167        const agent = await getClientAgent();


168        const fetchResult = await fetch(`https://aihorde.net/api/v2/generate/text/status/${taskId}`, {


169            headers: {


170                'Client-Agent': agent,


171            },


172        });


173


174        const data = await fetchResult.json();


175        console.info(`Horde task ${taskId} status:`, data);


176        return response.send(data);


177    } catch (error) {


178        console.error(error);


179        response.sendStatus(500);


180    }


181});


182


183router.post('/generate-text', async (request, response) => {


184    const apiKey = readSecret(request.user.directories, SECRET_KEYS.HORDE) || ANONYMOUS_KEY;


185    const url = 'https://aihorde.net/api/v2/generate/text/async';


186    const agent = await getClientAgent();


187


188    console.debug(request.body);


189    try {


190        const result = await fetch(url, {


191            method: 'POST',


192            body: JSON.stringify(request.body),


193            headers: {


194                'Content-Type': 'application/json',


195                'apikey': apiKey,


196                'Client-Agent': agent,


197            },


198        });


199


200        if (!result.ok) {


201            const message = await result.text();


202            console.error('Horde returned an error:', message);


203            return response.send({ error: { message } });


204        }


205


206        const data = await result.json();


207        return response.send(data);


208    } catch (error) {


209        console.error(error);


210        return response.send({ error: true });


211    }


212});


213


214router.post('/sd-samplers', async (_, response) => {


215    try {


216        const samplers = Object.values(ModelGenerationInputStableSamplers);


217        response.send(samplers);


218    } catch (error) {


219        console.error(error);


220        response.sendStatus(500);


221    }


222});


223


224router.post('/sd-models', async (_, response) => {


225    try {


226        const ai_horde = await getHordeClient();


227        const models = await ai_horde.getModels();


228        response.send(models);


229    } catch (error) {


230        console.error(error);


231        response.sendStatus(500);


232    }


233});


234


235router.post('/caption-image', async (request, response) => {


236    try {


237        const api_key_horde = readSecret(request.user.directories, SECRET_KEYS.HORDE) || ANONYMOUS_KEY;


238        const ai_horde = await getHordeClient();


239        const result = await ai_horde.postAsyncInterrogate({


240            source_image: request.body.image,


241            forms: [{ name: ModelInterrogationFormTypes.caption }],


242        }, { token: api_key_horde });


243


244        if (!result.id) {


245            console.error('Image interrogation request is not satisfyable:', result.message || 'unknown error');


246            return response.sendStatus(400);


247        }


248


249        const MAX_ATTEMPTS = 200;


250        const CHECK_INTERVAL = 3000;


251


252        for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {


253            await delay(CHECK_INTERVAL);


254            const status = await ai_horde.getInterrogationStatus(result.id);


255            console.info(status);


256


257            if (status.state === HordeAsyncRequestStates.done) {


258                if (status.forms === undefined) {


259                    console.error('Image interrogation request failed: no forms found.');


260                    return response.sendStatus(500);


261                }


262


263                console.debug('Image interrogation result:', status);


264                const caption = status?.forms[0]?.result?.caption || '';


265


266                if (!caption) {


267                    console.error('Image interrogation request failed: no caption found.');


268                    return response.sendStatus(500);


269                }


270


271                return response.send({ caption });


272            }


273


274            if (status.state === HordeAsyncRequestStates.faulted || status.state === HordeAsyncRequestStates.cancelled) {


275                console.error('Image interrogation request is not successful.');


276                return response.sendStatus(503);


277            }


278        }


279    } catch (error) {


280        console.error(error);


281        response.sendStatus(500);


282    }


283});


284


285router.post('/user-info', async (request, response) => {


286    const api_key_horde = readSecret(request.user.directories, SECRET_KEYS.HORDE);


287


288    if (!api_key_horde) {


289        return response.send({ anonymous: true });


290    }


291


292    try {


293        const ai_horde = await getHordeClient();


294        const sharedKey = await (async () => {


295            try {


296                return await ai_horde.getSharedKey(api_key_horde);


297            } catch {


298                return null;


299            }


300        })();


301        const user = await ai_horde.findUser({ token: api_key_horde });


302        return response.send({ user, sharedKey, anonymous: false });


303    } catch (error) {


304        console.error(error);


305        return response.sendStatus(500);


306    }


307});


308


309router.post('/generate-image', async (request, response) => {


310    if (!request.body.prompt) {


311        return response.sendStatus(400);


312    }


313


314    const MAX_ATTEMPTS = 200;


315    const CHECK_INTERVAL = 3000;


316    const PROMPT_THRESHOLD = 5000;


317


318    try {


319        const maxLength = PROMPT_THRESHOLD - String(request.body.negative_prompt).length - 5;


320        if (String(request.body.prompt).length > maxLength) {


321            console.warn('Stable Horde prompt is too long, truncating...');


322            request.body.prompt = String(request.body.prompt).substring(0, maxLength);


323        }


324


325        // Sanitize prompt if requested


326        if (request.body.sanitize) {


327            const sanitized = sanitizeHordeImagePrompt(request.body.prompt);


328


329            if (request.body.prompt !== sanitized) {


330                console.info('Stable Horde prompt was sanitized.');


331            }


332


333            request.body.prompt = sanitized;


334        }


335


336        const api_key_horde = readSecret(request.user.directories, SECRET_KEYS.HORDE) || ANONYMOUS_KEY;


337        console.debug('Stable Horde request:', request.body);


338


339        const ai_horde = await getHordeClient();


340        // noinspection JSCheckFunctionSignatures -- see @ts-ignore - use_gfpgan


341        const generation = await ai_horde.postAsyncImageGenerate(


342            {


343                prompt: `${request.body.prompt} ### ${request.body.negative_prompt}`,


344                params:


345                {


346                    sampler_name: request.body.sampler,


347                    hires_fix: request.body.enable_hr,


348                    // @ts-ignore - use_gfpgan param is not in the type definition, need to update to new ai_horde @ https://github.com/ZeldaFan0225/ai_horde/blob/main/index.ts


349                    use_gfpgan: request.body.restore_faces,


350                    cfg_scale: request.body.scale,


351                    steps: request.body.steps,


352                    width: request.body.width,


353                    height: request.body.height,


354                    karras: Boolean(request.body.karras),


355                    clip_skip: request.body.clip_skip,


356                    seed: request.body.seed >= 0 ? String(request.body.seed) : undefined,


357                    n: 1,


358                },


359                r2: false,


360                nsfw: request.body.nfsw,


361                models: [request.body.model],


362            },


363            { token: api_key_horde });


364


365        if (!generation.id) {


366            console.warn('Image generation request is not satisfyable:', generation.message || 'unknown error');


367            return response.sendStatus(400);


368        }


369


370        console.info('Horde image generation request:', generation);


371


372        const controller = new AbortController();


373        request.socket.removeAllListeners('close');


374        request.socket.on('close', function () {


375            console.warn('Horde image generation request aborted.');


376            controller.abort();


377            if (generation.id) ai_horde.deleteImageGenerationRequest(generation.id);


378        });


379


380        for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {


381            controller.signal.throwIfAborted();


382            await delay(CHECK_INTERVAL);


383            const check = await ai_horde.getImageGenerationCheck(generation.id);


384            console.info(check);


385


386            if (check.done) {


387                const result = await ai_horde.getImageGenerationStatus(generation.id);


388                if (result.generations === undefined) return response.sendStatus(500);


389                return response.send(result.generations[0].img);


390            }


391


392            /*


393            if (!check.is_possible) {


394                return response.sendStatus(503);


395            }


396            */


397


398            if (check.faulted) {


399                return response.sendStatus(500);


400            }


401        }


402


403        return response.sendStatus(504);


404    } catch (error) {


405        console.error(error);


406        return response.sendStatus(500);


407    }


408});