Skip to content

Commit 0065f05

Browse files
committed
progress, handle cached nodes
1 parent d44d5d1 commit 0065f05

File tree

10 files changed

+144
-17
lines changed

10 files changed

+144
-17
lines changed

packages/stablestudio-ui/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
"notistack": "^3.0.0-alpha.11",
4343
"query-string": "^8.1.0",
4444
"react": "^18.2.0",
45+
"react-circular-progressbar": "^2.1.0",
4546
"react-dom": "^18.2.0",
4647
"react-konva": "^18.2.3",
4748
"react-konva-utils": "^0.3.1",

packages/stablestudio-ui/src-tauri/src/server.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ impl Builder {
7878
Some("text/html".to_string()),
7979
None,
8080
),
81-
["api" | "prompt" | "object_info" | "view" | "history" | "queue" | "interrupt" | "extensions", ..] => (
81+
["api" | "prompt" | "object_info" | "view" | "upload" | "history" | "queue" | "interrupt" | "extensions", ..] => (
8282
None,
8383
None,
8484
Some(

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

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,12 @@ type State = {
116116

117117
unlisteners: (() => void)[];
118118
setUnlisteners: (unlisteners: (() => void)[]) => void;
119+
120+
runningPrompt: ID | null;
121+
setRunningPrompt: (promptID: ID | null) => void;
122+
123+
lastOuput: ComfyOutput | null;
124+
setLastOutput: (output: ComfyOutput | null) => void;
119125
};
120126

121127
export namespace Comfy {
@@ -142,6 +148,12 @@ export namespace Comfy {
142148

143149
unlisteners: [],
144150
setUnlisteners: (unlisteners) => set({ unlisteners }),
151+
152+
runningPrompt: null,
153+
setRunningPrompt: (runningPrompt) => set({ runningPrompt }),
154+
155+
lastOuput: null,
156+
setLastOutput: (lastOuput) => set({ lastOuput }),
145157
}));
146158

147159
export const registerListeners = async () => {
@@ -152,10 +164,26 @@ export namespace Comfy {
152164
api = get()?.api;
153165
}
154166

155-
api.addEventListener("executed", async ({ detail }) => {
167+
api.addEventListener("progress", ({ detail }) => {
168+
console.log("progress", detail);
169+
const runningPrompt = use.getState().runningPrompt;
170+
if (runningPrompt) {
171+
Generation.Image.Output.set({
172+
...Generation.Image.Output.get(runningPrompt),
173+
progress: detail.value / detail.max,
174+
});
175+
}
176+
});
177+
178+
api.addEventListener("b_preview", ({ detail }) => {
179+
console.log("b_preview", detail);
180+
});
181+
182+
const executed = async ({ detail }: any) => {
156183
const { output, prompt_id } = detail;
157184

158185
console.log("executed_in_comfy_domain", detail);
186+
use.getState().setRunningPrompt(null);
159187

160188
const newInputs: Record<ID, Generation.Image.Input> = {};
161189
const responses: Generation.Images = [];
@@ -194,7 +222,7 @@ export namespace Comfy {
194222
const newInput = {
195223
...Generation.Image.Input.initial(inputID),
196224
...input,
197-
seed: 0,
225+
seed: (input?.seed ?? 0) + images.indexOf(image),
198226
id: inputID,
199227
};
200228

@@ -211,11 +239,32 @@ export namespace Comfy {
211239
});
212240
responses.forEach(Generation.Image.add);
213241
Generation.Image.Output.received(prompt_id, responses);
242+
use.getState().setLastOutput(detail);
243+
};
244+
245+
api.addEventListener("executed", executed);
246+
api.addEventListener("execution_cached", async ({ detail }) => {
247+
const last: any = use.getState().lastOuput;
248+
console.log("execution_cached", detail, last);
249+
if (
250+
use.getState().runningPrompt === detail.prompt_id &&
251+
detail.nodes.includes(last?.node) &&
252+
last.output
253+
) {
254+
console.log("last", last);
255+
const d = { ...last, prompt_id: detail.prompt_id };
256+
await executed({ detail: d });
257+
}
214258
});
215259

216260
api.addEventListener("execution_error", ({ detail }) => {
217261
console.log("execution_error", detail);
218262
Generation.Image.Output.clear(detail.prompt_id);
263+
use.getState().setRunningPrompt(null);
264+
});
265+
266+
api.addEventListener("execution_start", ({ detail }) => {
267+
use.getState().setRunningPrompt(detail?.prompt_id);
219268
});
220269

221270
console.log("registered ComfyUI listeners");

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ export namespace Create {
5757
pluginInput.width = Math.ceil((pluginInput.width ?? 512) / 64) * 64;
5858
}
5959

60+
const startingSeed =
61+
input.seed === 0 ? Math.round(Math.random() * 100000) : input.seed;
62+
6063
Comfy.get()
6164
?.graph._nodes?.filter((node) => node.type === "StableStudioNode")
6265
.forEach((node) => {
@@ -73,10 +76,7 @@ export namespace Create {
7376
positive_prompt:
7477
input.prompts.find((p) => p.weight > 0)?.text ??
7578
node.stableValues.positive_prompt,
76-
seed:
77-
input.seed === 0
78-
? Math.round(Math.random() * 100000)
79-
: input.seed,
79+
seed: startingSeed,
8080
steps: input.steps,
8181
cfg: input.cfgScale ?? node.stableValues.cfg,
8282
sampler_name:
@@ -97,6 +97,7 @@ export namespace Create {
9797
...inputs,
9898
[prompt_id]: {
9999
...input,
100+
seed: startingSeed,
100101
id: prompt_id,
101102
},
102103
}));

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,32 @@
1+
import { useLocalStorage } from "react-use";
12
import { Generation } from "~/Generation";
23
import { Theme } from "~/Theme";
34

45
export function Dropdown({ id, className }: Styleable & { id: ID }) {
56
const { setInput, input } = Generation.Image.Input.use(id);
67
const { data: models, isLoading } = Generation.Image.Models.use();
78

9+
const [value, setValue] = useLocalStorage<string | undefined>(
10+
"default-model-id",
11+
undefined
12+
);
13+
useEffect(() => {
14+
if (value) {
15+
setInput((input) => {
16+
input.model = value;
17+
});
18+
}
19+
}, []);
20+
821
const onClick = useCallback(
922
(value: string) => {
1023
setInput((input) => {
1124
console.log("model", value);
1225
input.model = value;
26+
setValue(value);
1327
});
1428
},
15-
[setInput]
29+
[setInput, setValue]
1630
);
1731

1832
const options = useMemo(

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,11 @@ export type Output = {
1313

1414
requestedAt?: Date;
1515
completedAt?: Date;
16+
progress?: number;
1617

1718
count: number;
1819
imageIDs: ID[];
20+
progressImageIDs?: ID[];
1921

2022
exception?: Generation.Image.Exception;
2123
};
@@ -61,6 +63,7 @@ export function Output({ outputID, placeholder, divider }: Props) {
6163
key={keys("image", images.length, images.length - index)}
6264
placeholder={placeholder}
6365
image={image}
66+
progress={output?.progress}
6467
scale={1}
6568
example={
6669
Generation.Image.Prompt.Examples.images[
@@ -76,7 +79,14 @@ export function Output({ outputID, placeholder, divider }: Props) {
7679
{rendered}
7780
</div>
7881
);
79-
}, [count, images, input?.id, placeholder, exampleStartIndex]);
82+
}, [
83+
count,
84+
images,
85+
input?.id,
86+
placeholder,
87+
output?.progress,
88+
exampleStartIndex,
89+
]);
8090

8191
const controls = useMemo(
8292
() => (

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

Lines changed: 45 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
1+
import { CircularProgressbar } from "react-circular-progressbar";
12
import { Generation } from "~/Generation";
23
import { Theme } from "~/Theme";
4+
import "react-circular-progressbar/dist/styles.css";
35

46
import { Filter } from "./Filter";
57

@@ -10,6 +12,7 @@ export function SpecialEffects({
1012
example,
1113
onClick,
1214
input,
15+
progress,
1316
}: {
1417
showing?: boolean;
1518
loading?: boolean;
@@ -21,6 +24,7 @@ export function SpecialEffects({
2124
onClick?: () => void;
2225
input?: ID;
2326
border?: boolean;
27+
progress?: number;
2428
}) {
2529
if (!showing) return null;
2630
return (
@@ -74,13 +78,47 @@ export function SpecialEffects({
7478
</div>
7579
)}
7680
<div className="absolute flex h-full w-full items-center justify-center">
77-
{loading && (
78-
<Theme.Loading.Spinner
79-
className={classes(
80-
variant === "small" ? "h-1/2 w-1/2" : "h-10 w-10"
81-
)}
82-
/>
83-
)}
81+
{loading &&
82+
(typeof progress === "number" ? (
83+
<div
84+
className={classes(
85+
variant === "small" ? "h-1/2 w-1/2" : "h-10 w-10"
86+
)}
87+
>
88+
<CircularProgressbar
89+
value={progress}
90+
maxValue={1}
91+
styles={{
92+
// Customize the path, i.e. the "completed progress"
93+
path: {
94+
// Path color
95+
stroke: "#ffffff",
96+
// Whether to use rounded or flat corners on the ends - can use 'butt' or 'round'
97+
strokeLinecap: "butt",
98+
// Customize transition animation
99+
transition: "stroke-dashoffset 0.5s ease 0s",
100+
transformOrigin: "center center",
101+
strokeWidth: 8,
102+
},
103+
// Customize the circle behind the path, i.e. the "total progress"
104+
trail: {
105+
// Trail color
106+
stroke: "#4c4c4d",
107+
// Whether to use rounded or flat corners on the ends - can use 'butt' or 'round'
108+
strokeLinecap: "butt",
109+
transformOrigin: "center center",
110+
strokeWidth: 8,
111+
},
112+
}}
113+
/>
114+
</div>
115+
) : (
116+
<Theme.Loading.Spinner
117+
className={classes(
118+
variant === "small" ? "h-1/2 w-1/2" : "h-10 w-10"
119+
)}
120+
/>
121+
))}
84122
</div>
85123
</div>
86124
);

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ type Props = Styleable &
4141
example?: Generation.Image.Prompt.Examples.Example;
4242
hideControls?: boolean;
4343
placeholder?: boolean;
44+
progress?: number;
4445

4546
onClick?: () => void;
4647
onDelete?: () => void;
@@ -51,6 +52,7 @@ export function Image({
5152
// example,
5253
hideControls,
5354
placeholder,
55+
progress,
5456

5557
scale,
5658
preserveAspectRatio,
@@ -155,13 +157,14 @@ export function Image({
155157
<Image.SpecialEffects
156158
showing={shouldShowSpecialEffects}
157159
loading={!placeholder && shouldShowSpecialEffects}
160+
progress={progress}
158161
variant={(style.height ?? 512) < 48 ? "small" : undefined}
159162
// example={example}
160163
// onClick={example ? onTryTemplate : undefined}
161164
// input={currentInput?.id}
162165
/>
163166
),
164-
[placeholder, shouldShowSpecialEffects, style.height]
167+
[placeholder, progress, shouldShowSpecialEffects, style.height]
165168
);
166169

167170
return useMemo(

packages/stablestudio-ui/vite.config.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ export default defineConfig(({ mode }) => {
4848
"/queue": redirectComfy,
4949
"/history": redirectComfy,
5050
"/interrupt": redirectComfy,
51+
"/upload": redirectComfy,
5152
},
5253
},
5354

yarn.lock

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1263,6 +1263,7 @@ __metadata:
12631263
prettier-plugin-tailwindcss: ^0.2.1
12641264
query-string: ^8.1.0
12651265
react: ^18.2.0
1266+
react-circular-progressbar: ^2.1.0
12661267
react-dom: ^18.2.0
12671268
react-konva: ^18.2.3
12681269
react-konva-utils: ^0.3.1
@@ -6322,6 +6323,15 @@ __metadata:
63226323
languageName: node
63236324
linkType: hard
63246325

6326+
"react-circular-progressbar@npm:^2.1.0":
6327+
version: 2.1.0
6328+
resolution: "react-circular-progressbar@npm:2.1.0"
6329+
peerDependencies:
6330+
react: ^0.14.0 || ^15.0.0 || ^16.0.0 || ^17.0.0 || ^18.0.0
6331+
checksum: dc118010a8f94733daafac586c969f7e889ad0736d96a0fda79406ee62e9410848abfd3dee4887bb0ff46b99e1f8c86ee1afabc88dfbf4dcb95107b22de1d6d7
6332+
languageName: node
6333+
linkType: hard
6334+
63256335
"react-dom@npm:^18.2.0":
63266336
version: 18.2.0
63276337
resolution: "react-dom@npm:18.2.0"

0 commit comments

Comments
 (0)