BareGit
import { Component, h, render } from "/ui/vendor/preact.module.js";

const INITIAL_FORM = {
    prompt: "",
    negative_prompt: "",
    width: 832,
    height: 1248,
    batch_size: 1,
    steps: 10,
    cfg: 1.0,
    seed: -1,
    output_extension: "png",
    output_quality: 95,
};

function apiMessage(error, fallback) {
    if(error && typeof error.message === "string" && error.message)
    {
        return error.message;
    }
    return fallback;
}

async function readJson(response) {
    const text = await response.text();
    if(!text)
    {
        return null;
    }
    try
    {
        return JSON.parse(text);
    }
    catch(_error)
    {
        return null;
    }
}

function fieldValue(value, fallback) {
    if(value === "" || value === null || value === undefined)
    {
        return fallback;
    }
    return value;
}

function NumberField(props) {
    return h(
        "label",
        { class: "field" },
        h("span", { class: "field-label" }, props.label),
        h("input", {
            type: "number",
            name: props.name,
            min: props.min,
            max: props.max,
            step: props.step,
            value: props.value,
            disabled: props.disabled,
            onInput: props.onInput,
        }),
    );
}

function SelectField(props) {
    return h(
        "label",
        { class: "field" },
        h("span", { class: "field-label" }, props.label),
        h(
            "select",
            {
                name: props.name,
                value: props.value,
                disabled: props.disabled,
                onInput: props.onInput,
            },
            props.options.map((option) => h(
                "option",
                { key: option.value, value: option.value },
                option.label,
            )),
        ),
    );
}

function TextAreaField(props) {
    return h(
        "label",
        { class: "field field-wide" },
        h("span", { class: "field-label" }, props.label),
        h("textarea", {
            name: props.name,
            rows: props.rows,
            value: props.value,
            disabled: props.disabled,
            onInput: props.onInput,
        }),
    );
}

function ResolutionField(props) {
    const groups = {};
    for(const preset of props.presets)
    {
        if(!groups[preset.resolution])
        {
            groups[preset.resolution] = [];
        }
        groups[preset.resolution].push(preset);
    }

    return h(
        "label",
        { class: "field field-wide" },
        h("span", { class: "field-label" }, "Resolution"),
        h(
            "select",
            {
                name: "resolution",
                value: props.value,
                disabled: props.disabled,
                onInput: props.onInput,
            },
            Object.keys(groups).map((resolution) => h(
                "optgroup",
                { key: resolution, label: resolution },
                groups[resolution].map((preset) => h(
                    "option",
                    {
                        key: preset.width + "x" + preset.height,
                        value: preset.width + "x" + preset.height,
                    },
                    preset.label,
                )),
            )),
        ),
    );
}

function StatusMessage(props) {
    const class_name = props.error
        ? "status-message status-error"
        : "status-message";
    return h(
        "div",
        { class: class_name, role: props.error ? "alert" : null },
        props.children,
    );
}

function ResultGallery(props) {
    if(props.images.length === 0)
    {
        return h(
            "section",
            { class: "results empty-results", "aria-label": "Results" },
            h("div", { class: "empty-box" }, "No image generated"),
        );
    }

    return h(
        "section",
        { class: "results", "aria-label": "Results" },
        props.images.map((image, index) => {
            const url = "data:" + image.mime_type + ";base64," + image.data;
            return h(
                "figure",
                { class: "result-item", key: index },
                h(
                    "a",
                    {
                        href: url,
                        target: "_blank",
                        rel: "noreferrer",
                    },
                    h("img", { src: url, alt: "Generated image" }),
                ),
                h("figcaption", null, "Seed ", String(image.seed)),
            );
        }),
    );
}

function GenerationForm(props) {
    const form = props.form;
    const disabled = props.disabled;
    const numeric = [
        ["steps", "Steps", 1, 1],
        ["cfg", "CFG", 0, 0.1],
        ["batch_size", "Batch", 1, 1],
    ];

    return h(
        "form",
        { class: "generation-form", onSubmit: props.onSubmit },
        h(TextAreaField, {
            label: "Prompt",
            name: "prompt",
            rows: 7,
            value: form.prompt,
            disabled,
            onInput: props.onInput,
        }),
        h(TextAreaField, {
            label: "Negative prompt",
            name: "negative_prompt",
            rows: 4,
            value: form.negative_prompt,
            disabled,
            onInput: props.onInput,
        }),
        h(ResolutionField, {
            presets: props.resolution_presets,
            value: form.width + "x" + form.height,
            disabled,
            onInput: props.onResolutionInput,
        }),
        h(
            "div",
            { class: "field-grid" },
            numeric.map((item) => h(NumberField, {
                key: item[0],
                name: item[0],
                label: item[1],
                min: item[2],
                step: item[3],
                value: form[item[0]],
                disabled,
                onInput: props.onInput,
            })),
            h(SelectField, {
                name: "output_extension",
                label: "Format",
                value: form.output_extension,
                options: props.output_formats,
                disabled,
                onInput: props.onInput,
            }),
            h(NumberField, {
                name: "output_quality",
                label: "Quality",
                min: 1,
                max: 100,
                step: 1,
                value: form.output_quality,
                disabled,
                onInput: props.onInput,
            }),
            h(
                "div",
                { class: "seed-row" },
                h(NumberField, {
                    name: "seed",
                    label: "Seed",
                    step: 1,
                    value: form.seed,
                    disabled: disabled || form.seed === -1,
                    onInput: props.onInput,
                }),
                h(
                    "label",
                    { class: "random-seed" },
                    h("input", {
                        type: "checkbox",
                        checked: form.seed === -1,
                        disabled,
                        onChange: props.onRandomSeed,
                    }),
                    h("span", null, "Random"),
                ),
            ),
        ),
        h(
            "button",
            {
                class: "primary-action",
                type: "submit",
                disabled,
            },
            props.is_generating ? "Generating" : "Generate",
        ),
    );
}

class App extends Component {
    constructor() {
        super();
        this.state = {
            defaults: null,
            health: null,
            form: { ...INITIAL_FORM },
            resolution_presets: [],
            output_formats: [
                { value: "png", label: "PNG" },
                { value: "jpg", label: "JPEG" },
                { value: "webp", label: "WebP" },
                { value: "avif", label: "AVIF" },
            ],
            is_loading_defaults: true,
            is_generating: false,
            error_message: "",
            images: [],
        };
        this.handleInput = this.handleInput.bind(this);
        this.handleResolutionInput = this.handleResolutionInput.bind(this);
        this.handleRandomSeed = this.handleRandomSeed.bind(this);
        this.handleSubmit = this.handleSubmit.bind(this);
    }

    componentDidMount() {
        this.loadInitialState();
    }

    async loadInitialState() {
        try
        {
            const [defaults_response, health_response] = await Promise.all([
                fetch("/ui/api/defaults"),
                fetch("/ui/api/health"),
            ]);
            const defaults = await readJson(defaults_response);
            const health = await readJson(health_response);
            if(!defaults_response.ok || !defaults)
            {
                throw new Error("Could not load defaults");
            }
            const resolution_presets = defaults.resolution_presets || [];
            const output_formats = defaults.output_formats
                || this.state.output_formats;
            const form_defaults = { ...defaults };
            delete form_defaults.resolution_presets;
            delete form_defaults.output_formats;
            const form = { ...INITIAL_FORM, ...form_defaults, prompt: "" };
            const is_valid_resolution = resolution_presets.some(
                (preset) => (
                    preset.width === form.width
                    && preset.height === form.height
                ),
            );
            if(!is_valid_resolution && resolution_presets.length > 0)
            {
                form.width = resolution_presets[0].width;
                form.height = resolution_presets[0].height;
            }
            this.setState({
                defaults,
                health: health_response.ok ? health : null,
                form,
                resolution_presets,
                output_formats,
                is_loading_defaults: false,
            });
        }
        catch(_error)
        {
            this.setState({
                is_loading_defaults: false,
                error_message: "Could not load defaults.",
            });
        }
    }

    handleInput(event) {
        const target = event.currentTarget;
        const name = target.name;
        let value = target.value;
        if(target.type === "number")
        {
            value = value === "" ? "" : Number(value);
        }
        this.setState((state) => ({
            form: { ...state.form, [name]: value },
        }));
    }

    handleResolutionInput(event) {
        const value = event.currentTarget.value;
        const preset = this.state.resolution_presets.find(
            (item) => value === item.width + "x" + item.height,
        );
        if(!preset)
        {
            return;
        }
        this.setState((state) => ({
            form: {
                ...state.form,
                width: preset.width,
                height: preset.height,
            },
        }));
    }

    handleRandomSeed(event) {
        const checked = event.currentTarget.checked;
        this.setState((state) => ({
            form: {
                ...state.form,
                seed: checked ? -1 : 0,
            },
        }));
    }

    requestBody() {
        const form = this.state.form;
        return {
            prompt: form.prompt,
            negative_prompt: form.negative_prompt,
            width: fieldValue(form.width, null),
            height: fieldValue(form.height, null),
            batch_size: fieldValue(form.batch_size, null),
            steps: fieldValue(form.steps, null),
            cfg: fieldValue(form.cfg, null),
            seed: fieldValue(form.seed, -1),
            output_extension: form.output_extension,
            output_quality: fieldValue(form.output_quality, null),
        };
    }

    async handleSubmit(event) {
        event.preventDefault();
        this.setState({
            error_message: "",
            is_generating: true,
        });

        try
        {
            const response = await fetch("/ui/api/txt2img", {
                method: "POST",
                headers: { "Content-Type": "application/json" },
                body: JSON.stringify(this.requestBody()),
            });
            const data = await readJson(response);
            if(!response.ok)
            {
                throw new Error(apiMessage(
                    data,
                    "Request failed. Check the server log for details.",
                ));
            }
            this.setState({ images: data.images || [] });
        }
        catch(error)
        {
            this.setState({
                error_message: apiMessage(
                    error,
                    "Request failed. Check the server log for details.",
                ),
            });
        }
        finally
        {
            this.setState({ is_generating: false });
        }
    }

    render(_props, state) {
        const disabled = state.is_loading_defaults || state.is_generating;
        const status_text = state.is_generating
            ? "Generating"
            : state.is_loading_defaults ? "Loading" : "Ready";
        const health_text = state.health
            ? state.health.api_profile + " / " + state.health.model_residency
            : "Server";

        return h(
            "div",
            { class: "app-shell" },
            h(
                "header",
                { class: "top-bar" },
                h("div", { class: "brand" }, "diffusion-cli"),
                h(
                    "div",
                    { class: "server-state" },
                    h("span", null, health_text),
                    h("span", { class: "state-pill" }, status_text),
                ),
            ),
            h(
                "main",
                { class: "workspace" },
                h(
                    "section",
                    { class: "form-panel", "aria-label": "Generation" },
                    state.error_message
                        ? h(StatusMessage, { error: true },
                            state.error_message)
                        : h(StatusMessage, null, status_text),
                    h(GenerationForm, {
                        form: state.form,
                        resolution_presets: state.resolution_presets,
                        output_formats: state.output_formats,
                        disabled,
                        is_generating: state.is_generating,
                        onInput: this.handleInput,
                        onResolutionInput: this.handleResolutionInput,
                        onRandomSeed: this.handleRandomSeed,
                        onSubmit: this.handleSubmit,
                    }),
                ),
                h(ResultGallery, { images: state.images }),
            ),
        );
    }
}

render(h(App), document.getElementById("app"));