Merge pull request #24 from jtydhr88/webui-plugin

Support local inference via the A1111 plugin

README.md

@@ -4,18 +4,16 @@
 <img src="./misc/GenerateScreenshot.png" style="width: 400px; max-width: 600px; flex-grow: 1;" />
 <img src="./misc/EditScreenshot.png" style="width: 400px; max-width: 600px; flex-grow: 1;" />
 
-<h3>
-<span>👋 Welcome to StableStudio, the open-source version of <a href="https://dreamstudio.ai" target="_blank">DreamStudio</a></span>
-<br />
-<br />
-<span>[ <a href="./packages/stablestudio-ui/README.md">🎨 UI README</a> ]</span>
-<span>[ <a href="./packages/stablestudio-plugin/README.md" href="./packages/stablestudio-ui/README.md">🔌 Plugins README</a> ]</span>
-<span>[ <a href="https://discord.gg/stablediffusion" target="_blank">🎮 Discord</a> ]</span>
-<span>[ <a href="https://github.com/Stability-AI/StableStudio/issues">🛟 Bugs & Support</a> ]</span>
-<span>[ <a href="https://github.com/Stability-AI/StableStudio/discussions">💬 Discussion</a> ]</span>
-</h3>
+<h3>👋 Welcome to StableStudio, the open-source version of <a href="https://dreamstudio.ai" target="_blank">DreamStudio</a>!</h3>
 
-<hr />
+**🗺 Contents – [🚀 Quick Start](#quick-start) · [ℹ️ About](#about) · [🙋 FAQ](#faq) · [🧑‍💻 Contributing](#contributing)**
+
+**📚 Documentation – [🎨 UI](./packages/stablestudio-ui/README.md) · [🔌 Plugins](./packages/stablestudio-plugin/README.md) · <a href="https://platform.stability.ai" target="_blank">⚡️ platform.stability.ai</a>**
+
+**🔗 Links – <a href="https://discord.com/channels/1002292111942635562/1108055793674227782" target="_blank">🎮 Discord</a> · <a href="https://dreamstudio.ai" target="_blank">🌈 DreamStudio</a> · <a href="https://github.com/Stability-AI/StableStudio/issues">🛟 Bugs & Support</a> · <a href="https://github.com/Stability-AI/StableStudio/discussions">💬 Discussion</a>**
+
+<br />
+<br />
 
 </div>
 

package.json

@@ -15,6 +15,7 @@
     "stablestudio-plugin-webui": "yarn workspace @stability/stablestudio-plugin-webui",
     "stablestudio-ui": "yarn workspace @stability/stablestudio-ui",
     "dev:use-example-plugin": "cross-env VITE_USE_EXAMPLE_PLUGIN=true yarn dev",
+    "dev:use-webui-plugin": "cross-env VITE_USE_WEBUI_PLUGIN=true yarn dev",
     "dev": "yarn workspaces foreach --all --interlaced --verbose --parallel --jobs unlimited run dev",
     "build": "yarn workspaces foreach --all --interlaced --verbose --jobs unlimited run build",
     "clean": "yarn workspaces foreach --all --interlaced --verbose --parallel --jobs unlimited run clean && rimraf node_modules"

packages/stablestudio-plugin-webui/README.md

@@ -1 +1,67 @@
-Soon™️
+<div align="center">
+
+# 🔌 [`stable-diffusion-webui`](https://github.com/AUTOMATIC1111/stable-diffusion-webui) Plugin
+
+**🗺 Contents – [ℹ️ About](#about) · [⚙️ Usage](#usage) · [⭐️ Features](#features)**
+
+**[⬆️ Top-Level README](../../README.md)**
+
+![Electric1111](../../misc/Electric1111.png)
+
+</div>
+
+# <a id="about" href="#about">ℹ️ About</a>
+
+This plugin enables StableStudio to run using [`stable-diffusion-webui`](https://github.com/AUTOMATIC1111/stable-diffusion-webui), which means you can generate images entirely on your own machine!
+
+Thanks goes to [Terry Jia](https://github.com/jtydhr88) for the original work on this plugin.
+
+# <a id="usage" href="#usage">⚙️ Usage</a>
+
+1. First, you'll need to configure your local installation of `stable-diffusion-webui` to run without the UI and with CORS enabled.
+
+   **Windows**
+
+   Edit the command line arguments within `webui-user.bat`:
+
+   ```
+   set COMMANDLINE_ARGS=--nowebui --cors-allow-origins=http://localhost:3000
+   ```
+
+   **Mac**
+
+   Edit the command line arguments within `webui-macos-env.sh`:
+
+   ```
+   export COMMANDLINE_ARGS="--nowebui --cors-allow-origins=http://localhost:3000"
+   ```
+
+2. Start `stable-diffusion-webui` and look for `INFO: Uvicorn running on http://127.0.0.1:7861`.
+
+   You can make sure everything is running correctly by checking to see if [`http://127.0.0.1:7861/docs`](http://127.0.0.1:7861/docs) displays API documentation.
+
+3. Within your installation of StableStudio, run `yarn dev:use-webui-plugin`.
+
+   _**That's it!**_ 🎉 You should now be able to generate images using your local machine.
+
+## <a id="image-history" href="#image-history">💾 Image History</a>
+
+To persist your image history, you'll need to install the [`sd-webui-StableStudio`](https://github.com/jtydhr88/sd-webui-StableStudio) extension for `stable-diffusion-webui`.
+
+> 🛑 Be wary installing third-party extensions for `stable-diffusion-webui`, it's always a good idea to check before running untrusted code.
+
+# <a id="features" href="#features">⭐️ Features</a>
+
+Missing something? Please [let us know](https://github.com/Stability-AI/StableStudio/issues/new/choose)!
+
+- [x] Text-to-image
+- [x] Image-to-image
+- [x] Basic features (prompt, negative prompt, steps, batch size, image size)
+- [x] Model selection
+- [x] Sampler selection
+- [x] Masking, in-painting, and out-painting
+- [x] Settings storage
+- [x] Accurate plugin status
+- [x] [Loading existing images]("#image-history)
+- [x] Upscaling
+- [ ] Lora support

packages/stablestudio-plugin-webui/src/Utilities.ts

@@ -0,0 +1,184 @@
+import { StableDiffusionInput } from "@stability/stablestudio-plugin";
+
+export function base64ToBlob(base64: string, contentType = ""): Promise<Blob> {
+  return fetch(`data:${contentType};base64,${base64}`).then((res) =>
+    res.blob()
+  );
+}
+
+export function blobToBase64(blob: Blob): Promise<string> {
+  return new Promise((resolve, reject) => {
+    const reader = new FileReader();
+
+    reader.onloadend = () => resolve(reader.result as string);
+    reader.onerror = reject;
+
+    reader.readAsDataURL(blob);
+  });
+}
+
+export async function fetchOptions(baseUrl: string | undefined) {
+  const optionsResponse = await fetch(`${baseUrl}/sdapi/v1/options`, {
+    method: "GET",
+    headers: {
+      "Content-Type": "application/json",
+    },
+  });
+
+  return await optionsResponse.json();
+}
+
+export async function setOptions(baseUrl: string | undefined, options: any) {
+  const optionsResponse = await fetch(`${baseUrl}/sdapi/v1/options`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(options),
+  });
+
+  return await optionsResponse.json();
+}
+
+export async function getImageInfo(
+  baseUrl: string | undefined,
+  base64image: any
+) {
+  const imageInfoResponse = await fetch(`${baseUrl}/sdapi/v1/png-info`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify({ image: base64image }),
+  });
+
+  const imageInfoJson = await imageInfoResponse.json();
+
+  const info = imageInfoJson.info.split("\n");
+
+  const data: any = {};
+
+  if (info.length === 0) {
+    return data;
+  }
+
+  data.prompt = info[0];
+
+  let detailIndex = 1;
+
+  if (info.length === 3) {
+    data.nagtivePrompt = info[1].split(":")[1].trim();
+
+    detailIndex = 2;
+  }
+
+  const details = info[detailIndex].split(",");
+
+  details.map((detail: any) => {
+    const detailInfo = detail.trim().split(":");
+
+    data[detailInfo[0]] = detailInfo[1].trim();
+  });
+
+  return data;
+}
+
+export async function testForHistoryPlugin(webuiHostUrl: string) {
+  // timeout after 1 second
+  const finished = Promise.race([
+    fetch(`${webuiHostUrl}/StableStudio/get-generated-images`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        limit: 1,
+      }),
+    }),
+    new Promise((_, reject) =>
+      setTimeout(() => reject(new Error("Request timed out")), 1000)
+    ),
+  ]);
+
+  try {
+    await finished;
+    return (finished as any).ok;
+  } catch (error) {
+    return false;
+  }
+}
+
+export async function constructPayload(
+  options: {
+    input?: StableDiffusionInput | undefined;
+    count?: number | undefined;
+  },
+  isUpscale = false,
+  upscaler: string | undefined
+) {
+  const { sampler, prompts, initialImage, maskImage, width, height, steps } =
+    options?.input ?? {};
+
+  // Construct payload
+  const data: any = {
+    seed: options?.input?.seed === 0 ? -1 : options?.input?.seed,
+    cfgScale: options?.input?.cfgScale ?? 7,
+  };
+
+  if (isUpscale) {
+    /*
+      Upscaling values
+    */
+
+    data.upscaling_resize_w = width ?? 512;
+    data.upscaling_resize_h = height ?? 512;
+    data.upscaler_1 = upscaler;
+  } else {
+    /*
+      regular image generation values
+    */
+
+    data.width = width ?? 512;
+    data.height = height ?? 512;
+
+    data.sampler_name = sampler?.name ?? "";
+    data.sampler_index = sampler?.name ?? "";
+
+    data.prompt =
+      prompts?.find((p) => (p.text && (p.weight ?? 0) > 0) ?? 0 > 0)?.text ??
+      "";
+    data.negative_prompt =
+      prompts?.find((p) => (p.text && (p.weight ?? 0) < 0) ?? 0 < 0)?.text ??
+      "";
+
+    data.steps = steps ?? 20;
+    data.batch_size = options?.count;
+    data.save_images = true;
+  }
+
+  if (initialImage?.weight && !isUpscale) {
+    data.denoising_strength = 1 - initialImage.weight;
+  }
+
+  if (initialImage?.blob) {
+    const initImgB64 = await blobToBase64(initialImage?.blob);
+
+    if (isUpscale) {
+      data.image = initImgB64.split(",")[1];
+    } else {
+      data.init_images = [initImgB64.split(",")[1]];
+    }
+  }
+
+  if (maskImage?.blob) {
+    const maskImgB64 = await blobToBase64(maskImage?.blob);
+
+    data.mask = maskImgB64.split(",")[1];
+
+    data.inpainting_mask_invert = 1; // Mask mode
+    data.inpainting_fill = 1; // Masked content
+    data.inpaint_full_res = false; // Inpaint area
+  }
+
+  return data;
+}

packages/stablestudio-plugin-webui/src/index.ts

@@ -1,3 +1,392 @@
 import * as StableStudio from "@stability/stablestudio-plugin";
+import { StableDiffusionImage } from "@stability/stablestudio-plugin";
 
-export const createPlugin = StableStudio.createPlugin(() => ({}));
+import {
+  base64ToBlob,
+  constructPayload,
+  fetchOptions,
+  getImageInfo,
+  setOptions,
+  testForHistoryPlugin,
+} from "./Utilities";
+
+const manifest = {
+  name: "stable-diffusion-webui",
+  author: "Terry Jia",
+  link: "https://github.com/jtydhr88",
+  icon: `${window.location.origin}/DummyImage.png`,
+  version: "0.0.0",
+  license: "MIT",
+  description:
+    "This plugin uses [`stable-diffusion-webui`](https://github.com/AUTOMATIC1111/stable-diffusion-webui) as its back-end for inference",
+};
+
+const webuiUpscalers = [
+  {
+    label: "None",
+    value: "None",
+  },
+  {
+    label: "Lanczos",
+    value: "Lanczos",
+  },
+  {
+    label: "Nearest",
+    value: "Nearest",
+  },
+  {
+    label: "ESRGAN_4x",
+    value: "ESRGAN_4x",
+  },
+  {
+    label: "LDSR",
+    value: "LDSR",
+  },
+  {
+    label: "R-ESRGAN 4x+",
+    value: "R-ESRGAN 4x+",
+  },
+  {
+    label: "R-ESRGAN 4x+ Anime6B",
+    value: "R-ESRGAN 4x+ Anime6B",
+  },
+  {
+    label: "ScuNET GAN",
+    value: "ScuNET GAN",
+  },
+  {
+    label: "ScuNET PSNR",
+    value: "ScuNET PSNR",
+  },
+  {
+    label: "SwinIR_4x",
+    value: "SwinIR_4x",
+  },
+];
+
+const getNumber = (strValue: string | null, defaultValue: number) => {
+  let retValue = defaultValue;
+
+  if (strValue) {
+    retValue = Number(strValue);
+  }
+
+  return retValue;
+};
+
+const getStableDiffusionDefaultCount = () => 4;
+export const createPlugin = StableStudio.createPlugin<{
+  settings: {
+    baseUrl: StableStudio.PluginSettingString;
+    upscaler: StableStudio.PluginSettingString;
+    historyImagesCount: StableStudio.PluginSettingNumber;
+  };
+}>(({ set, get }) => {
+  const webuiLoad = (
+    webuiHostUrl?: string
+  ): Pick<
+    StableStudio.Plugin,
+    | "createStableDiffusionImages"
+    | "getStatus"
+    | "getStableDiffusionModels"
+    | "getStableDiffusionSamplers"
+    | "getStableDiffusionDefaultCount"
+    | "getStableDiffusionDefaultInput"
+    | "getStableDiffusionExistingImages"
+  > => {
+    webuiHostUrl = webuiHostUrl ?? "http://127.0.0.1:7861";
+
+    return {
+      createStableDiffusionImages: async (options) => {
+        if (!options) {
+          throw new Error("options is required");
+        }
+
+        // fetch the current webui options (model/sampler/etc)
+        const webUIOptions = await fetchOptions(webuiHostUrl);
+
+        const { model, sampler, initialImage } = options?.input ?? {};
+        options.count = options?.count ?? getStableDiffusionDefaultCount();
+
+        // quickly save the sampler and model name to local storage
+        if (sampler?.name) {
+          localStorage.setItem("webui-saved-sampler", sampler.name);
+        }
+
+        if (model) {
+          localStorage.setItem("webui-saved-model", model);
+        }
+
+        // little hacky until StableStudio is better with upscaling
+        const isUpscale =
+          options?.input?.initialImage?.weight === 1 &&
+          model === "esrgan-v1-x2plus";
+
+        // WebUI doesn't have the right model loaded, switch the model
+        if (model && model !== webUIOptions.sd_model_checkpoint && !isUpscale) {
+          localStorage.setItem("webui-saved-model", model);
+          const modelResponse = await setOptions(webuiHostUrl, {
+            sd_model_checkpoint: model,
+          });
+
+          if (modelResponse.ok) {
+            console.log("applied model");
+          }
+        }
+
+        // Construct payload for webui
+        const data = await constructPayload(
+          options,
+          isUpscale,
+          get().settings.upscaler.value
+        );
+
+        // Send payload to webui
+        const response = await fetch(
+          initialImage
+            ? isUpscale
+              ? `${webuiHostUrl}/sdapi/v1/extra-single-image`
+              : `${webuiHostUrl}/sdapi/v1/img2img`
+            : `${webuiHostUrl}/sdapi/v1/txt2img`,
+          {
+            method: "POST",
+            headers: {
+              "Content-Type": "application/json",
+            },
+            body: JSON.stringify(data),
+          }
+        );
+
+        const responseData = await response.json();
+
+        const images = [];
+        const createdAt = new Date();
+
+        if (isUpscale) {
+          // Upscaling only returns one image
+          const blob = await base64ToBlob(responseData.image, "image/jpeg");
+
+          const image = {
+            id: `${Math.random() * 10000000}`,
+            createdAt: createdAt,
+            blob: blob,
+            input: {
+              model: model ?? "",
+            },
+          };
+
+          images.push(image);
+        } else {
+          // Image generation returns an array of images
+          const startIndex =
+            responseData.images.length > data.batch_size ? 1 : 0;
+
+          for (let i = startIndex; i < responseData.images.length; i++) {
+            const blob = await base64ToBlob(
+              responseData.images[i],
+              "image/jpeg"
+            );
+
+            const image: StableDiffusionImage = {
+              id: `${Math.random() * 10000000}`,
+              createdAt,
+              blob,
+              input: {
+                prompts: options?.input?.prompts ?? [],
+                steps: options?.input?.steps ?? 0,
+                seed: responseData.images[i].seed,
+                model: model ?? "",
+                width: options?.input?.width ?? 512,
+                height: options?.input?.height ?? 512,
+                cfgScale: options?.input?.cfgScale ?? 7,
+                sampler: sampler ?? { id: "", name: "" },
+              },
+            };
+
+            images.push(image);
+          }
+        }
+
+        return {
+          id: `${Math.random() * 10000000}`,
+          images: images,
+        };
+      },
+
+      getStableDiffusionModels: async () => {
+        const response = await fetch(`${webuiHostUrl}/sdapi/v1/sd-models`);
+        const responseData = await response.json();
+
+        return responseData.map((model: any) => ({
+          id: model.title,
+          name: model.model_name,
+        }));
+      },
+
+      getStatus: async () => {
+        const optionsResponse = await fetch(`${webuiHostUrl}/sdapi/v1/options`);
+        const hasWebuiHistoryPlugin = await testForHistoryPlugin(
+          `${webuiHostUrl}`
+        );
+
+        return optionsResponse.ok
+          ? {
+              indicator: hasWebuiHistoryPlugin ? "success" : "info",
+              text: `Ready ${
+                hasWebuiHistoryPlugin ? "with" : "without"
+              } history plugin`,
+            }
+          : {
+              indicator: "error",
+              text: "unable to connect webui on " + webuiHostUrl,
+            };
+      },
+    };
+  };
+
+  const webuiHostUrl =
+    localStorage.getItem("webui-host-url") ?? "http://127.0.0.1:7861";
+
+  return {
+    ...webuiLoad(webuiHostUrl),
+
+    getStableDiffusionDefaultCount: () => 4,
+
+    getStableDiffusionDefaultInput: () => {
+      return {
+        steps: 20,
+        sampler: {
+          id: localStorage.getItem("webui-saved-sampler") ?? "",
+          name: localStorage.getItem("webui-saved-sampler") ?? "",
+        },
+        model: localStorage.getItem("webui-saved-model") ?? "",
+      };
+    },
+
+    getStableDiffusionSamplers: async () => {
+      const response = await fetch(`${webuiHostUrl}/sdapi/v1/samplers`);
+      const responseData = await response.json();
+
+      return responseData.map((sampler: any) => ({
+        id: sampler.name,
+        name: sampler.name,
+      }));
+    },
+
+    getStableDiffusionExistingImages: async () => {
+      const existingImagesResponse = await fetch(
+        `${webuiHostUrl}/StableStudio/get-generated-images`,
+        {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+          },
+          body: JSON.stringify({
+            limit: get().settings.historyImagesCount.value,
+          }),
+        }
+      );
+
+      if (!existingImagesResponse.ok) {
+        console.warn("unable to get existing data from webui");
+      }
+
+      const responseData = await existingImagesResponse.json();
+
+      const images = [];
+
+      for (let i = 0; i < responseData.length; i++) {
+        const imageInfo = await getImageInfo(
+          webuiHostUrl,
+          responseData[i].content
+        );
+
+        const blob = await base64ToBlob(responseData[i].content, "image/jpeg");
+
+        const timestampInSeconds = responseData[i].create_date;
+        const timestampInMilliseconds = timestampInSeconds * 1000;
+        const createdAt = new Date(timestampInMilliseconds);
+
+        const stableDiffusionImage = {
+          id: responseData[i].image_name,
+          createdAt: createdAt,
+          blob: blob,
+          input: {
+            prompts: [
+              {
+                text: imageInfo["prompt"],
+                weight: imageInfo["CFG scale"],
+              },
+            ],
+            style: "",
+            steps: Number(imageInfo["Steps"]) ?? -1,
+            seed: Number(imageInfo["Seed"]) ?? -1,
+            model: imageInfo["Model"] ?? "",
+            width: responseData[i].width,
+            height: responseData[i].height,
+          },
+        };
+
+        images.push(stableDiffusionImage);
+      }
+
+      return [
+        {
+          id: `${Math.random() * 10000000}`,
+          images: images,
+        },
+      ];
+    },
+
+    settings: {
+      baseUrl: {
+        type: "string",
+        title: "Host URL",
+        placeholder: "http://127.0.0.1:7861",
+        value: localStorage.getItem("webui-host-url") ?? "",
+        description:
+          "The URL of the `stable-diffusion-webui` host, usually http://127.0.0.1:7861",
+      },
+
+      upscaler: {
+        type: "string",
+        title: "Upscaler 1",
+        options: webuiUpscalers,
+        value: localStorage.getItem("upscaler1") ?? webuiUpscalers[0].value,
+        description:
+          "Select the upscaler used when downloading images at more than 1x size",
+      },
+
+      historyImagesCount: {
+        type: "number",
+        title: "History image count",
+        description: "How many images should be fetched from local history?",
+        min: 0,
+        max: 50,
+        step: 1,
+        variant: "slider",
+        value: getNumber(localStorage.getItem("historyImagesCount"), 20),
+      },
+    },
+
+    setSetting: (key, value) => {
+      set(({ settings }) => ({
+        settings: {
+          ...settings,
+          [key]: { ...settings[key], value: value as string },
+        },
+      }));
+
+      if (key === "baseUrl" && typeof value === "string") {
+        localStorage.setItem("webui-host-url", value);
+        set((plugin) => ({ ...plugin, ...webuiLoad(value) }));
+      } else if (key === "upscaler" && typeof value === "string") {
+        localStorage.setItem("upscaler1", value);
+      } else if (key === "historyImagesCount" && typeof value === "number") {
+        localStorage.setItem("historyImagesCount", value.toString());
+      }
+    },
+
+    manifest,
+  };
+});

packages/stablestudio-plugin/README.md

@@ -20,7 +20,25 @@ In order to make StableStudio easier to extend, we've ripped out the "back-end"
 
 This means you can implement an entirely different inference stack, StableStudio doesn't care if it's local or a hosted API.
 
-## ⭐️ Features
+## ⭐️ First-Party Plugins
+
+There are currently three first-party plugins which are maintained in this repository:
+
+- [`stablestudio-plugin-stability`](../stablestudio-plugin-stability/src/index.ts) – The default plugin which uses [Stability's API](https://platform.stability.ai) for inference.
+
+  ```bash
+  yarn dev
+  ```
+
+- [`stablestudio-plugin-webui`](../stablestudio-plugin-webui/README.md) – This plugin uses [`stable-diffusion-webui`](https://github.com/AUTOMATIC1111/stable-diffusion-webui) for inference.
+
+  ```bash
+  yarn dev:use-webui-plugin
+  ```
+
+We are still figuring out a more scalable strategy for third-party plugins, [let us know what you think](https://github.com/Stability-AI/StableStudio/issues/3)!
+
+## ⚡️ Features
 
 We're hoping this list expands [over time](#future), but here's what's available right now...
 

packages/stablestudio-ui/src/Environment/index.tsx

@@ -8,6 +8,7 @@ declare global {
   interface ImportMetaEnv {
     readonly VITE_GIT_HASH: string;
     readonly VITE_USE_EXAMPLE_PLUGIN: string;
+    readonly VITE_USE_WEBUI_PLUGIN: string;
   }
 }
 
@@ -20,6 +21,7 @@ export namespace Environment {
   const variables = {
     VITE_GIT_HASH: import.meta.env.VITE_GIT_HASH,
     VITE_USE_EXAMPLE_PLUGIN: import.meta.env.VITE_USE_EXAMPLE_PLUGIN ?? "false",
+    VITE_USE_WEBUI_PLUGIN: import.meta.env.VITE_USE_WEBUI_PLUGIN ?? "false",
   } as const;
 
   export function get(name: VariableName): string {

packages/stablestudio-ui/src/Plugin/index.tsx

@@ -1,6 +1,7 @@
 import * as StableStudio from "@stability/stablestudio-plugin";
 import * as StableStudioPluginExample from "@stability/stablestudio-plugin-example";
 import * as StableStudioPluginStability from "@stability/stablestudio-plugin-stability";
+import * as StableStudioPluginWebUI from "@stability/stablestudio-plugin-webui";
 
 import { Environment } from "~/Environment";
 import { Generation } from "~/Generation";
@@ -115,6 +116,8 @@ namespace State {
     const { createPlugin: createRootPlugin } =
       Environment.get("USE_EXAMPLE_PLUGIN") === "true"
         ? StableStudioPluginExample
+        : Environment.get("USE_WEBUI_PLUGIN") === "true"
+        ? StableStudioPluginWebUI
         : StableStudioPluginStability;
 
     return {