Make samplers promisable and ID-based to support arbitrary schemes outside of gRPC

Author: cruhl

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

@@ -26,9 +26,9 @@ const getStableDiffusionDefaultInputFromPrompt = (prompt: string) => ({
   ],
 
   model: "stable-diffusion-xl-beta-v2-2-2",
+  sampler: { id: "0", name: "DDIM" },
   style: "enhance",
 
-  sampler: { value: 0, name: "DDIM" },
   width: 512,
   height: 512,
 
@@ -147,6 +147,8 @@ export const createPlugin = StableStudio.createPlugin<{
           );
         }
 
+        console.log(input.sampler.id ? parseInt(input.sampler.id) : 0);
+
         const imageParams = Generation.ImageParameters.create({
           width: BigInt(width),
           height: BigInt(height),
@@ -158,7 +160,7 @@ export const createPlugin = StableStudio.createPlugin<{
           transform: Generation.TransformType.create({
             type: {
               oneofKind: "diffusion",
-              diffusion: input.sampler.value ?? 0,
+              diffusion: input.sampler.id ? parseInt(input.sampler.id) : 0,
             },
           }),
 
@@ -411,55 +413,6 @@ export const createPlugin = StableStudio.createPlugin<{
         );
       },
 
-      getStableDiffusionSamplers: () => {
-        return [
-          {
-            value: 0,
-            name: "DDIM",
-          },
-          {
-            value: 1,
-            name: "DDPM",
-          },
-          {
-            value: 2,
-            name: "K_EULER",
-          },
-          {
-            value: 3,
-            name: "K_EULER_ANCESTRAL",
-          },
-          {
-            value: 4,
-            name: "K_HEUN",
-          },
-          {
-            value: 5,
-            name: "K_DPM_2",
-          },
-          {
-            value: 6,
-            name: "K_DPM_2_ANCESTRAL",
-          },
-          {
-            value: 7,
-            name: "K_LMS",
-          },
-          {
-            value: 8,
-            name: "K_DPMPP_2S_ANCESTRAL",
-          },
-          {
-            value: 9,
-            name: "K_DPMPP_2M",
-          },
-          {
-            value: 10,
-            name: "K_DPMPP_SDE",
-          },
-        ];
-      },
-
       getStableDiffusionModels: async () => {
         const request = await engines.listEngines({});
         const allEngines = await request.response.engine;
@@ -484,6 +437,20 @@ export const createPlugin = StableStudio.createPlugin<{
       localStorage.getItem("stability-apiKey") ?? undefined
     ),
 
+    getStableDiffusionSamplers: () => [
+      { id: "0", name: "DDIM" },
+      { id: "1", name: "DDPM" },
+      { id: "2", name: "K_EULER" },
+      { id: "3", name: "K_EULER_ANCESTRAL" },
+      { id: "4", name: "K_HEUN" },
+      { id: "5", name: "K_DPM_2" },
+      { id: "6", name: "K_DPM_2_ANCESTRAL" },
+      { id: "7", name: "K_LMS" },
+      { id: "8", name: "K_DPMPP_2S_ANCESTRAL" },
+      { id: "9", name: "K_DPMPP_2M" },
+      { id: "10", name: "K_DPMPP_SDE" },
+    ],
+
     getStableDiffusionStyles: () => [
       {
         id: "enhance",

packages/stablestudio-plugin/src/Plugin.ts

@@ -181,13 +181,16 @@ export type Plugin<P extends PluginTypeHelper = PluginTypeHelperDefault> = {
     StableDiffusionModel[] | undefined
   >;
 
+  /** If more than one `StableDiffusionSampler`, you can return them via this function and they will be presented as a dropdown in the UI */
+  getStableDiffusionSamplers?: () => MaybePromise<
+    StableDiffusionSampler[] | undefined
+  >;
+
   /** If more than one `StableDiffusionStyle`, you can return them via this function and they will be presented as a dropdown in the UI */
   getStableDiffusionStyles?: () => MaybePromise<
     StableDiffusionStyle[] | undefined
   >;
 
-  getStableDiffusionSamplers?: () => 
-    StableDiffusionSampler[] | undefined;
   /** Determines the default count passed to `createStableDiffusionImages` */
   getStableDiffusionDefaultCount?: () => number | undefined;
 
@@ -243,10 +246,6 @@ export type StableDiffusionPrompt = {
   weight?: number;
 };
 
-export declare type StableDiffusionSampler = {
-    name?: string;
-    value?: number;
-}
 /** This controls how a `StableDiffusionModel` is displayed in the UI */
 export type StableDiffusionModel = {
   id: ID;
@@ -263,6 +262,12 @@ export type StableDiffusionStyle = {
   image?: URLString;
 };
 
+/** This controls how a `StableDiffusionSampler` is displayed in the UI */
+export declare type StableDiffusionSampler = {
+  id: ID;
+  name?: string;
+};
+
 export type StableDiffusionInputImage = {
   blob?: Blob;
 

packages/stablestudio-ui/src/Generation/Image/Model/Dropdown.tsx

@@ -55,8 +55,8 @@ export function Dropdown({ id, className }: Styleable & { id: ID }) {
       options={options}
       anchor="bottom"
     >
-      {models === undefined && (
-        <div className="flex flex-col items-center justify-center py-32 px-16">
+      {!models && (
+        <div className="flex flex-col items-center justify-center px-16 py-32">
           <div className="text-muted-white pb-3">Loading models...</div>
         </div>
       )}

packages/stablestudio-ui/src/Generation/Image/Sampler/Dropdown.tsx

@@ -3,41 +3,40 @@ import { Theme } from "~/Theme";
 
 export function Dropdown({ id, className }: Styleable & { id: ID }) {
   const { setInput, input } = Generation.Image.Input.use(id);
-  const { samplers } = Generation.Image.Sampler.use();
+  const { data: samplers } = Generation.Image.Samplers.use();
 
   const options = useMemo(
     () => [
-      ...(samplers ?? []).map(({ value, name }) => ({
+      ...(samplers ?? []).map(({ id, name }) => ({
         name: name,
-        value: value,
+        value: id,
       })),
     ],
     [samplers]
   );
 
   const onClick = useCallback(
-    (value: number) => {
+    (id: ID) => {
       setInput((input) => {
-        console.log("sampler", value);
-        input.sampler = { value: value, name: options.at(value)?.name };
+        input.sampler = samplers?.find((sampler) => sampler.id === id);
       });
     },
-    [setInput, options]
+    [setInput, samplers]
   );
 
   if (!input) return null;
   return (
     <Theme.Popout
       title="Sampler"
       label="Sampler"
-      placeholder={"Select a Sampler"}
-      value={input.sampler?.value}
+      placeholder={"Select a sampler"}
+      value={input.sampler?.id}
       className={className}
       onClick={onClick}
       options={options}
       anchor="bottom"
     >
-      {samplers === undefined && (
+      {!samplers && (
         <div className="flex flex-col items-center justify-center px-16 py-32">
           <div className="text-muted-white pb-3">Loading samplers...</div>
         </div>

packages/stablestudio-ui/src/Generation/Image/Sampler/Samplers.tsx

@@ -0,0 +1,23 @@
+import * as StableStudio from "@stability/stablestudio-plugin";
+import * as ReactQuery from "@tanstack/react-query";
+
+import { Plugin } from "~/Plugin";
+
+export type Samplers = StableStudio.StableDiffusionSampler[];
+export namespace Samplers {
+  export const use = () => {
+    const getStableDiffusionSamplers = Plugin.use(
+      ({ getStableDiffusionSamplers }) => getStableDiffusionSamplers
+    );
+
+    return ReactQuery.useQuery({
+      enabled: !!getStableDiffusionSamplers,
+
+      queryKey: ["Generation.Image.Samplers.use"],
+      queryFn: async () => (await getStableDiffusionSamplers?.()) ?? [],
+    });
+  };
+
+  export const useAreEnabled = () =>
+    Plugin.use(({ getStableDiffusionModels }) => !!getStableDiffusionModels);
+}

packages/stablestudio-ui/src/Generation/Image/Sampler/index.tsx

@@ -1,32 +1,15 @@
 import * as StableStudio from "@stability/stablestudio-plugin";
 
-import { Plugin } from "~/Plugin";
-
 import { Dropdown } from "./Dropdown";
 
+export * from "./Samplers";
+
 export type Sampler = StableStudio.StableDiffusionSampler[];
 
 export declare namespace Sampler {
   export { Dropdown };
 }
 
 export namespace Sampler {
-  export const use = () => {
-    const getStableDiffusionSamplers = Plugin.use(
-      ({ getStableDiffusionSamplers }) => getStableDiffusionSamplers
-    );
-
-    const samplers = useMemo(
-      () => getStableDiffusionSamplers?.(),
-      [getStableDiffusionSamplers]
-    );
-    return { samplers };
-  };
-
-  export const useAreEnabled = () =>
-    Plugin.use(
-      ({ getStableDiffusionSamplers }) => !!getStableDiffusionSamplers
-    );
-
   Sampler.Dropdown = Dropdown;
 }

packages/stablestudio-ui/src/Generation/Image/Sidebar/Advanced.tsx

@@ -8,7 +8,7 @@ export function Advanced({
 }: App.Sidebar.Section.Props & { id: ID }) {
   const { setInput, input } = Generation.Image.Input.use(id);
   const areModelsEnabled = Generation.Image.Models.useAreEnabled();
-  const areSamplersEnabled = Generation.Image.Sampler.useAreEnabled();
+  const areSamplersEnabled = Generation.Image.Samplers.useAreEnabled();
 
   const onPromptStrengthChange = useCallback(
     (cfgScale: number) => {

packages/stablestudio-ui/src/Generation/Image/index.tsx

@@ -14,7 +14,7 @@ import { Modal } from "./Modal";
 import { Model, Models } from "./Model";
 import { Output, Outputs } from "./Output";
 import { Prompt, Prompts } from "./Prompt";
-import { Sampler } from "./Sampler";
+import { Sampler, Samplers } from "./Sampler";
 import { Search } from "./Search";
 import { Session } from "./Session";
 import { Sidebar } from "./Sidebar";
@@ -210,6 +210,7 @@ export declare namespace Image {
     Prompt,
     Prompts,
     Sampler,
+    Samplers,
     Search,
     Session,
     Sidebar,
@@ -242,6 +243,7 @@ export namespace Image {
   Image.Prompt = Prompt;
   Image.Prompts = Prompts;
   Image.Sampler = Sampler;
+  Image.Samplers = Samplers;
   Image.Search = Search;
   Image.Session = Session;
   Image.Sidebar = Sidebar;

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

@@ -1,12 +1,10 @@
-// import { App } from "~/App";
 import { Shortcut } from "~/Shortcut";
 import { Theme } from "~/Theme";
 
 import { Image, Images } from "./Image";
 
 export function Generation() {
   const createDream = Generation.Image.Session.useCreateDream();
-  // App.Breadcrumbs.useSet(["defaultProject", "generate"]);
 
   Shortcut.use(
     useMemo(