From 88045a0cdbe5405cbadf04dcd1c86735b8b2d0ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=A5=E6=B5=AA?= <757078144@qq.com> Date: Tue, 1 Oct 2024 17:09:08 +0800 Subject: [PATCH] =?UTF-8?q?=E7=AD=89=E5=BE=85=E9=80=89=E6=8B=A9=EF=BC=8C?= =?UTF-8?q?=E5=BE=AA=E7=8E=AF=E6=98=BE=E7=A4=BA=E4=BC=98=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- js/MultiTextConcatenate.js | 13 +- js/StatusInfo.js | 2 +- js/WaitImagSelector.js | 230 ++++++++++++++++++++++++++++++++++ js/image_chooser/messaging.js | 34 +++++ js/image_chooser/preview.js | 90 +++++++++++++ js/image_chooser/prompt.js | 114 +++++++++++++++++ js/image_chooser/state.js | 55 ++++++++ py/DoWhileEnd.py | 13 ++ py/LamCommonPrint.py | 30 ++++- py/WaitImagSelector.py | 93 ++++++++++++++ py/src/utils/chooser.py | 52 ++++++++ 11 files changed, 722 insertions(+), 4 deletions(-) create mode 100644 js/WaitImagSelector.js create mode 100644 js/image_chooser/messaging.js create mode 100644 js/image_chooser/preview.js create mode 100644 js/image_chooser/prompt.js create mode 100644 js/image_chooser/state.js create mode 100644 py/WaitImagSelector.py create mode 100644 py/src/utils/chooser.py diff --git a/js/MultiTextConcatenate.js b/js/MultiTextConcatenate.js index d4609f5..888a81e 100644 --- a/js/MultiTextConcatenate.js +++ b/js/MultiTextConcatenate.js @@ -9,12 +9,23 @@ app.registerExtension({ var showNames=["IfInnerExecute",'MultiIntFormula','MultiParamFormula','ForInnerEnd','DoWhileEnd',"LamSwitcherCase"] if (showNames.indexOf(nodeData.name)>=0) { const onDrawForeground = nodeType.prototype.onDrawForeground; + let oldText='' nodeType.prototype.onDrawForeground = function (ctx) { try { const r = onDrawForeground?.apply?.(this, arguments); const v = app.nodeOutputs?.[this.id + ""]; if (!this.flags.collapsed && v) { - const text = v.value[0] + ""; + let text='' + if(['ForInnerEnd','DoWhileEnd'].indexOf(nodeData.name)>=0){ + if(v.value[0].indexOf('DOWHILE')!=-1){ + text=v.value[0].split('DOWHILE')[1] + oldText=text + }else{ + text=oldText + } + }else{ + text = v.value[0] + ""; + } ctx.save(); ctx.font = "bold 12px sans-serif"; ctx.fillStyle = "dodgerblue"; diff --git a/js/StatusInfo.js b/js/StatusInfo.js index 54e60a0..5b80a39 100644 --- a/js/StatusInfo.js +++ b/js/StatusInfo.js @@ -7,7 +7,7 @@ app.registerExtension({ name: "StatusInfo", async beforeRegisterNodeDef(nodeType, nodeData, app) { var names=["Image2Video",'Video2TalkingFace','VideoAddAudio','Image2TalkingFace','ForEnd','LoadVideo','VideoFaceFusion', - 'Text2AutioEdgeTts','VideoRoopFaceSwap','PromptTranslator','LamCommonPrint','OutDoWhileEnd'] + 'Text2AutioEdgeTts','VideoRoopFaceSwap','PromptTranslator','LamCommonPrint','LamCommonPrintNoOutput','OutDoWhileEnd'] if (names.indexOf(nodeData.name)>=0) { // When the node is created we want to add a readonly text widget to display the text const onNodeCreated = nodeType.prototype.onNodeCreated; diff --git a/js/WaitImagSelector.js b/js/WaitImagSelector.js new file mode 100644 index 0000000..76a0d98 --- /dev/null +++ b/js/WaitImagSelector.js @@ -0,0 +1,230 @@ +import { app } from "../../../../scripts/app.js"; +import { api } from "../../../../scripts/api.js"; +import { ComfyDialog, $el } from "../../../../scripts/ui.js"; + +import { restart_from_here } from "./image_chooser/prompt.js"; +import { hud, FlowState } from "./image_chooser/state.js"; +import { send_cancel, send_message, send_onstart, skip_next_restart_message } from "./image_chooser/messaging.js"; +import { display_preview_images, additionalDrawBackground, click_is_in_image } from "./image_chooser/preview.js"; + + +class chooserImageDialog extends ComfyDialog { + + constructor() { + super(); + this.node = null + this.select_index = [] + this.dialog_div = null + } + + show(image,node){ + this.select_index = [] + this.node = node + + const images_div = image.map((img, index) => { + const imgEl = $el('img', { + src: img.src, + onclick: _ => { + if(this.select_index.includes(index)){ + this.select_index = this.select_index.filter(i => i !== index) + imgEl.classList.remove('selected') + } else { + this.select_index.push(index) + imgEl.classList.add('selected') + } + if (node.selected.has(index)) node.selected.delete(index); + else node.selected.add(index); + } + }) + return imgEl + }) + super.show($el('div.easyuse-chooser-dialog',[ + $el('h5.easyuse-chooser-dialog-title', '选择图像继续'), + $el('div.easyuse-chooser-dialog-images',images_div) + ])) + } + createButtons() { + const btns = super.createButtons(); + btns[0].onclick = _ => { + if (FlowState.running()) { send_cancel();} + super.close() + } + btns.unshift($el('button', { + type: 'button', + textContent: "选择已选图像", + onclick: _ => { + if (FlowState.paused()) { + send_message(this.node.id, [...this.node.selected, -1, ...this.node.anti_selected]); + } + if (FlowState.idle()) { + skip_next_restart_message(); + restart_from_here(this.node.id).then(() => { send_message(this.node.id, [...this.node.selected, -1, ...this.node.anti_selected]); }); + } + super.close() + } + })) + return btns + } + +} + +function progressButtonPressed() { + const node = app.graph._nodes_by_id[this.node_id]; + if (node) { + const selected = [...node.selected] + if(selected?.length>0){ + node.setProperty('values',selected) + } + if (FlowState.paused()) { + send_message(node.id, [...node.selected, -1, ...node.anti_selected]); + } + if (FlowState.idle()) { + skip_next_restart_message(); + restart_from_here(node.id).then(() => { send_message(node.id, [...node.selected, -1, ...node.anti_selected]); }); + } + } +} + +function cancelButtonPressed() { + + if (FlowState.running()) { send_cancel();} +} + +function enable_disabling(button) { + Object.defineProperty(button, 'clicked', { + get : function() { return this._clicked; }, + set : function(v) { this._clicked = (v && this.name!=''); } + }) +} + +function disable_serialize(widget) { + if (!widget.options) widget.options = { }; + widget.options.serialize = false; +} + +app.registerExtension({ + name:'comfy.WaitImagSelector', + init() { + window.addEventListener("beforeunload", send_cancel, true); + }, + setup(app) { + + const draw = LGraphCanvas.prototype.draw; + LGraphCanvas.prototype.draw = function() { + if (hud.update()) { + app.graph._nodes.forEach((node)=> { if (node.update) { node.update(); } }) + } + draw.apply(this,arguments); + } + + + function easyuseImageChooser(event) { + const {node,image,isKSampler} = display_preview_images(event); + if(isKSampler) { + const dialog = new chooserImageDialog(); + dialog.show(image,node) + } + } + api.addEventListener("lam-wait-image-select", easyuseImageChooser); + /* + If a run is interrupted, send a cancel message (unless we're doing the cancelling, to avoid infinite loop) + */ + const original_api_interrupt = api.interrupt; + api.interrupt = function () { + if (FlowState.paused() && !FlowState.cancelling) send_cancel(); + original_api_interrupt.apply(this, arguments); + } + + /* + At the start of execution + */ + function on_execution_start(event) { + if (send_onstart()) { + app.graph._nodes.forEach((node)=> { + if (node.selected || node.anti_selected) { + node.selected.clear(); + node.anti_selected.clear(); + node.update(); + } + }) + } + } + api.addEventListener("execution_start", on_execution_start); + }, + + async nodeCreated(node, app) { + + if(node.comfyClass == 'WaitImagSelector'){ + node.setProperty('values',[]) + + /* A property defining the top of the image when there is just one */ + if(node?.imageIndex === undefined){ + Object.defineProperty(node, 'imageIndex', { + get : function() { return null; }, + set: function (v) {node.overIndex= v}, + }) + } + if(node?.imagey === undefined){ + Object.defineProperty(node, 'imagey', { + get : function() { return null; }, + set: function (v) {return node.widgets[node.widgets.length-1].last_y+LiteGraph.NODE_WIDGET_HEIGHT;}, + }) + } + + /* Capture clicks */ + const org_onMouseDown = node.onMouseDown; + node.onMouseDown = function( e, pos, canvas ) { + if (e.isPrimary) { + const i = click_is_in_image(node, pos); + if (i>=0) { this.imageClicked(i); } + } + return (org_onMouseDown && org_onMouseDown.apply(this, arguments)); + } + + node.send_button_widget = node.addWidget("button", "确认", "", progressButtonPressed); + //node.cancel_button_widget = node.addWidget("button", "", "", cancelButtonPressed); + //enable_disabling(node.cancel_button_widget); + enable_disabling(node.send_button_widget); + //disable_serialize(node.cancel_button_widget); + disable_serialize(node.send_button_widget); + + } + }, + + beforeRegisterNodeDef(nodeType, nodeData, app) { + if(nodeData?.name == 'WaitImagSelector'){ + + const onDrawBackground = nodeType.prototype.onDrawBackground; + nodeType.prototype.onDrawBackground = function(ctx) { + onDrawBackground.apply(this, arguments); + additionalDrawBackground(this, ctx); + } + + nodeType.prototype.imageClicked = function (imageIndex) { + if (nodeType?.comfyClass==="WaitImagSelector") { + if (this.selected.has(imageIndex)) this.selected.delete(imageIndex); + else this.selected.add(imageIndex); + this.update(); + } + } + + const update = nodeType.prototype.update; + nodeType.prototype.update = function() { + if (update) update.apply(this,arguments); + if (this.send_button_widget) { + this.send_button_widget.node_id = this.id; + const selection = ( this.selected ? this.selected.size : 0 ) + ( this.anti_selected ? this.anti_selected.size : 0 ) + const maxlength = this.imgs?.length || 0; + if (FlowState.paused_here(this.id) && selection>0) { + this.send_button_widget.name = (selection>1) ? "确认选择 (" + selection + '/' + maxlength +")" : "确认选择"; + } + } + // if (this.cancel_button_widget) { + // const isRunning = FlowState.running() + // this.cancel_button_widget.name = isRunning ? "取消选择" : ""; + // } + this.setDirtyCanvas(true,true); + } + } + } +}) \ No newline at end of file diff --git a/js/image_chooser/messaging.js b/js/image_chooser/messaging.js new file mode 100644 index 0000000..8c5e654 --- /dev/null +++ b/js/image_chooser/messaging.js @@ -0,0 +1,34 @@ +import { api } from "../../../../scripts/api.js"; +import { FlowState } from "./state.js"; + +function send_message_from_pausing_node(message) { + const id = app.runningNodeId; + send_message(id, message); +} + +function send_message(id, message) { + const body = new FormData(); + body.append('message',message); + body.append('id', id); + api.fetchApi("/lam/image_chooser_message", { method: "POST", body, }); +} + +function send_cancel() { + send_message(-1,'__cancel__'); + FlowState.cancelling = true; + api.interrupt(); + FlowState.cancelling = false; +} + +var skip_next = 0; +function skip_next_restart_message() { skip_next += 1; } +function send_onstart() { + if (skip_next>0) { + skip_next -= 1; + return false; + } + send_message(-1,'__start__'); + return true; +} + +export { send_message_from_pausing_node, send_cancel, send_message, send_onstart, skip_next_restart_message } \ No newline at end of file diff --git a/js/image_chooser/preview.js b/js/image_chooser/preview.js new file mode 100644 index 0000000..6f8fe44 --- /dev/null +++ b/js/image_chooser/preview.js @@ -0,0 +1,90 @@ +import { app } from "../../../../scripts/app.js"; + +const kSampler = ['easy kSampler', 'easy kSamplerTiled', 'easy fullkSampler'] + +function display_preview_images(event) { + const node = app.graph._nodes_by_id[event.detail.id]; + if (node) { + node.selected = new Set(); + node.anti_selected = new Set(); + const image = showImages(node, event.detail.urls); + return {node,image,isKSampler:kSampler.includes(node.type)} + } else { + console.log(`Image Chooser Preview - failed to find ${event.detail.id}`) + } +} + +function showImages(node, urls) { + node.imgs = []; + urls.forEach((u)=> { + const img = new Image(); + node.imgs.push(img); + img.onload = () => { app.graph.setDirtyCanvas(true); }; + img.src = `/view?filename=${encodeURIComponent(u.filename)}&type=temp&subfolder=${app.getPreviewFormatParam()}` + }) + node.setSizeForImage?.(); + return node.imgs +} + +function drawRect(node, s, ctx) { + const padding = 1; + var rect; + if (node.imageRects) { + rect = node.imageRects[s]; + } else { + const y = node.imagey; + rect = [padding,y+padding,node.size[0]-2*padding,node.size[1]-y-2*padding]; + } + ctx.strokeRect(rect[0]+padding, rect[1]+padding, rect[2]-padding*2, rect[3]-padding*2); +} + +function additionalDrawBackground(node, ctx) { + if (!node.imgs) return; + if (node.imageRects) { + for (let i = 0; i < node.imgs.length; i++) { + // delete underlying image + ctx.fillStyle = "#000"; + ctx.fillRect(...node.imageRects[i]) + // draw the new one + const img = node.imgs[i]; + const cellWidth = node.imageRects[i][2]; + const cellHeight = node.imageRects[i][3]; + + let wratio = cellWidth/img.width; + let hratio = cellHeight/img.height; + var ratio = Math.min(wratio, hratio); + + let imgHeight = ratio * img.height; + let imgWidth = ratio * img.width; + + const imgX = node.imageRects[i][0] + (cellWidth - imgWidth)/2; + const imgY = node.imageRects[i][1] + (cellHeight - imgHeight)/2; + const cell_padding = 2; + ctx.drawImage(img, imgX+cell_padding, imgY+cell_padding, imgWidth-cell_padding*2, imgHeight-cell_padding*2); + + } + } + ctx.lineWidth = 2; + ctx.strokeStyle = "green"; + node?.selected?.forEach((s) => { drawRect(node,s, ctx) }) + ctx.strokeStyle = "#F88"; + node?.anti_selected?.forEach((s) => { drawRect(node,s, ctx) }) +} + +function click_is_in_image(node, pos) { + if (node.imgs?.length>1) { + for (var i = 0; i 0 && dx < node.imageRects[i][2] && + dy > 0 && dy < node.imageRects[i][3] ) { + return i; + } + } + } else if (node.imgs?.length==1) { + if (pos[1]>node.imagey) return 0; + } + return -1; +} + +export { display_preview_images, additionalDrawBackground, click_is_in_image } \ No newline at end of file diff --git a/js/image_chooser/prompt.js b/js/image_chooser/prompt.js new file mode 100644 index 0000000..c9c89b6 --- /dev/null +++ b/js/image_chooser/prompt.js @@ -0,0 +1,114 @@ +import { app } from "../../../../scripts/app.js"; + +function links_with(p, node_id, down, up) { + const links_with = []; + p.workflow.links.forEach((l) => { + if (down && l[1]===node_id && !links_with.includes(l[3])) links_with.push(l[3]) + if (up && l[3]===node_id && !links_with.includes(l[1])) links_with.push(l[1]) + }); + return links_with; +} + +function _all_v_nodes(p, here_id) { + /* + Make a list of all downstream nodes. + */ + const downstream = []; + const to_process = [here_id] + while(to_process.length>0) { + const id = to_process.pop(); + downstream.push(id); + to_process.push( + ...links_with(p,id,true,false).filter((nid)=>{ + return !(downstream.includes(nid) || to_process.includes(nid)) + }) + ) + } + + /* + Now all upstream nodes from any of the downstream nodes (except us). + Put us on the result list so we don't flow up through us + */ + to_process.push(...downstream.filter((n)=>{ return n!=here_id})); + const back_upstream = [here_id]; + while(to_process.length>0) { + const id = to_process.pop(); + back_upstream.push(id); + to_process.push( + ...links_with(p,id,false,true).filter((nid)=>{ + return !(back_upstream.includes(nid) || to_process.includes(nid)) + }) + ) + } + + const keep = []; + keep.push(...downstream); + keep.push(...back_upstream.filter((n)=>{return !keep.includes(n)})); + + console.log(`Nodes to keep: ${keep}`); + return keep; +} + +async function all_v_nodes(here_id) { + const p = structuredClone(await app.graphToPrompt()); + const all_nodes = []; + p.workflow.nodes.forEach((node)=>{all_nodes.push(node.id)}) + p.workflow.links = p.workflow.links.filter((l)=>{ return (all_nodes.includes(l[1]) && all_nodes.includes(l[3]))} ) + return _all_v_nodes(p,here_id); +} + +async function restart_from_here(here_id, go_down_to_chooser=false) { + const p = structuredClone(await app.graphToPrompt()); + /* + Make a list of all nodes, and filter out links that are no longer valid + */ + const all_nodes = []; + p.workflow.nodes.forEach((node)=>{all_nodes.push(node.id)}) + p.workflow.links = p.workflow.links.filter((l)=>{ return (all_nodes.includes(l[1]) && all_nodes.includes(l[3]))} ) + + /* Move downstream to a chooser */ + if (go_down_to_chooser) { + while (!app.graph._nodes_by_id[here_id].isChooser) { + here_id = links_with(p, here_id, true, false)[0]; + } + } + + const keep = _all_v_nodes(p, here_id); + + /* + Filter p.workflow.nodes and p.workflow.links + */ + p.workflow.nodes = p.workflow.nodes.filter((node) => { + if (node.id===here_id) node.inputs.forEach((i)=>{i.link=null}) // remove our upstream links + return (keep.includes(node.id)) // only keep keepers + }) + p.workflow.links = p.workflow.links.filter((l) => {return (keep.includes(l[1]) && keep.includes(l[3]))}) + + /* + Filter the p.output object to only include nodes we're keeping + */ + const new_output = {} + for (const [key, value] of Object.entries(p.output)) { + if (keep.includes(parseInt(key))) new_output[key] = value; + } + /* + Filter the p.output entry for the start node to remove any list (ie link) inputs + */ + const new_inputs = {}; + for (const [key, value] of Object.entries(new_output[here_id.toString()].inputs)) { + if (!Array.isArray(value)) new_inputs[key] = value; + } + new_output[here_id.toString()].inputs = new_inputs; + + p.output = new_output; + + // temporarily hijack graph_to_prompt with a version that restores the old one but returns this prompt + const gtp_was = app.graphToPrompt; + app.graphToPrompt = () => { + app.graphToPrompt = gtp_was; + return p; + } + app.queuePrompt(0); +} + +export { restart_from_here, all_v_nodes } \ No newline at end of file diff --git a/js/image_chooser/state.js b/js/image_chooser/state.js new file mode 100644 index 0000000..69437d7 --- /dev/null +++ b/js/image_chooser/state.js @@ -0,0 +1,55 @@ +import { app } from "../../../../scripts/app.js"; + + +class HUD { + constructor() { + this.current_node_id = undefined; + this.class_of_current_node = null; + this.current_node_is_chooser = false; + } + + update() { + if (app.runningNodeId==this.current_node_id) return false; + + this.current_node_id = app.runningNodeId; + + if (this.current_node_id) { + this.class_of_current_node = app.graph?._nodes_by_id[app.runningNodeId.toString()]?.comfyClass; + this.current_node_is_chooser = this.class_of_current_node === "easy imageChooser" + } else { + this.class_of_current_node = undefined; + this.current_node_is_chooser = false; + } + return true; + } +} + +const hud = new HUD(); + + +class FlowState { + constructor(){} + static idle() { + return (!app.runningNodeId); + } + static paused() { + return true; + } + static paused_here(node_id) { + return (FlowState.paused() && FlowState.here(node_id)) + } + static running() { + return (!FlowState.idle()); + } + static here(node_id) { + return (app.runningNodeId==node_id); + } + static state() { + if (FlowState.paused()) return "Paused"; + if (FlowState.running()) return "Running"; + return "Idle"; + } + static cancelling = false; +} + +export { hud, FlowState} \ No newline at end of file diff --git a/py/DoWhileEnd.py b/py/DoWhileEnd.py index 9e2b096..08e559d 100644 --- a/py/DoWhileEnd.py +++ b/py/DoWhileEnd.py @@ -32,15 +32,27 @@ class DoWhileEnd: CATEGORY = "lam" + def getTaskOvers(self,nodeId): + if nodeId.find("Recurse")!= -1: + recId=nodeId.split('Recurse')[0] + teskSize=nodeId.count(recId)+1 + return 'DOWHILE'+','.join([str(i+1) for i in range(max(0,teskSize-10),teskSize)]) + elif nodeId.find(".")!= -1: + return 'DOWHILE'+'1,2' + else: + return 'DOWHILE'+'1' + def for_end_fun(self,start,ANY,obj,stop=1,dynprompt=None, unique_id=None, **kwargs): assert dynprompt is not None graph = GraphBuilder() open_node = start[0] + overIds=self.getTaskOvers(open_node) if not ANY: values = [[open_node,1],obj] for i in range(NUM_FLOW_SOCKETS): values.append(kwargs.get(f"initial_value{i}", None)) return { + 'ui':{'value':[overIds]}, "result": tuple(values), "expand": graph.finalize(), } @@ -84,6 +96,7 @@ class DoWhileEnd: assert my_clone is not None result = map(lambda x: my_clone.out(x), range(NUM_FLOW_SOCKETS+2)) return { + 'ui':{'value':[overIds]}, "result": tuple(result), "expand": graph.finalize(), } diff --git a/py/LamCommonPrint.py b/py/LamCommonPrint.py index 63f37ce..72178a7 100644 --- a/py/LamCommonPrint.py +++ b/py/LamCommonPrint.py @@ -22,11 +22,37 @@ class LamCommonPrint: def common_print(self,obj): return {"ui": {"text": str(obj)}, "result": (obj, )} + +class LamCommonPrintNoOutput: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "obj": (AlwaysEqualProxy("*"),), + } + } + RETURN_TYPES = (AlwaysEqualProxy("*"),) + + RETURN_NAMES = ('obj',) + + FUNCTION = "common_print" + + OUTPUT_NODE = False + + CATEGORY = "lam" + + def common_print(self,obj): + return {"ui": {"text": str(obj)}, "result": (obj, )} NODE_CLASS_MAPPINGS = { - "LamCommonPrint": LamCommonPrint + "LamCommonPrint": LamCommonPrint, + "LamCommonPrintNoOutput":LamCommonPrintNoOutput } NODE_DISPLAY_NAME_MAPPINGS = { - "LamCommonPrint": "通用打印输出" + "LamCommonPrint": "通用打印输出", + "LamCommonPrintNoOutput": "通用打印输出(非输出节点)" } diff --git a/py/WaitImagSelector.py b/py/WaitImagSelector.py new file mode 100644 index 0000000..6eea3cc --- /dev/null +++ b/py/WaitImagSelector.py @@ -0,0 +1,93 @@ +import folder_paths +from nodes import SaveImage +import random +from PIL import Image, ImageOps +import os +import numpy as np +import json +from PIL.PngImagePlugin import PngInfo +from comfy.cli_args import args +from server import PromptServer +from .src.utils.chooser import ChooserMessage, ChooserCancelled + + +class WaitImagSelector(SaveImage): + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + self.prefix_append = "_temp_" + \ + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") + for x in range(5)) + self.compress_level = 1 + + @classmethod + def INPUT_TYPES(s): + return {"required": + {"images": ("IMAGE", ), }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "unique_id": "UNIQUE_ID"}, + } + + RETURN_TYPES = ("BOOL","IMAGE",) + RETURN_NAMES = ("是否没选","选择图片",) + FUNCTION = "save_images" + CATEGORY = "image" + OUTPUT_NODE = True + + last_ic = {} + @classmethod + def IS_CHANGED(cls, unique_id, **kwargs): + return cls.last_ic[unique_id.split(".")[-1]] + + def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, unique_id=None): + id = unique_id.split(".")[-1] + if id not in ChooserMessage.stash: + ChooserMessage.stash[id] = {} + my_stash = ChooserMessage.stash[id] + + filename_prefix += self.prefix_append + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( + filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) + results = list() + for (batch_number, image) in enumerate(images): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = None + if not args.disable_metadata: + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + filename_with_batch_num = filename.replace( + "%batch_num%", str(batch_number)) + file = f"{filename_with_batch_num}_{counter:05}_.png" + img.save(os.path.join(full_output_folder, file), + pnginfo=metadata, compress_level=self.compress_level) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + PromptServer.instance.send_sync( + "lam-wait-image-select", {"id": id, "urls": results}) + try: + selections = ChooserMessage.waitForMessage(id, asList=True) + if len(selections)==1 and selections[0]==-1: + return {"ui": {"images": results},"result": (True,None,)} + selImgs = images[selections[:-1]] + return {"ui": {"images": results},"result": (False,selImgs,)} + except ChooserCancelled: + return {"ui": {"images": results},"result": (True,None,)} + + + +NODE_CLASS_MAPPINGS = { + "WaitImagSelector": WaitImagSelector +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "WaitImagSelector": "等待图片选择器" +} diff --git a/py/src/utils/chooser.py b/py/src/utils/chooser.py new file mode 100644 index 0000000..44ea509 --- /dev/null +++ b/py/src/utils/chooser.py @@ -0,0 +1,52 @@ +from server import PromptServer +from aiohttp import web +import time + +class ChooserCancelled(Exception): + pass + +class ChooserMessage: + stash = {} + messages = {} + cancelled = False + + @classmethod + def addMessage(cls, id, message): + if message == '__cancel__': + cls.messages = {} + cls.cancelled = True + elif message == '__start__': + cls.messages = {} + cls.stash = {} + cls.cancelled = False + else: + cls.messages[str(id)] = message + + @classmethod + def waitForMessage(cls, id, period=0.1, asList=False): + sid = str(id) + while not (sid in cls.messages) and not ("-1" in cls.messages): + if cls.cancelled: + cls.cancelled = False + raise ChooserCancelled() + time.sleep(period) + if cls.cancelled: + cls.cancelled = False + raise ChooserCancelled() + message = cls.messages.pop(str(id), None) or cls.messages.pop("-1") + try: + if asList: + return [int(x.strip()) for x in message.split(",")] + else: + return int(message.strip()) + except ValueError: + print( + f"ERROR IN IMAGE_CHOOSER - failed to parse '${message}' as ${'comma separated list of ints' if asList else 'int'}") + return [-1] if asList else -1 + + +@PromptServer.instance.routes.post('/lam/image_chooser_message') +async def make_image_selection(request): + post = await request.post() + ChooserMessage.addMessage(post.get("id"), post.get("message")) + return web.json_response({}) \ No newline at end of file