From e9657bb20f33d3cc62cea2c4ff2cb9937af8a7cf Mon Sep 17 00:00:00 2001 From: kaj <40004347+KAJdev@users.noreply.github.com> Date: Mon, 24 Jul 2023 17:01:14 -0800 Subject: [PATCH] workflows --- packages/stablestudio-ui/public/file-code.svg | 7 + .../stablestudio-ui/src-tauri/src/main.rs | 4 + .../src-tauri/src/workflows.rs | 94 +++++++++ packages/stablestudio-ui/src/Comfy/index.tsx | 4 +- .../src/Generation/Image/Sidebar/index.tsx | 13 +- .../src/Generation/Image/Workflow/Modal.tsx | 128 ++++++++++++ .../src/Generation/Image/Workflow/index.tsx | 188 ++++++++++++++++++ .../src/Generation/Image/index.tsx | 3 + .../stablestudio-ui/src/Theme/Icon/SVGs.tsx | 64 ++++++ .../stablestudio-ui/src/Theme/Icon/index.tsx | 6 + .../src/Theme/Popout/index.tsx | 106 ++++++---- 11 files changed, 571 insertions(+), 46 deletions(-) create mode 100644 packages/stablestudio-ui/public/file-code.svg create mode 100644 packages/stablestudio-ui/src-tauri/src/workflows.rs create mode 100644 packages/stablestudio-ui/src/Generation/Image/Workflow/Modal.tsx create mode 100644 packages/stablestudio-ui/src/Generation/Image/Workflow/index.tsx diff --git a/packages/stablestudio-ui/public/file-code.svg b/packages/stablestudio-ui/public/file-code.svg new file mode 100644 index 0000000..b911940 --- /dev/null +++ b/packages/stablestudio-ui/public/file-code.svg @@ -0,0 +1,7 @@ + + + + + + + diff --git a/packages/stablestudio-ui/src-tauri/src/main.rs b/packages/stablestudio-ui/src-tauri/src/main.rs index 8165fe0..017a02c 100644 --- a/packages/stablestudio-ui/src-tauri/src/main.rs +++ b/packages/stablestudio-ui/src-tauri/src/main.rs @@ -18,6 +18,7 @@ use tauri_plugin_upload; mod server; mod settings; mod show_path; +mod workflows; static WINDOW: OnceLock = OnceLock::new(); @@ -84,6 +85,9 @@ fn main() { show_path::show_in_folder, settings::get_setting, settings::set_setting, + workflows::fetch_workflows, + workflows::save_workflow, + workflows::delete_workflow, ]) .build(context) .expect("error while building tauri application") diff --git a/packages/stablestudio-ui/src-tauri/src/workflows.rs b/packages/stablestudio-ui/src-tauri/src/workflows.rs new file mode 100644 index 0000000..3776363 --- /dev/null +++ b/packages/stablestudio-ui/src-tauri/src/workflows.rs @@ -0,0 +1,94 @@ +use std::{ + fs, + io::{BufReader, BufWriter}, +}; + +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize)] +pub struct Workflow { + pub name: String, + pub description: Option, + pub icon: Option, + pub serialized_workflow: String, + pub created_at: String, +} + +#[tauri::command] +pub fn fetch_workflows(handle: tauri::AppHandle) -> Result, String> { + let mut workflows: Vec = Vec::new(); + let mut path = handle.path_resolver().app_data_dir().unwrap(); + path.push("workflows"); + + // create directory if it doesn't exist + std::fs::create_dir_all(&path).unwrap(); + + for entry in fs::read_dir(path).unwrap() { + let entry = entry.unwrap(); + let path = entry.path(); + let metadata = fs::metadata(&path).unwrap(); + if metadata.is_file() { + let workflow_file = fs::File::open(path).unwrap(); + let reader = BufReader::new(workflow_file); + if let Ok(workflow_data) = serde_json::from_reader(reader) { + workflows.push(workflow_data); + } + } + } + Ok(workflows) +} + +#[tauri::command] +pub fn save_workflow( + handle: tauri::AppHandle, + serialized_workflow: String, + name: String, + description: Option, + icon: Option, +) -> Result<(), String> { + let mut workflow_path = handle.path_resolver().app_data_dir().unwrap(); + workflow_path.push("workflows"); + + // create directory if it doesn't exist + std::fs::create_dir_all(&workflow_path).unwrap(); + + let file_name = name.replace(" ", "_").to_lowercase(); + workflow_path.push(format!("{file_name}.json")); + + let workflow_file = fs::File::create(workflow_path).unwrap(); + let writer = BufWriter::new(workflow_file); + + match serde_json::to_writer( + writer, + &Workflow { + name, + description, + icon, + serialized_workflow, + created_at: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + .to_string(), + }, + ) { + Ok(_) => Ok(()), + Err(e) => Err(e.to_string()), + } +} + +#[tauri::command] +pub fn delete_workflow(handle: tauri::AppHandle, name: String) -> Result<(), String> { + let mut workflow_path = handle.path_resolver().app_data_dir().unwrap(); + workflow_path.push("workflows"); + + // create directory if it doesn't exist + std::fs::create_dir_all(&workflow_path).unwrap(); + + workflow_path.push(format!("{name}.json")); + + match fs::remove_file(workflow_path) { + Ok(_) => Ok(()), + Err(e) => Err(e.to_string()), + } +} diff --git a/packages/stablestudio-ui/src/Comfy/index.tsx b/packages/stablestudio-ui/src/Comfy/index.tsx index 38aab55..1907f90 100644 --- a/packages/stablestudio-ui/src/Comfy/index.tsx +++ b/packages/stablestudio-ui/src/Comfy/index.tsx @@ -1,4 +1,3 @@ -import * as StableStudio from "@stability/stablestudio-plugin"; import { useLocation } from "react-router-dom"; import { create } from "zustand"; import { shallow } from "zustand/shallow"; @@ -7,7 +6,7 @@ import { Generation } from "~/Generation"; export type ComfyApp = { setup: () => void; registerNodes: () => void; - loadGraphData: (graph: Graph) => void; + loadGraphData: (graph?: Graph) => void; graphToPrompt: () => Promise<{ workflow: any; prompt: any; @@ -18,6 +17,7 @@ export type ComfyApp = { clean: () => void; api: ComfyAPI; graph: { + serialize: () => any; _nodes: { title: string; type: string; diff --git a/packages/stablestudio-ui/src/Generation/Image/Sidebar/index.tsx b/packages/stablestudio-ui/src/Generation/Image/Sidebar/index.tsx index 4767b9b..5cabc27 100644 --- a/packages/stablestudio-ui/src/Generation/Image/Sidebar/index.tsx +++ b/packages/stablestudio-ui/src/Generation/Image/Sidebar/index.tsx @@ -77,13 +77,12 @@ export namespace Sidebar { const areStylesEnabled = Generation.Image.Styles.useAreEnabled(); return ( <> - {areStylesEnabled && ( - -
- -
-
- )} + +
+ {areStylesEnabled && } + +
+
{variant === "generate" && ( diff --git a/packages/stablestudio-ui/src/Generation/Image/Workflow/Modal.tsx b/packages/stablestudio-ui/src/Generation/Image/Workflow/Modal.tsx new file mode 100644 index 0000000..42f4168 --- /dev/null +++ b/packages/stablestudio-ui/src/Generation/Image/Workflow/Modal.tsx @@ -0,0 +1,128 @@ +import { Generation } from "~/Generation"; +import { GlobalState } from "~/GlobalState"; +import { Theme } from "~/Theme"; + +export function Modal() { + const { + image, + open, + name, + description, + setImage, + setName, + setOpen, + setDescription, + } = Modal.State.use(); + const [saving, setSaving] = useState(false); + + return ( + setOpen(false)} + > + + setOpen(false)}> + + Save workflow + + +
+
+
+ Name + +
+
+ Description + +
+
+ Icon + { + const file = e.target.files?.[0]; + if (!file) return; + + const reader = new FileReader(); + reader.onload = (e) => { + const result = e.target?.result; + if (typeof result !== "string") return; + + setImage(result); + }; + reader.readAsDataURL(file); + }} + /> +
+
+ { + if (!image || !name) return; + + setSaving(true); + try { + await Generation.Image.Workflow.saveWorkflow( + name, + image, + description + ); + } finally { + setSaving(false); + } + setOpen(false); + }} + > + Save + +
+
+
+ ); +} + +export namespace Modal { + export type State = { + image?: string | null; + name: string; + open: boolean; + description: string; + + setImage: (image?: string) => void; + setName: (name: string) => void; + setOpen: (open: boolean) => void; + setDescription: (description: string) => void; + }; + + export namespace State { + export const use = GlobalState.create((set) => ({ + image: null, + name: "", + open: false, + description: "", + + setImage: (image) => set({ image }), + setName: (name) => set({ name }), + setOpen: (open) => set({ open, name: "", image: null, description: "" }), + setDescription: (description) => set({ description }), + })); + } +} diff --git a/packages/stablestudio-ui/src/Generation/Image/Workflow/index.tsx b/packages/stablestudio-ui/src/Generation/Image/Workflow/index.tsx new file mode 100644 index 0000000..5fa60ee --- /dev/null +++ b/packages/stablestudio-ui/src/Generation/Image/Workflow/index.tsx @@ -0,0 +1,188 @@ +import { invoke } from "@tauri-apps/api/tauri"; +import { create } from "zustand"; +import { Comfy } from "~/Comfy"; +import { Theme } from "~/Theme"; + +import { Modal } from "./Modal"; + +export type Workflow = { + name: string; + icon?: string; + description?: string; + serialized_workflow: string; + created_at: string; +}; + +export declare namespace Workflow { + export { Modal }; +} + +export namespace Workflow { + export const saveWorkflow = async ( + name: string, + icon: string, + description: string + ) => { + const comfyApp = Comfy.get(); + if (!comfyApp) return; + + const serializedWorkflow = JSON.stringify(comfyApp.graph.serialize(), null); + try { + await invoke("save_workflow", { + name, + icon, + serializedWorkflow, + description, + }); + } catch (error) { + console.error(error); + } + + console.log("saved! reloading workflows"); + + try { + Workflow.fetchWorkflows(); + } catch (error) { + console.error(error); + } + }; + + function SlightBox({ + onClick, + name, + icon, + }: { + onClick: (e: any) => void; + name: string; + icon: React.ReactNode; + }) { + return ( +
+
+ {icon} +
+

+ {name} +

+
+ ); + } + + export const fetchWorkflows = async () => { + const response = await invoke("fetch_workflows"); + Workflow.use.getState().setWorkflows(response as Workflow[]); + }; + + export const applyWorkflow = async (workflowId: string) => { + const workflow = use + .getState() + .workflows.find((workflow) => workflow.name === workflowId); + + if (!workflow) return; + + const comfyApp = Comfy.get(); + if (!comfyApp) return; + + // turn serialized_workflow string into a file + const blob = new Blob([workflow.serialized_workflow], { + type: "text/json", + }); + const file = new File([blob], "workflow.json", { type: "text/json" }); + + // load the file into comfy + comfyApp.handleFile(file); + }; + + export const use = create<{ + workflows: Workflow[]; + setWorkflows: (workflows: Workflow[]) => void; + }>((set) => ({ + workflows: [], + setWorkflows: (workflows) => set({ workflows }), + })); + + export function Dropdown({ className }: Styleable) { + const { workflows, setWorkflows } = Workflow.use(); + + useEffect(() => { + fetchWorkflows(); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []); + + const options = useMemo( + () => [ + { + value: " save", + image: "/save.svg", + component: (onClick: any) => ( + { + Modal.State.use.getState().setOpen(true); + onClick(e); + }} + name="Save workflow" + icon={} + /> + ), + }, + + ...(workflows ?? []) + .sort((a, b) => + new Date(a.created_at) > new Date(b.created_at) ? -1 : 1 + ) + .map((workflow: Workflow) => ({ + value: workflow.name, + name: workflow.name, + image: workflow.icon ?? "/file-code.svg", + onDelete: async () => { + await invoke("delete_workflow", { + name: workflow.name.replace(/ /g, "_").toLowerCase(), + }); + setWorkflows(workflows.filter((w) => w.name !== workflow.name)); + }, + })), + + { + value: " default", + image: "/rotate.svg", + component: (onClick: any) => ( + { + Comfy.get()?.loadGraphData(); + onClick(e); + }} + name="Load default" + icon={} + /> + ), + }, + ], + [setWorkflows, workflows] + ); + + return ( + <> + + + + ); + } +} diff --git a/packages/stablestudio-ui/src/Generation/Image/index.tsx b/packages/stablestudio-ui/src/Generation/Image/index.tsx index 5b71f45..e3656f8 100644 --- a/packages/stablestudio-ui/src/Generation/Image/index.tsx +++ b/packages/stablestudio-ui/src/Generation/Image/index.tsx @@ -24,6 +24,7 @@ import { Style, Styles } from "./Style"; import { TopBar } from "./TopBar"; import { Upscale, Upscales } from "./Upscale"; import { Variations } from "./Variation"; +import { Workflow } from "./Workflow"; export * from "./Images"; @@ -227,6 +228,7 @@ export declare namespace Image { Upscales, Upscale, Exception, + Workflow, }; } @@ -260,6 +262,7 @@ export namespace Image { Image.Upscales = Upscales; Image.Upscale = Upscale; Image.Exception = Exception; + Image.Workflow = Workflow; export const get = (id: ID): Image | undefined => Images.State.use(({ images }) => images[id]); diff --git a/packages/stablestudio-ui/src/Theme/Icon/SVGs.tsx b/packages/stablestudio-ui/src/Theme/Icon/SVGs.tsx index cbaa5a9..208365f 100644 --- a/packages/stablestudio-ui/src/Theme/Icon/SVGs.tsx +++ b/packages/stablestudio-ui/src/Theme/Icon/SVGs.tsx @@ -425,6 +425,70 @@ export function Upscale(props: Props) { )); } +export function Rotate(props: Props) { + return defaults(props)(({ width, height, color, ...props }) => ( + + + + + )); +} + +export function Save(props: Props) { + return defaults(props)(({ width, height, color, ...props }) => ( + + + + + + )); +} + const defaults = (props: Props) => (render: (props: Props) => JSX.Element): JSX.Element => diff --git a/packages/stablestudio-ui/src/Theme/Icon/index.tsx b/packages/stablestudio-ui/src/Theme/Icon/index.tsx index 1acce91..b735392 100644 --- a/packages/stablestudio-ui/src/Theme/Icon/index.tsx +++ b/packages/stablestudio-ui/src/Theme/Icon/index.tsx @@ -77,6 +77,8 @@ import { Instagram, ModelIcon, Rectangle, + Rotate, + Save, Scale, ShareIcon, SlidersIcon, @@ -169,6 +171,8 @@ export declare namespace Icon { Upscale, Keyboard, ChevronsLeftRight, + Rotate, + Save, }; } @@ -262,6 +266,8 @@ export namespace Icon { Icon.Upscale = makeComponent(Upscale); Icon.Keyboard = makeComponent(Keyboard); Icon.ChevronsLeftRight = makeComponent(ChevronsLeftRight); + Icon.Rotate = makeComponent(Rotate); + Icon.Save = makeComponent(Save); export function Invisible(props: Props) { return ( diff --git a/packages/stablestudio-ui/src/Theme/Popout/index.tsx b/packages/stablestudio-ui/src/Theme/Popout/index.tsx index c6480c2..113802a 100644 --- a/packages/stablestudio-ui/src/Theme/Popout/index.tsx +++ b/packages/stablestudio-ui/src/Theme/Popout/index.tsx @@ -5,11 +5,18 @@ import { Box } from "~/Geometry"; import { Theme } from "~/Theme"; export type Option = { - name: React.ReactNode; value: any; image?: string; disabled?: boolean; -}; + onDelete?: () => void; +} & ( + | { + name: React.ReactNode; + } + | { + component: (onClick: (value: any) => void) => React.ReactNode; + } +); export function Popout({ onClick, @@ -109,7 +116,11 @@ export function Popout({ /> )}

- {valueOption?.name ?? placeholder ?? "Select"} + {(valueOption && "name" in valueOption + ? valueOption?.name + : undefined) ?? + placeholder ?? + "Select"}

@@ -154,40 +165,61 @@ function Floating({ > {!!children ? children - : options.map((option, index) => ( -
onClick(option.value) - } - className={classes( - "group flex cursor-pointer flex-col rounded duration-100", - option.disabled && "opacity-muted cursor-not-allowed" - )} - > - {option.image ? ( - Preset Image - ) : ( - hasImages && ( -
- -
- ) - )} -

- {option.name} -

-
- ))} + : options.map( + (option, index) => + ("component" in option && option.component(onClick)) || ( +
onClick(option.value) + } + className={classes( + "group relative flex cursor-pointer flex-col rounded duration-100", + option.disabled && "opacity-muted cursor-not-allowed" + )} + > + {option.onDelete && ( + +
{ + e.stopPropagation(); + option.onDelete?.(); + }} + className="group/delete pointer-events-none absolute -right-4 -top-4 h-fit w-fit scale-50 cursor-pointer rounded-xl bg-zinc-900 p-1 text-white/75 opacity-0 duration-150 ease-in-out hover:text-white group-hover:pointer-events-auto group-hover:-right-2 group-hover:-top-2 group-hover:scale-100 group-hover:opacity-100" + > + +
+
+ )} + {option.image ? ( + Preset Image + ) : ( + hasImages && ( +
+ +
+ ) + )} +

+ {"name" in option && option.name} +

+
+ ) + )} );