progress, handle cached nodes

Stability-AI · Jul 23, 2023 · 0065f05 · 0065f05
1 parent d44d5d1
commit 0065f05
Show file tree

Hide file tree

Showing 10 changed files with 144 additions and 17 deletions.
diff --git a/packages/stablestudio-ui/package.json b/packages/stablestudio-ui/package.json
@@ -42,6 +42,7 @@
     "notistack": "^3.0.0-alpha.11",
     "query-string": "^8.1.0",
     "react": "^18.2.0",
+    "react-circular-progressbar": "^2.1.0",
     "react-dom": "^18.2.0",
     "react-konva": "^18.2.3",
     "react-konva-utils": "^0.3.1",

diff --git a/packages/stablestudio-ui/src-tauri/src/server.rs b/packages/stablestudio-ui/src-tauri/src/server.rs
@@ -78,7 +78,7 @@ impl Builder {
                                     Some("text/html".to_string()),
                                     None,
                                 ),
-                                ["api" | "prompt" | "object_info" | "view" | "history" | "queue" | "interrupt" | "extensions", ..] => (
+                                ["api" | "prompt" | "object_info" | "view" | "upload" |  "history" | "queue" | "interrupt" | "extensions", ..] => (
                                     None,
                                     None,
                                     Some(

diff --git a/packages/stablestudio-ui/src/Comfy/index.tsx b/packages/stablestudio-ui/src/Comfy/index.tsx
@@ -116,6 +116,12 @@ type State = {
 
   unlisteners: (() => void)[];
   setUnlisteners: (unlisteners: (() => void)[]) => void;
+
+  runningPrompt: ID | null;
+  setRunningPrompt: (promptID: ID | null) => void;
+
+  lastOuput: ComfyOutput | null;
+  setLastOutput: (output: ComfyOutput | null) => void;
 };
 
 export namespace Comfy {
@@ -142,6 +148,12 @@ export namespace Comfy {
 
     unlisteners: [],
     setUnlisteners: (unlisteners) => set({ unlisteners }),
+
+    runningPrompt: null,
+    setRunningPrompt: (runningPrompt) => set({ runningPrompt }),
+
+    lastOuput: null,
+    setLastOutput: (lastOuput) => set({ lastOuput }),
   }));
 
   export const registerListeners = async () => {
@@ -152,10 +164,26 @@ export namespace Comfy {
       api = get()?.api;
     }
 
-    api.addEventListener("executed", async ({ detail }) => {
+    api.addEventListener("progress", ({ detail }) => {
+      console.log("progress", detail);
+      const runningPrompt = use.getState().runningPrompt;
+      if (runningPrompt) {
+        Generation.Image.Output.set({
+          ...Generation.Image.Output.get(runningPrompt),
+          progress: detail.value / detail.max,
+        });
+      }
+    });
+
+    api.addEventListener("b_preview", ({ detail }) => {
+      console.log("b_preview", detail);
+    });
+
+    const executed = async ({ detail }: any) => {
       const { output, prompt_id } = detail;
 
       console.log("executed_in_comfy_domain", detail);
+      use.getState().setRunningPrompt(null);
 
       const newInputs: Record<ID, Generation.Image.Input> = {};
       const responses: Generation.Images = [];
@@ -194,7 +222,7 @@ export namespace Comfy {
         const newInput = {
           ...Generation.Image.Input.initial(inputID),
           ...input,
-          seed: 0,
+          seed: (input?.seed ?? 0) + images.indexOf(image),
           id: inputID,
         };
 
@@ -211,11 +239,32 @@ export namespace Comfy {
       });
       responses.forEach(Generation.Image.add);
       Generation.Image.Output.received(prompt_id, responses);
+      use.getState().setLastOutput(detail);
+    };
+
+    api.addEventListener("executed", executed);
+    api.addEventListener("execution_cached", async ({ detail }) => {
+      const last: any = use.getState().lastOuput;
+      console.log("execution_cached", detail, last);
+      if (
+        use.getState().runningPrompt === detail.prompt_id &&
+        detail.nodes.includes(last?.node) &&
+        last.output
+      ) {
+        console.log("last", last);
+        const d = { ...last, prompt_id: detail.prompt_id };
+        await executed({ detail: d });
+      }
     });
 
     api.addEventListener("execution_error", ({ detail }) => {
       console.log("execution_error", detail);
       Generation.Image.Output.clear(detail.prompt_id);
+      use.getState().setRunningPrompt(null);
+    });
+
+    api.addEventListener("execution_start", ({ detail }) => {
+      use.getState().setRunningPrompt(detail?.prompt_id);
     });
 
     console.log("registered ComfyUI listeners");

diff --git a/packages/stablestudio-ui/src/Generation/Image/Create/index.tsx b/packages/stablestudio-ui/src/Generation/Image/Create/index.tsx
@@ -57,6 +57,9 @@ export namespace Create {
         pluginInput.width = Math.ceil((pluginInput.width ?? 512) / 64) * 64;
       }
 
+      const startingSeed =
+        input.seed === 0 ? Math.round(Math.random() * 100000) : input.seed;
+
       Comfy.get()
         ?.graph._nodes?.filter((node) => node.type === "StableStudioNode")
         .forEach((node) => {
@@ -73,10 +76,7 @@ export namespace Create {
             positive_prompt:
               input.prompts.find((p) => p.weight > 0)?.text ??
               node.stableValues.positive_prompt,
-            seed:
-              input.seed === 0
-                ? Math.round(Math.random() * 100000)
-                : input.seed,
+            seed: startingSeed,
             steps: input.steps,
             cfg: input.cfgScale ?? node.stableValues.cfg,
             sampler_name:
@@ -97,6 +97,7 @@ export namespace Create {
           ...inputs,
           [prompt_id]: {
             ...input,
+            seed: startingSeed,
             id: prompt_id,
           },
         }));

diff --git a/packages/stablestudio-ui/src/Generation/Image/Model/Dropdown.tsx b/packages/stablestudio-ui/src/Generation/Image/Model/Dropdown.tsx
@@ -1,18 +1,32 @@
+import { useLocalStorage } from "react-use";
 import { Generation } from "~/Generation";
 import { Theme } from "~/Theme";
 
 export function Dropdown({ id, className }: Styleable & { id: ID }) {
   const { setInput, input } = Generation.Image.Input.use(id);
   const { data: models, isLoading } = Generation.Image.Models.use();
 
+  const [value, setValue] = useLocalStorage<string | undefined>(
+    "default-model-id",
+    undefined
+  );
+  useEffect(() => {
+    if (value) {
+      setInput((input) => {
+        input.model = value;
+      });
+    }
+  }, []);
+
   const onClick = useCallback(
     (value: string) => {
       setInput((input) => {
         console.log("model", value);
         input.model = value;
+        setValue(value);
       });
     },
-    [setInput]
+    [setInput, setValue]
   );
 
   const options = useMemo(

diff --git a/packages/stablestudio-ui/src/Generation/Image/Output/index.tsx b/packages/stablestudio-ui/src/Generation/Image/Output/index.tsx
@@ -13,9 +13,11 @@ export type Output = {
 
   requestedAt?: Date;
   completedAt?: Date;
+  progress?: number;
 
   count: number;
   imageIDs: ID[];
+  progressImageIDs?: ID[];
 
   exception?: Generation.Image.Exception;
 };
@@ -61,6 +63,7 @@ export function Output({ outputID, placeholder, divider }: Props) {
         key={keys("image", images.length, images.length - index)}
         placeholder={placeholder}
         image={image}
+        progress={output?.progress}
         scale={1}
         example={
           Generation.Image.Prompt.Examples.images[
@@ -76,7 +79,14 @@ export function Output({ outputID, placeholder, divider }: Props) {
         {rendered}
       </div>
     );
-  }, [count, images, input?.id, placeholder, exampleStartIndex]);
+  }, [
+    count,
+    images,
+    input?.id,
+    placeholder,
+    output?.progress,
+    exampleStartIndex,
+  ]);
 
   const controls = useMemo(
     () => (

diff --git a/packages/stablestudio-ui/src/Generation/Image/SpecialEffects/index.tsx b/packages/stablestudio-ui/src/Generation/Image/SpecialEffects/index.tsx
@@ -1,5 +1,7 @@
+import { CircularProgressbar } from "react-circular-progressbar";
 import { Generation } from "~/Generation";
 import { Theme } from "~/Theme";
+import "react-circular-progressbar/dist/styles.css";
 
 import { Filter } from "./Filter";
 
@@ -10,6 +12,7 @@ export function SpecialEffects({
   example,
   onClick,
   input,
+  progress,
 }: {
   showing?: boolean;
   loading?: boolean;
@@ -21,6 +24,7 @@ export function SpecialEffects({
   onClick?: () => void;
   input?: ID;
   border?: boolean;
+  progress?: number;
 }) {
   if (!showing) return null;
   return (
@@ -74,13 +78,47 @@ export function SpecialEffects({
         </div>
       )}
       <div className="absolute flex h-full w-full items-center justify-center">
-        {loading && (
-          <Theme.Loading.Spinner
-            className={classes(
-              variant === "small" ? "h-1/2 w-1/2" : "h-10 w-10"
-            )}
-          />
-        )}
+        {loading &&
+          (typeof progress === "number" ? (
+            <div
+              className={classes(
+                variant === "small" ? "h-1/2 w-1/2" : "h-10 w-10"
+              )}
+            >
+              <CircularProgressbar
+                value={progress}
+                maxValue={1}
+                styles={{
+                  // Customize the path, i.e. the "completed progress"
+                  path: {
+                    // Path color
+                    stroke: "#ffffff",
+                    // Whether to use rounded or flat corners on the ends - can use 'butt' or 'round'
+                    strokeLinecap: "butt",
+                    // Customize transition animation
+                    transition: "stroke-dashoffset 0.5s ease 0s",
+                    transformOrigin: "center center",
+                    strokeWidth: 8,
+                  },
+                  // Customize the circle behind the path, i.e. the "total progress"
+                  trail: {
+                    // Trail color
+                    stroke: "#4c4c4d",
+                    // Whether to use rounded or flat corners on the ends - can use 'butt' or 'round'
+                    strokeLinecap: "butt",
+                    transformOrigin: "center center",
+                    strokeWidth: 8,
+                  },
+                }}
+              />
+            </div>
+          ) : (
+            <Theme.Loading.Spinner
+              className={classes(
+                variant === "small" ? "h-1/2 w-1/2" : "h-10 w-10"
+              )}
+            />
+          ))}
       </div>
     </div>
   );

diff --git a/packages/stablestudio-ui/src/Generation/Image/index.tsx b/packages/stablestudio-ui/src/Generation/Image/index.tsx
@@ -41,6 +41,7 @@ type Props = Styleable &
     example?: Generation.Image.Prompt.Examples.Example;
     hideControls?: boolean;
     placeholder?: boolean;
+    progress?: number;
 
     onClick?: () => void;
     onDelete?: () => void;
@@ -51,6 +52,7 @@ export function Image({
   // example,
   hideControls,
   placeholder,
+  progress,
 
   scale,
   preserveAspectRatio,
@@ -155,13 +157,14 @@ export function Image({
       <Image.SpecialEffects
         showing={shouldShowSpecialEffects}
         loading={!placeholder && shouldShowSpecialEffects}
+        progress={progress}
         variant={(style.height ?? 512) < 48 ? "small" : undefined}
         // example={example}
         // onClick={example ? onTryTemplate : undefined}
         // input={currentInput?.id}
       />
     ),
-    [placeholder, shouldShowSpecialEffects, style.height]
+    [placeholder, progress, shouldShowSpecialEffects, style.height]
   );
 
   return useMemo(

diff --git a/packages/stablestudio-ui/vite.config.ts b/packages/stablestudio-ui/vite.config.ts
@@ -48,6 +48,7 @@ export default defineConfig(({ mode }) => {
         "/queue": redirectComfy,
         "/history": redirectComfy,
         "/interrupt": redirectComfy,
+        "/upload": redirectComfy,
       },
     },
 

diff --git a/yarn.lock b/yarn.lock
@@ -1263,6 +1263,7 @@ __metadata:
     prettier-plugin-tailwindcss: ^0.2.1
     query-string: ^8.1.0
     react: ^18.2.0
+    react-circular-progressbar: ^2.1.0
     react-dom: ^18.2.0
     react-konva: ^18.2.3
     react-konva-utils: ^0.3.1
@@ -6322,6 +6323,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"react-circular-progressbar@npm:^2.1.0":
+  version: 2.1.0
+  resolution: "react-circular-progressbar@npm:2.1.0"
+  peerDependencies:
+    react: ^0.14.0 || ^15.0.0 || ^16.0.0 || ^17.0.0 || ^18.0.0
+  checksum: dc118010a8f94733daafac586c969f7e889ad0736d96a0fda79406ee62e9410848abfd3dee4887bb0ff46b99e1f8c86ee1afabc88dfbf4dcb95107b22de1d6d7
+  languageName: node
+  linkType: hard
+
 "react-dom@npm:^18.2.0":
   version: 18.2.0
   resolution: "react-dom@npm:18.2.0"