add ImageListReplace

2024-03-26 00:29:21 +08:00
parent d938adf3cc
commit f3866ede89
6 changed files with 342 additions and 38 deletions
@@ -603,7 +603,7 @@ from .nodes.Mask import MaskListReplace,MaskListMerge,OutlineMask,FeatheredMask

 from .nodes.Style import ApplyVisualStylePrompting,StyleAlignedReferenceSampler,StyleAlignedBatchAlign,StyleAlignedSampleReferenceLatents

-from .nodes.Video import LoadVideoAndSegment
+from .nodes.Video import LoadVideoAndSegment,ImageListReplace


 # 要导出的所有节点及其名称的字典
@@ -682,7 +682,8 @@ NODE_CLASS_MAPPINGS = {
    "StyleAlignedBatchAlign_": StyleAlignedBatchAlign,
    "LoadVideoAndSegment_":LoadVideoAndSegment,
    "ListSplit_":ListSplit,
-    "MaskListReplace_":MaskListReplace
+    "MaskListReplace_":MaskListReplace,
+    "ImageListReplace_":ImageListReplace
    # "LaMaInpainting":LaMaInpainting
    # "GamePal":GamePal
 }
@@ -726,6 +727,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "MaskListMerge_":"MaskList to Mask",
    "ListSplit_":"Split List",
    "MaskListReplace_":"MaskList Replace",
+    "ImageListReplace_":"ImageList Replace",
    "SwitchByIndex":"List Switch By Index",
    "GLIGENTextBoxApply_Advanced":"GLIGEN TextBox Apply ♾️Mixlab",
    "GridDisplayAndSave":"Grid Display And Save",
@@ -108,32 +108,32 @@ class MaskListReplace:
            "mask_replace": ("MASK",),
            "start_index":("INT", {"default": 0, "min": 0, "step": 1}),
            "end_index":("INT", {"default": 0, "min": 0, "step": 1}),
-            "reverse": ("BOOLEAN", {"default": False}),
+            "invert": ("BOOLEAN", {"default": False}),
        }
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "run"
-    CATEGORY = "♾️Mixlab/Mask"
+    CATEGORY = "♾️Mixlab/Video"

    INPUT_IS_LIST = True
    OUTPUT_IS_LIST = (True,)

-    def run(self, masks,mask_replace,start_index,end_index,reverse):
+    def run(self, masks,mask_replace,start_index,end_index,invert):
        mask_replace=mask_replace[0]
        start_index=start_index[0]
        end_index=end_index[0]
-        reverse=reverse[0]
+        invert=invert[0]
        
        new_masks=[]
        for i in range(len(masks)):
            if i>=start_index and i<=end_index:
-                if reverse:
+                if invert:
                    new_masks.append(masks[i])
                else:
                    new_masks.append(mask_replace)
            else:
-                if reverse:
+                if invert:
                    new_masks.append(mask_replace)
                else:
                    new_masks.append(masks[i])
@@ -34,6 +34,48 @@ if ffmpeg_path is None:
    except:
        print("ffmpeg could not be found. Outputs that require it have been disabled")

+# Tensor to PIL
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+# Convert PIL to Tensor
+def pil2tensor(image):
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+def count_files(directory):
+    count = 0
+    for root, dirs, files in os.walk(directory):
+        count += len(files)
+    return count
+
+def create_temp_file(image):
+    output_dir = folder_paths.get_temp_directory()
+
+    c=count_files(output_dir)
+
+    (
+            full_output_folder,
+            filename,
+            counter,
+            subfolder,
+            _,
+        ) = folder_paths.get_save_image_path('temp_', output_dir)
+
+    
+    image=tensor2pil(image)
+ 
+    image_file = f"{filename}_{c}_{counter:05}.png"
+     
+    image_path=os.path.join(full_output_folder, image_file)
+
+    image.save(image_path,compress_level=4)
+
+    return [{
+                "filename": image_file,
+                "subfolder": subfolder,
+                "type": "temp"
+                }]
+

 def split_list(lst, chunk_size, transition_size):
    result = []
@@ -50,6 +92,59 @@ def split_list(lst, chunk_size, transition_size):
 # result = split_list(images, chunk_size, transition_size)
 # print(result)

+class ImageListReplace:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "images": ("IMAGE",),
+            "image_replace": ("IMAGE",),
+            "start_index":("INT", {"default": 0, "min": 0, "step": 1}),
+            "end_index":("INT", {"default": 0, "min": 0, "step": 1}),
+            "invert": ("BOOLEAN", {"default": False}),
+        }
+        }
+
+    RETURN_TYPES = ("IMAGE","IMAGE",)
+    RETURN_NAMES = ("images","select_images",)
+    FUNCTION = "run"
+    CATEGORY = "♾️Mixlab/Video"
+
+    OUTPUT_NODE = True
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True,True,)
+
+    def run(self, images,image_replace,start_index,end_index,invert):
+        image_replace=image_replace[0]
+        start_index=start_index[0]
+        end_index=end_index[0]
+        invert=invert[0]
+        
+        new_images=[]
+        select_images=[]
+        for i in range(len(images)):
+            if i>=start_index and i<=end_index:
+                if invert:
+                    new_images.append(images[i])
+                else:
+                    new_images.append(image_replace)
+                    select_images.append(images[i])
+            else:
+                if invert:
+                    new_images.append(image_replace)
+                    select_images.append(images[i])
+                else:
+                    new_images.append(images[i])
+
+        imss=[]
+        # print(len(images))
+        for i in range(len(images)):
+            ims=create_temp_file(images[i][0])
+            imss.append(ims[0])
+
+        image_replace=create_temp_file(image_replace)
+
+        return {"ui":{"_images": imss,"_image_replace":image_replace},"result": (new_images,select_images,)}
+

 class LoadVideoAndSegment:
    @classmethod
@@ -70,11 +165,11 @@ class LoadVideoAndSegment:

    CATEGORY = "♾️Mixlab/Video"

-    RETURN_TYPES = ("IMAGE", "INT",)
-    RETURN_NAMES = ("image_batch", "frame_count",)
+    RETURN_TYPES = ("IMAGE", "INT","INT",)
+    RETURN_NAMES = ("image_batch", "frame_count","segment_count",)
    FUNCTION = "load_video"
    OUTPUT_NODE = True
-    OUTPUT_IS_LIST = (True,False,)
+    OUTPUT_IS_LIST = (True,False,False,)


    def is_gif(self, filename):
@@ -194,7 +289,7 @@ class LoadVideoAndSegment:

        # images = torch.from_numpy(np.stack(images))

-        return (imgs, len(imgs))
+        return (imgs, len(images),len(imgs),)

    @classmethod
    def IS_CHANGED(s, video, **kwargs):
@@ -202,7 +202,7 @@
            width: fit-content;
            max-width: 100%;
            margin-left: 12px;
-            min-height: 200px;
+            /*min-height: 200px;*/
        }

        .input_card {
@@ -2,12 +2,92 @@ import { app } from '../../../scripts/app.js'
 import { api } from '../../../scripts/api.js'
 import { ComfyWidgets } from '../../../scripts/widgets.js'

+import { $el } from '../../../scripts/ui.js'
+
+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 * 0.3 - MARGIN * 2}px`,
+    // background: '#EEEEEE',
+    display: 'flex',
+    flexDirection: 'column',
+    // alignItems: 'center',
+    justifyContent: 'space-around'
+  }
+}
+
 function videoUpload (node, inputName, inputData, app) {
  const imageWidget = node.widgets.find(w => w.name === 'video')
  let uploadWidget

+  const widget = {
+    type: 'div',
+    name: 'upload-preview',
+    draw (ctx, node, widget_width, y, widget_height) {
+      Object.assign(
+        this.div.style,
+        get_position_style(ctx, widget_width, 188, node.size[1]),
+        {
+          outline: '1px solid'
+        }
+      )
+    }
+  }
+
+  widget.div = $el('div', {})
+  widget.div.style.width = `120px`
+  document.body.appendChild(widget.div)
+  node.addCustomWidget(widget)
+  // console.log('#imageWidget', imageWidget)
  const displayDiv = document.createElement('video')
-  console.log('imageWidget', node)
+  displayDiv.controls = true
+  // displayDiv.style=`width:200px;height:200px`
+  imageWidget.callback = () => {
+    displayDiv.src = `/view?filename=${
+      imageWidget.value
+    }&type=input&subfolder=${''}&rand=${Math.random()}`
+
+    // displayDiv.onloadedmetadata = function () {
+    //   var frameCount = displayDiv.duration * displayDiv.webkitDecodedFrameCount
+    //   console.log('视频帧数：' + frameCount)
+    //   node.widgets.filter(w => w.name == 'video_segment_frames')[0].value =
+    //     frameCount
+    // }
+  }
+
+  if (imageWidget.value) {
+    displayDiv.src = `/view?filename=${
+      imageWidget.value
+    }&type=input&subfolder=${''}&rand=${Math.random()}`
+  }
+
+  widget.div.appendChild(displayDiv)
+
+  const onRemoved = node.onRemoved
+  node.onRemoved = () => {
+    widget.div.remove()
+    return onRemoved?.()
+  }

  var default_value = imageWidget.value
  Object.defineProperty(imageWidget, 'value', {
@@ -62,6 +142,10 @@ function videoUpload (node, inputName, inputData, app) {
        if (updateNode) {
          imageWidget.value = path
        }
+
+        return `/view?filename=${path}&type=input&subfolder=${
+          pasted ? 'pasted' : ''
+        }&rand=${Math.random()}`
      } else {
        alert(resp.status + ' - ' + resp.statusText)
      }
@@ -78,8 +162,23 @@ function videoUpload (node, inputName, inputData, app) {
    onchange: async () => {
      if (fileInput.files.length) {
        let file = fileInput.files[0]
-        console.log(file)
-        await uploadFile(file, true)
+
+        const url = await uploadFile(file, true)
+
+        // console.log('fileInput', file)
+        var reader = new FileReader()
+        reader.onload = function () {
+          displayDiv.src = url
+          displayDiv.onloadedmetadata = function () {
+            // var frameCount =
+            //   displayDiv.duration * displayDiv.webkitDecodedFrameCount
+            // console.log('视频帧数：' + frameCount)
+            // node.widgets.filter(
+            //   w => w.name == 'video_segment_frames'
+            // )[0].value = frameCount
+          }
+        }
+        reader.readAsDataURL(file)
      }
    }
  })
@@ -94,33 +193,140 @@ function videoUpload (node, inputName, inputData, app) {
 }
 ComfyWidgets.VIDEOUPLOAD_ = videoUpload

-
 app.registerExtension({
  name: 'Mixlab.Video.LoadVideoAndSegment_',
  async beforeRegisterNodeDef (nodeType, nodeData, app) {
    if (nodeData?.name == 'LoadVideoAndSegment_') {
-      nodeData.input.required.upload = ['VIDEOUPLOAD_'];
-
-    //   const onExecuted = nodeType.prototype.onExecuted
-    //   nodeType.prototype.onExecuted = function (message) {
-    //     onExecuted?.apply(this, arguments)
-    //     console.log(message)
-        
-    //     // try {
-    //     //   let a = this.widgets.filter(w => w.name === 'AppInfoRun')[0]
-    //     //   if (a) {
-    //     //     if (!a.value) a.value = 0
-    //     //     a.value += 1
-    //     //   }
-
-    //     //   const div = this.widgets.filter(w => w.div)[0].div
-    //     //   Array.from(
-    //     //     div.querySelectorAll('button'),
-    //     //     b => (b.style.background = 'yellow')
-    //     //   )
-    //     // } catch (error) {}
-    //   }
+      nodeData.input.required.upload = ['VIDEOUPLOAD_']
+    }
+  }
+})

+app.registerExtension({
+  name: 'Mixlab.Video.ImageListReplace',
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeData?.name == 'ImageListReplace_') {
+      const orig_nodeCreated = nodeType.prototype.onNodeCreated
+      nodeType.prototype.onNodeCreated = function () {
+        orig_nodeCreated?.apply(this, arguments)
+        const widget = {
+          type: 'div',
+          name: 'preview',
+          draw (ctx, node, widget_width, y, widget_height) {
+            Object.assign(
+              this.div.style,
+              get_position_style(ctx, widget_width, 188, node.size[1]),
+              {
+                outline: '1px solid',
+                display: 'flex',
+                flexWrap: 'wrap',
+                flexDirection: 'row',
+                justifyContent: 'flex-start'
+              }
+            )
+          }
+        }
+
+        widget.div = $el('div', {})
+        widget.div.style.width = `120px`
+        document.body.appendChild(widget.div)
+        this.addCustomWidget(widget)
+        // console.log('#ImageListReplace', widget)
+
+        const onRemoved = this.onRemoved
+        this.onRemoved = () => {
+          widget.div.remove()
+          return onRemoved?.()
+        }
+      }
+
+      const onExecuted = nodeType.prototype.onExecuted
+      nodeType.prototype.onExecuted = function (message) {
+        onExecuted?.apply(this, arguments)
+        // console.log('#ImageListReplace', message._images)
+
+        let _image_replace = message._image_replace[0]
+        _image_replace = `/view?filename=${_image_replace.filename}&type=${
+          _image_replace.type
+        }&subfolder=${_image_replace.subfolder}&rand=${Math.random()}`
+
+        let preview = this.widgets.filter(w => w.name == 'preview')[0]
+        preview.div.innerHTML = ''
+        for (const img_ of message._images) {
+          let img = new Image()
+          img.style = `width: 100px;
+          margin: 4px;`
+          img.src = `/view?filename=${img_.filename}&type=${
+            img_.type
+          }&subfolder=${img_.subfolder}&rand=${Math.random()}`
+          preview.div.appendChild(img)
+        }
+
+        let start_index = this.widgets.filter(w => w.name == 'start_index')[0]
+        let end_index = this.widgets.filter(w => w.name == 'end_index')[0]
+        let invert = this.widgets.filter(w => w.name == 'invert')[0]
+        let _sc = start_index.callback.bind(start_index)
+        let _ec = end_index.callback.bind(end_index)
+
+        const selectImages = () => {
+          // console.log(v)
+          let s = start_index.value,
+            e = end_index.value
+          let imgs = preview.div.querySelectorAll('img')
+          for (let index = 0; index < imgs.length; index++) {
+            if (invert.value) {
+              imgs[index].style.outline =
+                index >= s && index <= e ? 'none' : '4px solid #cbd3fe'
+            } else {
+              imgs[index].style.outline =
+                index >= s && index <= e ? '4px solid #cbd3fe' : 'none'
+            }
+          }
+        }
+
+        selectImages()
+
+        start_index.callback = v => {
+          let s = v,
+            e = end_index.value
+          let imgs = preview.div.querySelectorAll('img')
+          for (let index = 0; index < imgs.length; index++) {
+            if (invert.value) {
+              imgs[index].style.outline =
+                index >= s && index <= e ? 'none' : '4px solid #cbd3fe'
+            } else {
+              imgs[index].style.outline =
+                index >= s && index <= e ? '4px solid #cbd3fe' : 'none'
+            }
+          }
+
+          _sc(v)
+        }
+
+        end_index.callback = v => {
+          let s = start_index.value,
+            e = v
+          let imgs = preview.div.querySelectorAll('img')
+          for (let index = 0; index < imgs.length; index++) {
+            if (invert.value) {
+              imgs[index].style.outline =
+                index >= s && index <= e ? 'none' : '4px solid #cbd3fe'
+            } else {
+              imgs[index].style.outline =
+                index >= s && index <= e ? '4px solid #cbd3fe' : 'none'
+            }
+          }
+
+          _ec(v)
+        }
+
+        invert.callback = v => {
+          selectImages()
+        }
+
+        try {
+        } catch (error) {}
+      }
    }
  }
 })
@@ -83,6 +83,7 @@
    left: 6.5vh;
    z-index: 2;
    border: none;
+    display: none;
    box-shadow: inset 0 -0.5vh rgba(0, 0, 0, 0.1);
 }