Basic preview support

(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

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);

View File

@@ -651,6 +651,10 @@ export default class ComfyApp {
}
});
this.api.addEventListener("b_preview", (imageBlob: Blob) => {
queueState.previewUpdated(imageBlob);
});
const config = get(configState);
if (config.pollSystemStatsInterval > 0) {

View File

@@ -386,6 +386,9 @@
<span style="display: inline-flex !important; padding: 0 0.75rem;">
<Checkbox label="Auto-Add UI" bind:value={$uiState.autoAddUI}/>
</span>
<span style="display: inline-flex !important; padding: 0 0.75rem;">
<Checkbox label="Hide Previews" bind:value={$uiState.hidePreviews}/>
</span>
<!-- <span class="label" for="ui-edit-mode">
<BlockTitle>UI Edit mode</BlockTitle>
<select id="ui-edit-mode" name="ui-edit-mode" bind:value={$uiState.uiEditMode}>

View File

@@ -9,7 +9,8 @@ import ComfyWidgetNode from "./ComfyWidgetNode";
export interface ComfyGalleryProperties extends ComfyWidgetProperties {
index: number | null,
updateMode: "replace" | "append",
autoSelectOnUpdate: boolean
autoSelectOnUpdate: boolean,
showPreviews: boolean
}
export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetadata[]> {
@@ -18,7 +19,8 @@ export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetad
defaultValue: [],
index: 0,
updateMode: "replace",
autoSelectOnUpdate: true
autoSelectOnUpdate: true,
showPreviews: true
}
static slotLayout: SlotLayout = {

View File

@@ -615,6 +615,14 @@ const ALL_ATTRIBUTES: AttributesSpecList = [
validNodeTypes: ["ui/gallery"],
defaultValue: true
},
{
name: "showPreviews",
type: "boolean",
location: "nodeProps",
editable: true,
validNodeTypes: ["ui/gallery"],
defaultValue: true
},
// ImageUpload
{

View File

@@ -22,6 +22,7 @@ type QueueStateOps = {
executionCached: (promptID: PromptID, nodes: ComfyNodeID[]) => void,
executionError: (error: ComfyExecutionError) => CompletedQueueEntry | null,
progressUpdated: (progress: Progress) => void
previewUpdated: (imageBlob: Blob) => void
getQueueEntry: (promptID: PromptID) => QueueEntry | null;
afterQueued: (workflowID: WorkflowInstID, promptID: PromptID, number: number, prompt: SerializedPromptInputsAll, extraData: any) => void
queueItemDeleted: (type: QueueItemType, id: PromptID) => void;
@@ -88,6 +89,11 @@ export type QueueState = {
*/
runningNodeID: ComfyNodeID | null;
/*
* Currently executing prompt if any
*/
runningPromptID: PromptID | null;
/*
* Nodes which should be rendered as "executing" in the frontend (green border).
* This includes the running node and all its parent subgraphs
@@ -98,6 +104,12 @@ export type QueueState = {
* Progress for the current node reported by the frontend
*/
progress: Progress | null,
/*
* Image preview URL
*/
previewURL: string | null,
/**
* If true, user pressed the "Interrupt" button in the frontend. Disable the
* button and wait until the next prompt starts running to re-enable it
@@ -115,6 +127,7 @@ const store: Writable<QueueState> = writable({
runningNodeID: null,
executingNodes: new Set(),
progress: null,
preview: null,
isInterrupting: false
})
@@ -171,6 +184,19 @@ function progressUpdated(progress: Progress) {
})
}
function previewUpdated(imageBlob: Blob) {
console.debug("[queueState] previewUpdated", imageBlob?.type)
store.update(s => {
if (s.runningNodeID == null) {
s.previewURL = null;
return s;
}
s.previewURL = URL.createObjectURL(imageBlob);
return s;
})
}
function statusUpdated(status: ComfyAPIStatusResponse | null) {
console.debug("[queueState] statusUpdated", status)
store.update((s) => {
@@ -296,6 +322,7 @@ function executingUpdated(promptID: PromptID, runningNodeID: ComfyNodeID | null)
entry.nodesRan.add(runningNodeID)
}
s.runningNodeID = runningNodeID;
s.runningPromptID = promptID;
if (entry?.extraData?.workflowID) {
const workflow = workflowState.getWorkflow(entry.extraData.workflowID);
@@ -337,7 +364,9 @@ function executingUpdated(promptID: PromptID, runningNodeID: ComfyNodeID | null)
console.debug("[queueState] Could not find in pending! (executingUpdated)", promptID)
}
s.progress = null;
s.previewURL = null;
s.runningNodeID = null;
s.runningPromptID = null;
s.executingNodes.clear();
}
entry_ = entry;
@@ -362,7 +391,9 @@ function executionCached(promptID: PromptID, nodes: ComfyNodeID[]) {
}
s.isInterrupting = false; // TODO move to start
s.progress = null;
s.previewURL = null;
s.runningNodeID = null;
s.runningPromptID = null;
s.executingNodes.clear();
return s
})
@@ -380,7 +411,9 @@ function executionError(error: ComfyExecutionError): CompletedQueueEntry | null
console.error("[queueState] Could not find in pending! (executionError)", error.prompt_id)
}
s.progress = null;
s.previewURL = null;
s.runningNodeID = null;
s.runningPromptID = null;
s.executingNodes.clear();
return s
})
@@ -416,6 +449,7 @@ function executionStart(promptID: PromptID) {
}
s.isInterrupting = false;
s.runningNodeID = null;
s.runningPromptID = promptID;
s.executingNodes.clear();
return s
})
@@ -480,7 +514,9 @@ function queueCleared(type: QueueItemType) {
s.queuePending.set([]);
s.queueRemaining = 0;
s.runningNodeID = null;
s.runningPromptID = null;
s.progress = null;
s.previewURL = null;
s.executingNodes.clear();
}
else {
@@ -535,6 +571,7 @@ const queueStateStore: WritableQueueStateStore =
historyUpdated,
statusUpdated,
progressUpdated,
previewUpdated,
executionStart,
executingUpdated,
executionCached,

View File

@@ -10,6 +10,7 @@ export type UIState = {
autoAddUI: boolean,
uiUnlocked: boolean,
uiEditMode: UIEditMode,
hidePreviews: boolean,
reconnecting: boolean,
forceSaveUserState: boolean | null,
@@ -30,6 +31,7 @@ const store: Writable<UIState> = writable(
autoAddUI: true,
uiUnlocked: false,
uiEditMode: "widgets",
hidePreviews: false,
reconnecting: false,
forceSaveUserState: null,

View File

@@ -11,7 +11,11 @@
import { clamp, comfyBoxImageToComfyURL, type ComfyBoxImageMetadata } from "$lib/utils";
import { f7 } from "framework7-svelte";
import type { ComfyGalleryNode } from "$lib/nodes/widgets";
import { showMobileLightbox } from "$lib/components/utils";
import { showMobileLightbox } from "$lib/components/utils";
import queueState from "$lib/stores/queueState";
import uiState from "$lib/stores/uiState";
import { loadImage } from "./utils";
import Spinner from "$lib/components/Spinner.svelte";
export let widget: WidgetLayout | null = null;
export let isMobile: boolean = false;
@@ -25,6 +29,47 @@
$: widget && setNodeValue(widget);
function tagsMatch(tags: string[] | null): boolean {
if(tags != null && tags.length > 0)
return node.properties.tags.length > 0 && node.properties.tags.every(t => tags.includes(t));
else
return node.properties.tags.length === 0;
}
let previewURL: string | null;
let previewImage: HTMLImageElement | null = null;
let previewElem: HTMLImageElement | null = null
$: {
previewURL = $queueState.previewURL;
if (previewURL && $queueState.runningPromptID != null && !$uiState.hidePreviews && node.properties.showPreviews) {
const queueEntry = queueState.getQueueEntry($queueState.runningPromptID)
if (queueEntry != null) {
const tags = queueEntry.extraData?.extra_pnginfo?.comfyBoxPrompt?.subgraphs;
if (tagsMatch(tags)) {
loadImage(previewURL).then((img) => {
previewImage = img;
})
}
else {
previewImage = null;
}
}
else {
previewImage = null;
}
}
else {
previewImage = null;
}
}
function showPreview() {
}
function hidePreview() {
}
function setNodeValue(widget: WidgetLayout) {
if (widget) {
node = widget.node as ComfyGalleryNode
@@ -34,6 +79,8 @@
imageHeight = node.imageHeight
selected_image = node.selectedImage;
forceSelectImage = node.forceSelectImage;
previewURL = null;
previewImage = null;
if ($nodeValue != null) {
if (node.properties.index < 0 || node.properties.index >= $nodeValue.length) {
@@ -108,6 +155,11 @@
<div class="wrapper comfy-gallery-widget gradio-gallery" style={widget.attrs.style || ""}>
<Block variant="solid" padding={false}>
<div class="padding">
{#if previewImage}
<div class="comfy-gallery-preview" on:mouseover={hidePreview} on:mouseout={showPreview} >
<img src={previewImage.src} bind:this={previewElem} on:mouseout={showPreview} />
</div>
{/if}
<Gallery
value={images}
label={widget.attrs.title}
@@ -153,6 +205,29 @@
}
}
}
&:hover .comfy-gallery-preview {
opacity: 0%;
}
}
.comfy-gallery-preview {
position: absolute;
top: 0;
left: 0;
width: 100%;
height: 100%;
z-index: var(--layer-top);
pointer-events: none;
transition: opacity 0.1s linear;
opacity: 100%;
> img {
width: var(--size-full);
height: var(--size-full);
object-fit: contain;
border: 5px dashed var(--secondary-400);
}
}
.padding {