ImageGen: add BFL API for image generation

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

@@ -39,6 +39,7 @@ const UPDATE_INTERVAL = 1000;
 // This is a 1x1 transparent PNG
 const PNG_PIXEL = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=';
 const CUSTOM_STOP_EVENT = 'sd_stop_generation';
+const clamp = (value, min, max) => Math.min(Math.max(value, min), max);
 
 const sources = {
     extras: 'extras',
@@ -55,6 +56,7 @@ const sources = {
     blockentropy: 'blockentropy',
     huggingface: 'huggingface',
     nanogpt: 'nanogpt',
+    bfl: 'bfl',
 };
 
 const initiators = {
@@ -296,6 +298,9 @@ const defaultSettings = {
 
     // Stability AI settings
     stability_style_preset: 'anime',
+
+    // BFL API settings
+    bfl_upsampling: false,
 };
 
 const writePromptFieldsDebounced = debounce(writePromptFields, debounce_timeout.relaxed);
@@ -463,6 +468,7 @@ async function loadSettings() {
     $('#sd_stability_style_preset').val(extension_settings.sd.stability_style_preset);
     $('#sd_huggingface_model_id').val(extension_settings.sd.huggingface_model_id);
     $('#sd_function_tool').prop('checked', extension_settings.sd.function_tool);
+    $('#sd_bfl_upsampling').prop('checked', extension_settings.sd.bfl_upsampling);
 
     for (const style of extension_settings.sd.styles) {
         const option = document.createElement('option');
@@ -1089,15 +1095,14 @@ function onComfyWorkflowChange() {
     saveSettingsDebounced();
 }
 
-async function onStabilityKeyClick() {
+async function onApiKeyClick(popupText, secretKey) {
-    const popupText = 'Stability AI API Key:';
     const key = await callGenericPopup(popupText, POPUP_TYPE.INPUT, '', {
         customButtons: [{
             text: 'Remove Key',
             appendAtEnd: true,
             result: POPUP_RESULT.NEGATIVE,
             action: async () => {
-                await writeSecret(SECRET_KEYS.STABILITY, '');
+                await writeSecret(secretKey, '');
                 toastr.success('API Key removed');
                 await loadSettingOptions();
             },
@@ -1108,12 +1113,25 @@ async function onStabilityKeyClick() {
         return;
     }
 
-    await writeSecret(SECRET_KEYS.STABILITY, String(key));
+    await writeSecret(secretKey, String(key));
 
     toastr.success('API Key saved');
     await loadSettingOptions();
 }
 
+async function onStabilityKeyClick() {
+    return onApiKeyClick('Stability AI API Key:', SECRET_KEYS.STABILITY);
+}
+
+async function onBflKeyClick() {
+    return onApiKeyClick('BFL API Key:', SECRET_KEYS.BFL);
+}
+
+function onBflUpsamplingInput() {
+    extension_settings.sd.bfl_upsampling = !!$('#sd_bfl_upsampling').prop('checked');
+    saveSettingsDebounced();
+}
+
 function onStabilityStylePresetChange() {
     extension_settings.sd.stability_style_preset = String($('#sd_stability_style_preset').val());
     saveSettingsDebounced();
@@ -1238,6 +1256,7 @@ async function onModelChange() {
         sources.blockentropy,
         sources.huggingface,
         sources.nanogpt,
+        sources.bfl,
     ];
 
     if (cloudSources.includes(extension_settings.sd.source)) {
@@ -1459,6 +1478,9 @@ async function loadSamplers() {
         case sources.nanogpt:
             samplers = ['N/A'];
             break;
+        case sources.bfl:
+            samplers = ['N/A'];
+            break;
     }
 
     for (const sampler of samplers) {
@@ -1654,6 +1676,9 @@ async function loadModels() {
         case sources.nanogpt:
             models = await loadNanoGPTModels();
             break;
+        case sources.bfl:
+            models = await loadBflModels();
+            break;
     }
 
     for (const model of models) {
@@ -1680,6 +1705,17 @@ async function loadStabilityModels() {
     ];
 }
 
+async function loadBflModels() {
+    $('#sd_bfl_key').toggleClass('success', !!secret_state[SECRET_KEYS.BFL]);
+
+    return [
+        { value: 'flux-pro-1.1-ultra', text: 'flux-pro-1.1-ultra' },
+        { value: 'flux-pro-1.1', text: 'flux-pro-1.1' },
+        { value: 'flux-pro', text: 'flux-pro' },
+        { value: 'flux-dev', text: 'flux-dev' },
+    ];
+}
+
 async function loadPollinationsModels() {
     const result = await fetch('/api/sd/pollinations/models', {
         method: 'POST',
@@ -2027,6 +2063,9 @@ async function loadSchedulers() {
         case sources.nanogpt:
             schedulers = ['N/A'];
             break;
+        case sources.bfl:
+            schedulers = ['N/A'];
+            break;
     }
 
     for (const scheduler of schedulers) {
@@ -2112,6 +2151,9 @@ async function loadVaes() {
         case sources.nanogpt:
             vaes = ['N/A'];
             break;
+        case sources.bfl:
+            vaes = ['N/A'];
+            break;
     }
 
     for (const vae of vaes) {
@@ -2666,6 +2708,10 @@ async function sendGenerationRequest(generationType, prompt, additionalNegativeP
                 break;
             case sources.nanogpt:
                 result = await generateNanoGPTImage(prefixedPrompt, negativePrompt, signal);
+                break;
+            case sources.bfl:
+                result = await generateBflImage(prefixedPrompt, signal);
+                break;
         }
 
         if (!result.data) {
@@ -3370,7 +3416,7 @@ async function generateNanoGPTImage(prompt, negativePrompt, signal) {
             width: parseInt(extension_settings.sd.width),
             height: parseInt(extension_settings.sd.height),
             resolution: `${extension_settings.sd.width}x${extension_settings.sd.height}`,
-            showExplicitContent:  true,
+            showExplicitContent: true,
             nImages: 1,
         }),
     });
@@ -3384,6 +3430,38 @@ async function generateNanoGPTImage(prompt, negativePrompt, signal) {
     }
 }
 
+/**
+ * Generates an image using the BFL API.
+ * @param {string} prompt - The main instruction used to guide the image generation.
+ * @param {AbortSignal} signal - An AbortSignal object that can be used to cancel the request.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
+ */
+async function generateBflImage(prompt, signal) {
+    const result = await fetch('/api/sd/bfl/generate', {
+        method: 'POST',
+        headers: getRequestHeaders(),
+        signal: signal,
+        body: JSON.stringify({
+            prompt: prompt,
+            model: extension_settings.sd.model,
+            steps: clamp(extension_settings.sd.steps, 1, 50),
+            guidance: clamp(extension_settings.sd.scale, 1.5, 5),
+            width: clamp(extension_settings.sd.width, 256, 1440),
+            height: clamp(extension_settings.sd.height, 256, 1440),
+            prompt_upsampling: !!extension_settings.sd.bfl_upsampling,
+            seed: extension_settings.sd.seed >= 0 ? extension_settings.sd.seed : undefined,
+        }),
+    });
+
+    if (result.ok) {
+        const data = await result.json();
+        return { format: 'jpg', data: data.image };
+    } else {
+        const text = await result.text();
+        throw new Error(text);
+    }
+}
+
 async function onComfyOpenWorkflowEditorClick() {
     let workflow = await (await fetch('/api/sd/comfy/workflow', {
         method: 'POST',
@@ -3668,6 +3746,8 @@ function isValidState() {
             return secret_state[SECRET_KEYS.HUGGINGFACE];
         case sources.nanogpt:
             return secret_state[SECRET_KEYS.NANOGPT];
+        case sources.bfl:
+            return secret_state[SECRET_KEYS.BFL];
     }
 }
 
@@ -4338,6 +4418,8 @@ jQuery(async () => {
     $('#sd_stability_style_preset').on('change', onStabilityStylePresetChange);
     $('#sd_huggingface_model_id').on('input', onHFModelInput);
     $('#sd_function_tool').on('input', onFunctionToolInput);
+    $('#sd_bfl_key').on('click', onBflKeyClick);
+    $('#sd_bfl_upsampling').on('input', onBflUpsamplingInput);
 
     if (!CSS.supports('field-sizing', 'content')) {
         $('.sd_settings .inline-drawer-toggle').on('click', function () {

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

@@ -37,6 +37,7 @@
             </label>
             <label for="sd_source" data-i18n="Source">Source</label>
             <select id="sd_source">
+                <option value="bfl">BFL (Black Forest Labs)</option>
                 <option value="blockentropy">Block Entropy</option>
                 <option value="comfy">ComfyUI</option>
                 <option value="drawthings">DrawThings HTTP API</option>
@@ -234,6 +235,27 @@
                     </div>
                 </div>
             </div>
+
+            <div data-sd-source="bfl">
+                <div class="flex-container flexnowrap alignItemsBaseline marginBot5">
+                    <a href="https://api.bfl.ml/" target="_blank" rel="noopener noreferrer">
+                        <strong data-i18n="API Key">API Key</strong>
+                        <i class="fa-solid fa-share-from-square"></i>
+                    </a>
+                    <span class="expander"></span>
+                    <div id="sd_bfl_key" class="menu_button menu_button_icon">
+                        <i class="fa-fw fa-solid fa-key"></i>
+                        <span data-i18n="Click to set">Click to set</span>
+                    </div>
+                </div>
+                <label class="checkbox_label marginBot5" for="sd_bfl_upsampling" title="Whether to perform upsampling on the prompt. If active, automatically modifies the prompt for more creative generation.">
+                    <input id="sd_bfl_upsampling" type="checkbox" />
+                    <span data-i18n="Prompt Upsampling">
+                        Prompt Upsampling
+                    </span>
+                </label>
+            </div>
+
             <div class="flex-container">
                 <div class="flex1">
                     <label for="sd_model" data-i18n="Model">Model</label>
@@ -385,7 +407,7 @@
                 </label>
             </div>
 
-            <div data-sd-source="novel,togetherai,pollinations,comfy,drawthings,vlad,auto,horde,extras,stability,blockentropy" class="marginTop5">
+            <div data-sd-source="novel,togetherai,pollinations,comfy,drawthings,vlad,auto,horde,extras,stability,blockentropy,bfl" class="marginTop5">
                 <label for="sd_seed">
                     <span data-i18n="Seed">Seed</span>
                     <small data-i18n="(-1 for random)">(-1 for random)</small>

public/scripts/secrets.js

@@ -37,6 +37,7 @@ export const SECRET_KEYS = {
     CUSTOM_OPENAI_TTS: 'api_key_custom_openai_tts',
     NANOGPT: 'api_key_nanogpt',
     TAVILY: 'api_key_tavily',
+    BFL: 'api_key_bfl',
 };
 
 const INPUT_MAP = {

src/endpoints/secrets.js

@@ -49,6 +49,7 @@ export const SECRET_KEYS = {
     CUSTOM_OPENAI_TTS: 'api_key_custom_openai_tts',
     TAVILY: 'api_key_tavily',
     NANOGPT: 'api_key_nanogpt',
+    BFL: 'api_key_bfl',
 };
 
 // These are the keys that are safe to expose, even if allowKeysExposure is false

src/endpoints/stable-diffusion.js

@@ -1101,6 +1101,120 @@ nanogpt.post('/generate', jsonParser, async (request, response) => {
     }
 });
 
+const bfl = express.Router();
+
+bfl.post('/generate', jsonParser, async (request, response) => {
+    try {
+        const key = readSecret(request.user.directories, SECRET_KEYS.BFL);
+
+        if (!key) {
+            console.log('BFL key not found.');
+            return response.sendStatus(400);
+        }
+
+        const requestBody = {
+            prompt: request.body.prompt,
+            steps: request.body.steps,
+            guidance: request.body.guidance,
+            width: request.body.width,
+            height: request.body.height,
+            prompt_upsampling: request.body.prompt_upsampling,
+            seed: request.body.seed ?? null,
+            safety_tolerance: 6, // being least strict
+            output_format: 'jpeg',
+        };
+
+        function getClosestAspectRatio(width, height) {
+            const minAspect = 9 / 21;
+            const maxAspect = 21 / 9;
+            const currentAspect = width / height;
+
+            const gcd = (a, b) => b === 0 ? a : gcd(b, a % b);
+            const simplifyRatio = (w, h) => {
+                const divisor = gcd(w, h);
+                return `${w / divisor}:${h / divisor}`;
+            };
+
+            if (currentAspect < minAspect) {
+                const adjustedHeight = Math.round(width / minAspect);
+                return simplifyRatio(width, adjustedHeight);
+            } else if (currentAspect > maxAspect) {
+                const adjustedWidth = Math.round(height * maxAspect);
+                return simplifyRatio(adjustedWidth, height);
+            } else {
+                return simplifyRatio(width, height);
+            }
+        }
+
+        if (String(request.body.model).endsWith('-ultra')) {
+            requestBody.aspect_ratio = getClosestAspectRatio(request.body.width, request.body.height);
+            delete requestBody.steps;
+            delete requestBody.guidance;
+            delete requestBody.width;
+            delete requestBody.height;
+            delete requestBody.prompt_upsampling;
+        }
+
+        if (String(request.body.model).endsWith('-pro-1.1')) {
+            delete requestBody.steps;
+            delete requestBody.guidance;
+        }
+
+        console.log('BFL request:', requestBody);
+
+        const result = await fetch(`https://api.bfl.ml/v1/${request.body.model}`, {
+            method: 'POST',
+            body: JSON.stringify(requestBody),
+            headers: {
+                'Content-Type': 'application/json',
+                'x-key': key,
+            },
+        });
+
+        if (!result.ok) {
+            console.log('BFL returned an error.');
+            return response.sendStatus(500);
+        }
+
+        /** @type {any} */
+        const taskData = await result.json();
+        const { id } = taskData;
+
+        const MAX_ATTEMPTS = 100;
+        for (let i = 0; i < MAX_ATTEMPTS; i++) {
+            await delay(2500);
+
+            const statusResult = await fetch(`https://api.bfl.ml/v1/get_result?id=${id}`);
+
+            if (!statusResult.ok) {
+                const text = await statusResult.text();
+                console.log('BFL returned an error.', text);
+                return response.sendStatus(500);
+            }
+
+            /** @type {any} */
+            const statusData = await statusResult.json();
+
+            if (statusData?.status === 'Pending') {
+                continue;
+            }
+
+            if (statusData?.status === 'Ready') {
+                const { sample } = statusData.result;
+                const fetchResult = await fetch(sample);
+                const fetchData = await fetchResult.arrayBuffer();
+                const image = Buffer.from(fetchData).toString('base64');
+                return response.send({ image: image });
+            }
+
+            throw new Error('BFL failed to generate image.', { cause: statusData });
+        }
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
 router.use('/comfy', comfy);
 router.use('/together', together);
 router.use('/drawthings', drawthings);
@@ -1109,3 +1223,4 @@ router.use('/stability', stability);
 router.use('/blockentropy', blockentropy);
 router.use('/huggingface', huggingface);
 router.use('/nanogpt', nanogpt);
+router.use('/bfl', bfl);