neuro_llc/ComfyBox
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 1 Wiki Activity

Basic preview support

Browse Source 
(as of latest PR commit)
This commit is contained in:
space-nuko
2023-06-05 15:45:08 -05:00
parent fde480cb43
commit eb02561906
8 changed files with 198 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
src/lib/api.ts
View File

@@ -101,6 +101,7 @@ export type ComfyUIPromptExtraData = {
}
type ComfyAPIEvents = {
// JSON
status: (status: ComfyAPIStatusResponse | null, error?: Error | null) => void,
progress: (progress: Progress) => void,
reconnecting: () => void,
@@ -111,6 +112,9 @@ type ComfyAPIEvents = {
execution_cached: (promptID: PromptID, nodes: ComfyNodeID[]) => void,
execution_interrupted: (error: ComfyInterruptedError) => void,
execution_error: (error: ComfyExecutionError) => void,
// Binary
b_preview: (imageBlob: Blob) => void
}
export default class ComfyAPI {
@@ -126,7 +130,7 @@ export default class ComfyAPI {
}
/**
* Poll status for colab and other things that don't support websockets.
* Poll status for colab and other things that don't support websockets.
*/
private pollQueue() {
setInterval(async () => {
@@ -176,6 +180,7 @@ export default class ComfyAPI {
this.socket = new WebSocket(
`ws${window.location.protocol === "https:" ? "s" : ""}://${hostname}:${port}/ws${existingSession}`
);
this.socket.binaryType = "arraybuffer";
this.socket.addEventListener("open", () => {
opened = true;
@@ -204,38 +209,64 @@ export default class ComfyAPI {
this.socket.addEventListener("message", (event) => {
try {
const msg = JSON.parse(event.data);
switch (msg.type) {
case "status":
if (msg.data.sid) {
this.clientId = msg.data.sid;
sessionStorage["Comfy.SessionId"] = this.clientId;
}
this.eventBus.emit("status", { execInfo: { queueRemaining: msg.data.status.exec_info.queue_remaining } });
break;
case "progress":
this.eventBus.emit("progress", msg.data as Progress);
break;
case "executing":
this.eventBus.emit("executing", msg.data.prompt_id, msg.data.node);
break;
case "executed":
this.eventBus.emit("executed", msg.data.prompt_id, msg.data.node, msg.data.output);
break;
case "execution_start":
this.eventBus.emit("execution_start", msg.data.prompt_id);
break;
case "execution_cached":
this.eventBus.emit("execution_cached", msg.data.prompt_id, msg.data.nodes);
break;
case "execution_interrupted":
this.eventBus.emit("execution_interrupted", msg.data);
break;
case "execution_error":
this.eventBus.emit("execution_error", msg.data);
break;
default:
console.warn("Unhandled message:", event.data);
if (event.data instanceof ArrayBuffer) {
const view = new DataView(event.data);
const eventType = view.getUint32(0);
const buffer = event.data.slice(4);
switch (eventType) {
case 1:
const view2 = new DataView(event.data);
const imageType = view2.getUint32(0)
let imageMime: string
switch (imageType) {
case 1:
default:
imageMime = "image/jpeg";
break;
case 2:
imageMime = "image/png"
}
const imageBlob = new Blob([buffer.slice(4)], { type: imageMime });
this.eventBus.emit("b_preview", imageBlob);
break;
default:
throw new Error(`Unknown binary websocket message of type ${eventType}`);
}
}
else {
const msg = JSON.parse(event.data);
switch (msg.type) {
case "status":
if (msg.data.sid) {
this.clientId = msg.data.sid;
sessionStorage["Comfy.SessionId"] = this.clientId;
}
this.eventBus.emit("status", { execInfo: { queueRemaining: msg.data.status.exec_info.queue_remaining } });
break;
case "progress":
this.eventBus.emit("progress", msg.data as Progress);
break;
case "executing":
this.eventBus.emit("executing", msg.data.prompt_id, msg.data.node);
break;
case "executed":
this.eventBus.emit("executed", msg.data.prompt_id, msg.data.node, msg.data.output);
break;
case "execution_start":
this.eventBus.emit("execution_start", msg.data.prompt_id);
break;
case "execution_cached":
this.eventBus.emit("execution_cached", msg.data.prompt_id, msg.data.nodes);
break;
case "execution_interrupted":
this.eventBus.emit("execution_interrupted", msg.data);
break;
case "execution_error":
this.eventBus.emit("execution_error", msg.data);
break;
default:
console.warn("Unhandled message:", event.data);
}
}
} catch (error) {
console.error("Error handling message", event.data, error);