更新

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

ModelNode.py

@@ -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,)
-
-
-

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

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

__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
 }

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
-      }
-    }
   }
 }