better stdout, model and sampler options

Stability-AI · Jul 20, 2023 · 4d435a4 · 4d435a4
1 parent f4a9d4e
commit 4d435a4
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 63 deletions.
diff --git a/packages/stablestudio-ui/src-tauri/src/main.rs b/packages/stablestudio-ui/src-tauri/src/main.rs
@@ -2,6 +2,7 @@
 #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
 
 use std::collections::HashMap;
+use std::fmt::format;
 use std::fs::File;
 use std::sync::OnceLock;
 use tauri::api::process::CommandEvent;
@@ -97,7 +98,7 @@ fn emit_event(event: &str, data: Option<String>) {
                 .unwrap_or_else(|_| println!("[!!] event failed"));
         }
         None => {
-            println!("[!!] window locked");
+            println!("[!!] window not set");
         }
     }
 }
@@ -109,20 +110,21 @@ async fn watch_comfy(
         // check if output starts with "To see the GUI go to:"
         match i {
             Stdout(line) if line.len() > 1 => {
+                println!("[ComfyUI] stdout: {}", line);
+                emit_event("comfy-output", Some(format!("stdout:{line}")));
+
                 if line.starts_with("To see the GUI go to:") {
                     println!("Comfy launched successfully!");
                     return (rx, Ok("completed".to_string()));
                 }
-                println!("[ComfyUI] stdout: {}", line);
-                emit_event("comfy-stdout", Some(line));
             }
             Stderr(line) => {
                 println!("[ComfyUI] stderr: {}", line);
-                emit_event("comfy-stderr", Some(line));
+                emit_event("comfy-output", Some(format!("stderr:{line}")));
             }
             Error(line) => {
                 println!("[ComfyUI] error: {}", line);
-                emit_event("comfy-error", Some(line));
+                emit_event("comfy-output", Some(format!("error:{line}")));
             }
             Terminated(_) => {
                 println!("Comfy terminated!");

diff --git a/packages/stablestudio-ui/src/App/index.tsx b/packages/stablestudio-ui/src/App/index.tsx
@@ -1,3 +1,4 @@
+import { listen } from "@tauri-apps/api/event";
 import {
   BaseDirectory,
   exists,
@@ -8,6 +9,7 @@ import {
 import { appDataDir } from "@tauri-apps/api/path";
 import { invoke } from "@tauri-apps/api/tauri";
 import { download } from "tauri-plugin-upload";
+import { shallow } from "zustand/shallow";
 import { Comfy } from "~/Comfy";
 import { Router } from "~/Router";
 import { Shortcut } from "~/Shortcut";
@@ -105,6 +107,10 @@ export namespace App {
     const [isSetup, setIsSetup] = useState<SetupState>(SetupState.NotStarted);
     const [message, setMessage] = useState<string>("");
     const [progress, setProgress] = useState<number | null>(null);
+    const [setRunning, setUnlisteners, print] = Comfy.use(
+      (state) => [state.setRunning, state.setUnlisteners, state.print],
+      shallow
+    );
     const nonce = useRef<number>(0);
 
     const check = useCallback(async () => {
@@ -196,6 +202,14 @@ export namespace App {
       setMessage("Starting ComfyUI...");
       setProgress(null);
 
+      // add listener
+      const unlisten = await listen("comfy-output", (event) => {
+        const [t, ...d] = `${event.payload}`.split(":");
+        console.log("[COMFYUI]", t, d);
+        print(t, d.join(":"));
+      });
+      setUnlisteners([unlisten]);
+
       // start comfy
       const result = await invoke("launch_comfy", {
         path: `${appDataPath}/comfyui`,
@@ -206,8 +220,9 @@ export namespace App {
         throw new Error("Failed to launch ComfyUI");
       }
 
+      setRunning(true);
       setIsSetup(SetupState.ComfyRunning);
-    }, [isSetup]);
+    }, [isSetup, print, setRunning, setUnlisteners]);
 
     useEffect(() => {
       check();

diff --git a/packages/stablestudio-ui/src/Comfy/index.tsx b/packages/stablestudio-ui/src/Comfy/index.tsx
@@ -1,4 +1,3 @@
-import { listen } from "@tauri-apps/api/event";
 import { useLocation } from "react-router-dom";
 import { create } from "zustand";
 
@@ -60,37 +59,6 @@ export type Graph = {
 
 export function Comfy() {
   const location = useLocation();
-  const { print } = Comfy.use();
-
-  useEffect(() => {
-    let mounted = true;
-    const unlisteners = [] as (() => void)[];
-    (async () => {
-      if (mounted) {
-        unlisteners.push(
-          await listen("comfy-stdout", (event) => {
-            console.log("stdout", `${event.payload}`);
-            print("stdout", `${event.payload}`);
-          })
-        );
-        unlisteners.push(
-          await listen("comfy-stderr", (event) => {
-            console.log("stderr", `${event.payload}`);
-            print("stderr", `${event.payload}`);
-          })
-        );
-
-        mounted = false;
-      }
-    })();
-
-    return () => {
-      mounted = false;
-
-      for (const unlistener of unlisteners) unlistener();
-    };
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, []);
 
   return (
     <iframe
@@ -105,15 +73,21 @@ export function Comfy() {
   );
 }
 
-const MAX_STDOUT_LENGTH = 35;
+const MAX_STDOUT_LENGTH = 500;
 
 type State = {
   output: {
-    type: "stdout" | "stderr";
+    type: string;
     data: string;
   }[];
   max_lines: number;
-  print: (type: "stdout" | "stderr", data: string) => void;
+  print: (type: string, data: string) => void;
+
+  running: boolean;
+  setRunning: (running: boolean) => void;
+
+  unlisteners: (() => void)[];
+  setUnlisteners: (unlisteners: (() => void)[]) => void;
 };
 
 export namespace Comfy {
@@ -134,5 +108,11 @@ export namespace Comfy {
         if (output.length > MAX_STDOUT_LENGTH) output.shift();
         return { output };
       }),
+
+    running: false,
+    setRunning: (running) => set({ running }),
+
+    unlisteners: [],
+    setUnlisteners: (unlisteners) => set({ unlisteners }),
   }));
 }
diff --git a/packages/stablestudio-ui/src/Comfy/plugin.ts b/packages/stablestudio-ui/src/Comfy/plugin.ts
@@ -53,6 +53,32 @@ export const createPlugin = StableStudio.createPlugin<any>(({ set, get }) => ({
     };
   },
 
+  getStableDiffusionModels: async () => {
+    const resp = await fetch("/object_info", { cache: "no-cache" });
+    const jsonResp = await resp.json();
+
+    console.log(jsonResp);
+
+    return jsonResp?.CheckpointLoader?.input?.required?.ckpt_name?.[0]?.map(
+      (fileName: string) => ({
+        id: fileName,
+        name: fileName,
+      })
+    );
+  },
+
+  getStableDiffusionSamplers: async () => {
+    const resp = await fetch("/object_info", { cache: "no-cache" });
+    const jsonResp = await resp.json();
+
+    return jsonResp?.KSampler?.input?.required?.scheduler?.[0]?.map(
+      (name: string) => ({
+        id: name,
+        name,
+      })
+    );
+  },
+
   getStatus: () => {
     fetch("/comfyui", { cache: "no-cache" }).then((resp) => {
       set({

diff --git a/packages/stablestudio-ui/src/Settings/Manifest.tsx b/packages/stablestudio-ui/src/Settings/Manifest.tsx
@@ -6,14 +6,12 @@ import {
 import { shallow } from "zustand/shallow";
 import { Comfy } from "~/Comfy";
 
-import { Plugin } from "~/Plugin";
 import { Theme } from "~/Theme";
 
 import { Panel } from "./Panel";
 import { Setting } from "./Setting";
 
 export function Manifest({
-  id,
   manifest,
   pluginStatus,
   settings,
@@ -87,9 +85,18 @@ export function Manifest({
         </div>
         {output.length > 0 && (
           <>
-            <div className="whitespace-pre-wrap rounded bg-black/25 p-2 font-mono text-sm">
-              {output.map((line, index) => (
-                <p key={`${index}-${line}`}>{line.data}</p>
+            <div className="flex max-h-[25rem] flex-col-reverse overflow-y-auto whitespace-pre-wrap rounded bg-black/25 p-2 font-mono text-sm">
+              {[...output].reverse().map((line, index) => (
+                <p
+                  key={`${index}-${line}`}
+                  className={classes(
+                    "text-white",
+                    line.type === "stdout" && "text-green-200",
+                    line.type === "stderr" && "text-red-200"
+                  )}
+                >
+                  {line.data}
+                </p>
               ))}
             </div>
           </>
@@ -112,20 +119,3 @@ export function Manifest({
     </Panel>
   );
 }
-
-function MiniManifestField({
-  label,
-  value,
-  className,
-}: {
-  label: string;
-  value: string | undefined;
-  className?: string;
-}) {
-  return (
-    <p className={classes("opacity-muted text-sm", className)}>
-      <span className="select-none opacity-50">{label} </span>
-      {value}
-    </p>
-  );
-}