StableDiffusion Client

sdclient.js

class StableDiffusionClient {
    constructor(fetch) {
        this.fetch = fetch;
        this.settings = {};
    }

    get width() { return this.settings.width || 768; }
    get height() { return this.settings.height || 512; }

    parse(prompt, preset_pos, preset_neg) {
        const positives = (preset_pos || "").split(',').filter((x) => !!x.trim());
        const negatives = (preset_neg || "").split(',').filter((x) => !!x.trim());

        for(let tag of prompt.trim().split(',')) {
            tag = tag.trim();
            if(tag.startsWith('-')) {
                negatives.push(tag.slice(1));
            } else if(tag) {
                positives.push(tag);
            }
        }

        return {
            desired: positives.join(',').trim(),
            undesired: negatives.join(',').trim(),
        }
    }

    async generateFor(elem, prompt) {
        const document = elem.ownerDocument;
        const container = document.createElement('div');
        container.classList.add('sd-container');
        elem.replaceWith(container);

        const status_bar = document.createElement('p');
        status_bar.classList.add('sd-status');
        container.appendChild(status_bar);
        
        const images = document.createElement('div');
        images.classList.add('sd-images');
        container.appendChild(images);

        const image_count = this.settings.n_samples ?? 1;
        for(let i=0; i<image_count; ++i) {
            let error = null;
            const placeholder_elem = document.createElement('div');
            placeholder_elem.classList.add('sd-placeholder');
            placeholder_elem.style = `width: ${this.width}px; height: ${this.height}px; background-color: black;`;
            images.appendChild(placeholder_elem);
            let curr_elem = null;
            await this.generateOne(prompt, (kind, data) => {
                switch(kind) {
                    case 'progress':
                        status_bar.textContent = `Image ${i+1}/${image_count}: ${data}`;
                        break;
                    case 'preview':
                    case 'end':
                        if(curr_elem) {
                            curr_elem.src = data.src;
                        } else {
                            curr_elem = data;
                            curr_elem.alt = curr_elem.title = prompt;
                            curr_elem.classList.add('sd-image');
                            curr_elem.width = this.width;
                            curr_elem.height = this.height;
                            curr_elem.style = "display: inline-block;";
                            images.removeChild(placeholder_elem);
                            images.appendChild(curr_elem);
                        }
                        break;
                    case 'error':
                        error = data;
                        break;
                }
            });

            if(error) {
                console.error(error);
                throw error;
            }
        }

        container.removeChild(status_bar);
    }

    // imageCallback: (ind, 'preview'|'progress'|'end'|'error', data)
    async generate(prompt, imageCallback) {
        const images = [];
        for(let i=0; i<(this.settings.n_samples ?? 1); ++i) {
            images.push(await this.generateOne((kind, data) => imageCallback(i, kind, data)));
        }
        return images;
    }
    
    // imageCallback: ('preview'|'progress'|'end'|'error', data)
    async generateOne(prompt, imageCallback) {
        if(!this.settings.endpoint) {
            imageCallback('error', new Error("Endpoint not specified!"));
            return null;
        }

        prompt = prompt.replace(/_/g, ' ');

        const {desired, undesired} = this.parse(prompt, this.settings.desired, this.settings.undesired);
        if(!desired) return null;

        const request_id = `web-${Math.random().toString().slice(2)}`;
        const seed = Math.floor(Math.random() * 100_000_000);
        const body = {
            request_id,
            prompt: desired,
            width: this.width, height: this.height,
            cfg_scale: this.settings.scale || 8,
            sampler_name: this.settings.sampler || "DPM++ 2S a Karras",
            steps: this.settings.steps || 25,
            seed: seed ?? -1,
            n_iter: 1,
            styles: (this.settings.styles || "").split(",").map((x) => x.trim()).filter((x) => !!x),
            negative_prompt: undesired,
            override_settings: {
                sd_model_checkpoint: this.settings.checkpoint || "AbyssOrangeMix2-sfw",
                show_progress_every_n_steps: 10,
            },
        };

        if(this.settings.vae != null) {
            body.override_settings.sd_vae = this.settings.vae;
        }

        const headers = {
            "Content-Type": "application/json",
            'Origin': '',
        };

        if(this.settings.account) headers['Authorization'] = `Basic ${btoa(this.settings.account)}`;

        let preview_timer = null;
        try {
            console.log('StableDiffusion', desired, undesired);

            let preview_image_data = null;
            preview_timer = setInterval(async() => {
                if(preview_timer == null) return;
                const progress_status = await (await this.fetch(`${this.settings.endpoint}sdapi/v1/progress`, {
                    method: 'GET', credentials: 'include', headers,
                })).json();
                if(preview_timer == null || progress_status == null) return;

                const state = progress_status.state ?? {};
                let eta = progress_status.eta_relative;
                
                if(typeof eta === 'number') eta = Math.max(0, eta).toFixed(2);
        
                if(request_id !== state.request_id) {
                    imageCallback('progress', `Job in queue... (ETA: ${eta})`);
                    return;
                }

                const progress_percentage = progress_status.progress != null ? (progress_status.progress * 100).toFixed(1) : '??';
                const progress_step = `step ${(progress_status.progress >= 1 ? state.sampling_steps : state.sampling_step) ?? '?'} of ${state.sampling_steps || '?'}`;
                const progress_str = "$progress (ETA: $eta sec)".replace('$progress', `${progress_step} (${progress_percentage} %)`).replace('$eta', eta);

                imageCallback('progress', progress_str);
                if(progress_status.current_image !== preview_image_data) {
                    preview_image_data = progress_status.current_image;
                    const preview_image = new Image();
                    preview_image.src = 'data:image/png;base64,' + preview_image_data;
                    imageCallback('preview', preview_image);
                }
        
            }, 1000);
            
            const result = await this.fetch(`${this.settings.endpoint}sdapi/v1/txt2img`, {
                method: 'POST', credentials: 'include',
                headers, body: JSON.stringify(body)
            });

            clearInterval(preview_timer);
            preview_timer = null;

            const image_data = (await result.json()).images[0];
            const image = new Image();
            image.src = 'data:image/png;base64,' + image_data;
            
            imageCallback('end', image);
            return image;
        } catch(e) {
            console.error(e);
            imageCallback('error', e);

            return null;
        } finally {
            if(preview_timer != null) clearInterval(preview_timer);
        }
    }
}

window.StableDiffusionClient = StableDiffusionClient;

/*
Example prompt for ChatGPT:

`From now, when I request something along the line of "show me an image", display an image using Markdown syntax (do not enclose it in a code block),
with the URL "http://image.invalid" and the description of the image for the alt text.
Fill-in any unspecified details for the description, so that it would be vivid and concrete.`

(async function main(window, document) {
    if(".GPTE" in window) return;

    const client = new StableDiffusionClient(window.fetch);
    client.settings.endpoint = "<Insert your endpoint here>";
    client.settings.styles = "Default";
    client.settings.n_samples = 1;

    setInterval(() => {
        const images = document.querySelectorAll("div.markdown img[alt]:not(.sd-image)");
        for(const image of images) {
            client.generateFor(image, image.alt);
        }
    }, 500);
})(window, document);
*/