diff --git a/ImageNode.py b/ImageNode.py index 619d1f1..c2f525b 100644 --- a/ImageNode.py +++ b/ImageNode.py @@ -132,7 +132,7 @@ def doMask(image,mask,save_image=False,filename_prefix="Mixlab",invert="yes",sav -def load_image(fp): +def load_image(fp,white_bg=False): i = Image.open(fp) i = ImageOps.exif_transpose(i) image = i.convert("RGB") @@ -141,13 +141,17 @@ def load_image(fp): if 'A' in i.getbands(): mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 mask = 1. - torch.from_numpy(mask) + if white_bg==True: + nw = mask.unsqueeze(0).unsqueeze(-1).repeat(1, 1, 1, 3) + # 将mask的黑色部分对image进行白色处理 + image[nw == 1] = 1.0 else: mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") return (image,mask) # 获取图片s -def get_images_filepath(f): +def get_images_filepath(f,white_bg=False): images = [] if os.path.isdir(f): @@ -155,7 +159,7 @@ def get_images_filepath(f): for file in files: file_path = os.path.join(root, file) try: - (im,mask)=load_image(file_path) + (im,mask)=load_image(file_path,white_bg) images.append({ "image":im, "mask":mask, @@ -166,14 +170,14 @@ def get_images_filepath(f): elif os.path.isfile(f): try: - (im,mask)=load_image(f) + (im,mask)=load_image(f,white_bg) images.append({ "image":im, "mask":mask, "file_path":f }) except: - print('非图片',file_path) + print('非图片',f) else: print('路径不存在或无效',f) @@ -230,7 +234,8 @@ class SmoothMask: def run(self,mask,smoothness): # result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) # print(result.shape) - mask=mask.numpy() + if hasattr(mask,'numpy'): + mask=mask.numpy() images=[] for i in range(mask.shape[0]): m=mask[i] @@ -290,7 +295,10 @@ class FeatheredMask: if start_offset>0: mask = 1.0 - mask - image_np=mask.numpy() + if hasattr(mask,'numpy'): + image_np=mask.numpy() + else: + image_np=mask image = np.uint8(image_np * 255) # image = cv2.cvtColor(image_cv) @@ -449,6 +457,7 @@ class LoadImagesFromPath: "file_path": ("STRING",{"multiline": False,"default": ""}) }, "optional":{ + "white_bg": (["disable","enable"],), "newest_files": (["enable", "disable"],), "index_variable":("INT", { "default": -1, @@ -471,9 +480,9 @@ class LoadImagesFromPath: OUTPUT_IS_LIST = (True,True,) # 运行的函数 - def run(self,file_path,newest_files,index_variable,seed): - - images=get_images_filepath(file_path) + def run(self,file_path,white_bg,newest_files,index_variable,seed): + print(file_path) + images=get_images_filepath(file_path,white_bg=='enable') # 排序 sorted_files = sorted(images, key=lambda x: os.path.getmtime(x['file_path']), reverse=(newest_files=='enable')) diff --git a/ModelNode.py b/ModelNode.py deleted file mode 100644 index 64807a3..0000000 --- a/ModelNode.py +++ /dev/null @@ -1,72 +0,0 @@ -# 一些有用的模型 -# https://huggingface.co/kandinsky-community/kandinsky-2-1 - -from diffusers import AutoPipelineForText2Image -import torch -import os - - -directory = os.path.dirname(__file__) - - -class KandinskyModelLoad: - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "file_path": ("STRING",{"multiline": True,"default":os.path.join(directory,'model\kandinsky-2-1')}), - "image": ("IMAGE",), - } - } - - RETURN_TYPES = ('IMAGE',) - - FUNCTION = "run" - - # OUTPUT_IS_LIST = (True,) - - CATEGORY = "Mixlab/model" - - # 运行的函数 - def run(self,file_path,image): - print('#file_path',file_path) - - return (image,) - - - -class KandinskyModel: - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "file_path": ("STRING",{"multiline": False,"default":os.path.join(directory,'model\kandinsky-2-1')}), - "prompt": ("STRING",{"multiline": True,"default": "A alien cheeseburger creature eating itself, claymation, cinematic, moody lighting"}), - "negative_prompt ": ("STRING",{"multiline": True,"default": "low quality, bad quality"}) - } - } - - RETURN_TYPES = ('IMAGE',) - - FUNCTION = "run" - - OUTPUT_IS_LIST = (True,) - - CATEGORY = "Mixlab/model" - - # 运行的函数 - def run(self,file_path,prompt,negative_prompt): - print(file_path,prompt,negative_prompt) - pipe = AutoPipelineForText2Image.from_pretrained(file_path, - torch_dtype=torch.float16) - pipe.enable_model_cpu_offload() - - image = pipe(prompt=prompt, negative_prompt=negative_prompt, - prior_guidance_scale =1.0, height=768, width=768).images[0] - # image.save("cheeseburger_monster.png") - return (image,) - - - diff --git a/README.md b/README.md index bf76527..71198f9 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,9 @@ ## In progress. +！！ +> Now comfyui supports capturing screen pixel streams from any software and can be used for LCM-Lora integration. Let's get started with implementation and design! 💻🌐 + +https://github.com/shadowcz007/comfyui-mixlab-nodes/assets/12645064/3167aed0-cea0-41f2-9075-b05e0ed08536 ## Installation diff --git a/ScreenShareNode.py b/ScreenShareNode.py new file mode 100644 index 0000000..eb3f32e --- /dev/null +++ b/ScreenShareNode.py @@ -0,0 +1,53 @@ +import os,io +from PIL import Image, ImageOps +import numpy as np +import torch +import folder_paths + + +def load_image(fp,white_bg=False): + i = Image.open(fp) + image = i.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + if white_bg==True: + nw = mask.unsqueeze(0).unsqueeze(-1).repeat(1, 1, 1, 3) + # 将mask的黑色部分对image进行白色处理 + image[nw == 1] = 1.0 + else: + mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") + return (image,mask) + + + +class ScreenShareNode: + @classmethod + def INPUT_TYPES(s): + return { "required":{ + "image_path": ("CHEESE",) + }, + "optional":{ + "seed": ("INT", {"default": 1, "min": 0, "max": 0xffffffffffffffff}), + } } + + RETURN_TYPES = ('IMAGE','MASK') + + FUNCTION = "run" + + CATEGORY = "Mixlab/image" + + # INPUT_IS_LIST = True + OUTPUT_IS_LIST = (False,False,) + + # 运行的函数 + def run(self,image_path,seed): + np=os.path.join(folder_paths.get_temp_directory(),image_path) + print(seed,np) + + (im,mask)=load_image(np) + + return (im,mask) + diff --git a/__init__.py b/__init__.py index c5ad0c0..6795c6d 100644 --- a/__init__.py +++ b/__init__.py @@ -29,7 +29,7 @@ def is_installed(package, package_overwrite=None): from .PromptNode import RandomPrompt,RunWorkflow from .ImageNode import TransparentImage,LoadImagesFromPath,SmoothMask,FeatheredMask,SplitLongMask,ImagesCrop from .Vae import VAELoader,VAEDecode -# from .ModelNode import KandinskyModel,KandinskyModelLoad +from .ScreenShareNode import ScreenShareNode # 要导出的所有节点及其名称的字典 # 注意：名称应全局唯一 @@ -42,8 +42,8 @@ NODE_CLASS_MAPPINGS = { "SmoothMask":SmoothMask, "ImagesCrop":ImagesCrop, "VAELoaderConsistencyDecoder":VAELoader, - "VAEDecodeConsistencyDecoder":VAEDecode - # "RunWorkflow":RunWorkflow + "VAEDecodeConsistencyDecoder":VAEDecode, + "ScreenShare":ScreenShareNode # "KandinskyModelLoad":KandinskyModelLoad, # "KandinskyModel":KandinskyModel } diff --git a/web/javascript/main.js b/web/javascript/main.js index cf9a5b2..3720076 100644 --- a/web/javascript/main.js +++ b/web/javascript/main.js @@ -1,6 +1,217 @@ import { app } from '../../../scripts/app.js' import { api } from '../../../scripts/api.js' import { ComfyWidgets } from '../../../scripts/widgets.js' +import { $el } from '../../../scripts/ui.js' + +let api_host = '127.0.0.1:8188' +let api_base = '' +let url = `http://${api_host}${api_base}` + +async function getQueue () { + try { + const res = await fetch(`${url}/queue`) + const data = await res.json() + // console.log(data.queue_running,data.queue_pending) + return { + // Running action uses a different endpoint for cancelling + Running: data.queue_running.length, + Pending:data.queue_pending.length + } + } catch (error) { + console.error(error) + return { Running:0, Pending:0 } + } +} + +async function uploadFile (file) { + try { + const body = new FormData() + body.append('image', file) + body.append('overwrite', 'true') + body.append('type', 'temp') + + const resp = await fetch(`${url}/upload/image`, { + method: 'POST', + body + }) + + if (resp.status === 200) { + const data = await resp.json() + let path = data.name + if (data.subfolder) path = data.subfolder + '/' + path + return path + } else { + alert(resp.status + ' - ' + resp.statusText) + } + } catch (error) { + alert(error) + } +} + +async function shareScreenAndUpload (imgElement) { + try { + let webcamVideo = document.createElement('video') + const mediaStream = await navigator.mediaDevices.getDisplayMedia({ + video: true + }) + + webcamVideo.removeEventListener('timeupdate', videoTimeUpdateHandler) + webcamVideo.srcObject = mediaStream + webcamVideo.onloadedmetadata = () => { + webcamVideo.play() + webcamVideo.addEventListener('timeupdate', videoTimeUpdateHandler) + } + + async function videoTimeUpdateHandler () { + if (window._mixlab_screen_time) { + console.log('loading') + return + }; + + const {Pending}=await getQueue(); + if(Pending<5) document.querySelector('#queue-button').click(); + + const videoW = webcamVideo.videoWidth + const videoH = webcamVideo.videoHeight + const aspectRatio = videoW / videoH + const WIDTH = 512, + HEIGHT = Math.round(WIDTH / aspectRatio) + const canvas = new OffscreenCanvas(WIDTH, HEIGHT) + + const ctx = canvas.getContext('2d') + ctx.drawImage(webcamVideo, 0, 0, videoW, videoH, 0, 0, WIDTH, HEIGHT) + + const blob = await canvas.convertToBlob({ + type: 'image/jpeg', + quality: 1 + }) + + var reader = new FileReader() + reader.onload = function (event) { + // console.log(imgElement) + imgElement.src = event.target.result + // console.log(event.target.result) + } // data url! + var source = reader.readAsDataURL(blob) + + const file = new File([blob], `screenshot_mixlab.jpeg`) + window._mixlab_screen_time = true + window._mixlab_screen_imagePath = await uploadFile(file) + window._mixlab_screen_time = false + } + + // window._mixlab_screen_time = setInterval(() => { + // context.drawImage(videoTrack, 0, 0, canvas.width, canvas.height) + // }, 300) + } catch (error) { + alert('Error accessing screen stream: ' + error) + } +} + +/* +A method that returns the required style for the html +*/ +function get_position_style (ctx, widget_width, y, node_height) { + const MARGIN = 4 // the margin around the html element + + /* Create a transform that deals with all the scrolling and zooming */ + const elRect = ctx.canvas.getBoundingClientRect() + const transform = new DOMMatrix() + .scaleSelf( + elRect.width / ctx.canvas.width, + elRect.height / ctx.canvas.height + ) + .multiplySelf(ctx.getTransform()) + .translateSelf(MARGIN, MARGIN + y) + + return { + transformOrigin: '0 0', + transform: transform, + left: `0`, + top: `0`, + cursor: 'pointer', + position: 'absolute', + maxWidth: `${widget_width - MARGIN * 2}px`, + maxHeight: `${node_height - MARGIN * 2}px`, // we're assuming we have the whole height of the node + width: `${widget_width - MARGIN * 2}px`, + height: `${node_height - MARGIN * 2}px` + } +} + +app.registerExtension({ + name: 'Mixlab.image.ScreenShareNode', + getCustomWidgets (app) { + return { + CHEESE (node, inputName, inputData, app) { + // We return an object containing a field CHEESE which has a function (taking node, name, data, app) + const widget = { + type: inputData[0], // the type, CHEESE + name: inputName, // the name, slice + size: [128, 72], // a default size + draw (ctx, node, width, y) { + // a method to draw the widget (ctx is a CanvasRenderingContext2D) + }, + computeSize (...args) { + return [128, 72] // a method to compute the current size of the widget + }, + async serializeValue (nodeId, widgetIndex) { + return window._mixlab_screen_imagePath + } + } + // widget.something = something; // maybe adds stuff to it + node.addCustomWidget(widget) // adds it to the node + return widget // and returns it. + } + } + }, + async beforeRegisterNodeDef (nodeType, nodeData, app) { + if (nodeType.comfyClass == 'ScreenShare') { + /* + Hijack the onNodeCreated call to add our widget + */ + const orig_nodeCreated = nodeType.prototype.onNodeCreated + nodeType.prototype.onNodeCreated = function () { + orig_nodeCreated?.apply(this, arguments) + + const widget = { + type: 'HTML', // whatever + name: 'flying', // whatever + draw (ctx, node, widget_width, y, widget_height) { + Object.assign( + this.inputEl.style, + get_position_style(ctx, widget_width, y, node.size[1]) + ) // assign the required style when we are drawn + } + } + + /* + Create an html element and add it to the document. + Look at $el in ui.js for all the options here + */ + widget.inputEl = $el('img', { + src: 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAwAAAAMCAYAAABWdVznAAAAAXNSR0IArs4c6QAAALZJREFUKFOFkLERwjAQBPdbgBkInECGaMLUQDsE0AkRVRAYWqAByxldPPOWHwnw4OBGye1p50UDSoA+W2ABLPN7i+C5dyC6R/uiAUXRQCs0bXoNIu4QPQzAxDKxHoALOrZcqtiyR/T6CXw7+3IGHhkYcy6BOR2izwT8LptG8rbMiCRAUb+CQ6WzQVb0SNOi5Z2/nX35DRyb/ENazhpWKoGwrpD6nICp5c2qogc4of+c7QcrhgF4Aa/aoAFHiL+RAAAAAElFTkSuQmCC' + }) + // widget.inputEl = $el('button', { + // innerText: 'Start' + // }) + document.body.appendChild(widget.inputEl) + widget.inputEl.addEventListener('click', () => { + shareScreenAndUpload(widget.inputEl) + }) + // console.log('widget.inputEl',widget.inputEl) + + /* + Add the widget, make sure we clean up nicely, and we do not want to be serialized! + */ + this.addCustomWidget(widget) + this.onRemoved = function () { + widget.inputEl.remove() + } + this.serialize_widgets = false + } + } + } +}) // 和python实现一样 function run (mutable_prompt, immutable_prompt) { @@ -321,11 +532,11 @@ const node = { '[logging]', 'add custom node definitions', 'current nodes:', - defs + defs ) // 在这里进行 语言切换 for (const nodeName in defs) { - if(nodeName==='RandomPrompt'){ + if (nodeName === 'RandomPrompt') { // defs[nodeName].category // defs[nodeName].display_name } @@ -374,6 +585,54 @@ const node = { } } + if (nodeData.name === 'WSServer') { + // Create the button widget for selecting the files + // node.addWidget( + // 'button', + // 'choose file to upload', + // 'video', + // () => { + // console.log('click') + // } + // ) + // uploadWidget.serialize = false + // const onExecuted = nodeType.prototype.onExecuted + // nodeType.prototype.onExecuted = function (message) { + // const r = onExecuted?.apply?.(this, arguments) + // console.log('executed', message) + // const upload = this.widgets.filter(w => w.name === 'upload')[0] + // console.log('executed', this.widgets) + // // navigator.mediaDevices + // // .getDisplayMedia({ video: true }) + // // .then(stream => { + // // const videoElement = document.createElement('video') + // // videoElement.srcObject = stream + // // videoElement.autoplay = true + // // const canvasElement = document.createElement('canvas') + // // const context = canvasElement.getContext('2d') + // // videoElement.addEventListener('loadedmetadata', () => { + // // canvasElement.width = videoElement.videoWidth + // // canvasElement.height = videoElement.videoHeight + // // setInterval(async () => { + // // context.drawImage( + // // videoElement, + // // 0, + // // 0, + // // canvasElement.width, + // // canvasElement.height + // // ) + // // const imageData = canvasElement.toDataURL() + // // upload.value = await uploadScreenshot(imageData) + // // }, 200) + // // }) + // // }) + // // .catch(error => { + // // console.error('Error getting screen share:', error) + // // }) + // return r + // } + } + if (nodeData.name === 'RandomPrompt') { const onExecuted = nodeType.prototype.onExecuted nodeType.prototype.onExecuted = function (message) { @@ -409,15 +668,6 @@ const node = { return r } } - - if (nodeData.name === 'RunWorkflow') { - const onExecuted = nodeType.prototype.onExecuted - nodeType.prototype.onExecuted = function (message) { - const r = onExecuted?.apply?.(this, arguments) - - return r - } - } } }