chore:remove sim filter

this commit includes
-a stream-optimized node for detecting facial landmarks
-node for drawing the mesh on an image from the landmark data
-node for creating masks from the landmark data

.gitignore

@@ -0,0 +1,88 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+.pytest_cache/
+.coverage
+htmlcov/
+
+# Virtual Environment
+venv/
+env/
+ENV/
+.env/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# ML/AI specific
+*.pt
+*.pth
+*.onnx
+*.bin
+*.safetensors
+*.ckpt
+output/
+outputs/
+results/
+runs/
+logs/
+checkpoints/
+# IDE/Editor specific
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# OS specific
+.DS_Store
+Thumbs.db
+desktop.ini
+
+# ComfyUI specific
+temp/
+cache/
+input/
+output/
+
+# Media files that might be large
+*.jpg
+*.jpeg
+*.png
+*.gif
+*.bmp
+*.tif
+*.tiff
+*.mp4
+*.avi
+*.mov
+!examples/*.jpg
+!examples/*.png
+
+# Dependencies
+node_modules/
+
+# Logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*

__init__.py

@@ -0,0 +1,19 @@
+"""ComfyUI Stream Pack Nodes."""
+
+from .node_wrappers import feature_bank
+from .node_wrappers import facemesh
+from .node_wrappers import batched_denoising
+# Collect all NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS from submodules.
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
+# Import and update mappings from submodules.
+for module in [feature_bank, facemesh, batched_denoising]:
+    if hasattr(module, "NODE_CLASS_MAPPINGS"):
+        NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
+    if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"):
+        NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
+
+NODE_DISPLAY_NAME_MAPPINGS["StreamPack"] = "Stream Pack Nodes"
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

examples/workflows/sd15-tensorrt-api.json

@@ -0,0 +1,218 @@
+{
+  "1": {
+    "inputs": {
+      "image": "example.png",
+      "upload": "image"
+    },
+    "class_type": "LoadImage",
+    "_meta": {
+      "title": "Load Image"
+    }
+  },
+  "2": {
+    "inputs": {
+      "engine": "depth_anything_vitl14-fp16.engine",
+      "images": [
+        "1",
+        0
+      ]
+    },
+    "class_type": "DepthAnythingTensorrt",
+    "_meta": {
+      "title": "Depth Anything Tensorrt"
+    }
+  },
+  "3": {
+    "inputs": {
+      "unet_name": "static-dreamshaper8_SD15_$stat-b-1-h-512-w-512_00001_.engine",
+      "model_type": "SD15"
+    },
+    "class_type": "TensorRTLoader",
+    "_meta": {
+      "title": "TensorRT Loader"
+    }
+  },
+  "5": {
+    "inputs": {
+      "text": "the hulk",
+      "clip": [
+        "23",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "6": {
+    "inputs": {
+      "text": "",
+      "clip": [
+        "23",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "7": {
+    "inputs": {
+      "seed": 905056445574169,
+      "steps": 1,
+      "cfg": 1,
+      "sampler_name": "lcm",
+      "scheduler": "normal",
+      "denoise": 1,
+      "model": [
+        "3",
+        0
+      ],
+      "positive": [
+        "9",
+        0
+      ],
+      "negative": [
+        "9",
+        1
+      ],
+      "latent_image": [
+        "16",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "KSampler"
+    }
+  },
+  "8": {
+    "inputs": {
+      "control_net_name": "control_v11f1p_sd15_depth_fp16.safetensors"
+    },
+    "class_type": "ControlNetLoader",
+    "_meta": {
+      "title": "Load ControlNet Model"
+    }
+  },
+  "9": {
+    "inputs": {
+      "strength": 1,
+      "start_percent": 0,
+      "end_percent": 1,
+      "positive": [
+        "5",
+        0
+      ],
+      "negative": [
+        "6",
+        0
+      ],
+      "control_net": [
+        "10",
+        0
+      ],
+      "image": [
+        "2",
+        0
+      ]
+    },
+    "class_type": "ControlNetApplyAdvanced",
+    "_meta": {
+      "title": "Apply ControlNet"
+    }
+  },
+  "10": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": false,
+      "mode": "reduce-overhead",
+      "controlnet": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadControlNet",
+    "_meta": {
+      "title": "TorchCompileLoadControlNet"
+    }
+  },
+  "11": {
+    "inputs": {
+      "vae_name": "taesd"
+    },
+    "class_type": "VAELoader",
+    "_meta": {
+      "title": "Load VAE"
+    }
+  },
+  "13": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": true,
+      "mode": "reduce-overhead",
+      "compile_encoder": true,
+      "compile_decoder": true,
+      "vae": [
+        "11",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadVAE",
+    "_meta": {
+      "title": "TorchCompileLoadVAE"
+    }
+  },
+  "14": {
+    "inputs": {
+      "samples": [
+        "7",
+        0
+      ],
+      "vae": [
+        "13",
+        0
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "15": {
+    "inputs": {
+      "images": [
+        "14",
+        0
+      ]
+    },
+    "class_type": "PreviewImage",
+    "_meta": {
+      "title": "Preview Image"
+    }
+  },
+  "16": {
+    "inputs": {
+      "width": 512,
+      "height": 512,
+      "batch_size": 1
+    },
+    "class_type": "EmptyLatentImage",
+    "_meta": {
+      "title": "Empty Latent Image"
+    }
+  },
+  "23": {
+    "inputs": {
+      "clip_name": "CLIPText/model.fp16.safetensors",
+      "type": "stable_diffusion",
+      "device": "default"
+    },
+    "class_type": "CLIPLoader",
+    "_meta": {
+      "title": "Load CLIP"
+    }
+  }
+}

examples/workflows/sd15-tensorrt-batched-denoising.json

@@ -0,0 +1,376 @@
+{
+  "2": {
+    "inputs": {
+      "engine": "depth_anything_vitl14-fp16.engine",
+      "images": [
+        "26",
+        0
+      ]
+    },
+    "class_type": "DepthAnythingTensorrt",
+    "_meta": {
+      "title": "Depth Anything Tensorrt"
+    }
+  },
+  "3": {
+    "inputs": {
+      "unet_name": "static-dreamshaper8_SD15_$stat-b-1-h-512-w-512_00001_.engine",
+      "model_type": "SD15"
+    },
+    "class_type": "TensorRTLoader",
+    "_meta": {
+      "title": "TensorRT Loader"
+    }
+  },
+  "5": {
+    "inputs": {
+      "text": "beautiful pencil sketch, masterpiece ",
+      "clip": [
+        "18",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "6": {
+    "inputs": {
+      "text": "",
+      "clip": [
+        "18",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "8": {
+    "inputs": {
+      "control_net_name": "control_v11f1p_sd15_depth_fp16.safetensors"
+    },
+    "class_type": "ControlNetLoader",
+    "_meta": {
+      "title": "Load ControlNet Model"
+    }
+  },
+  "9": {
+    "inputs": {
+      "strength": 0.9,
+      "start_percent": 0,
+      "end_percent": 1,
+      "positive": [
+        "5",
+        0
+      ],
+      "negative": [
+        "6",
+        0
+      ],
+      "control_net": [
+        "10",
+        0
+      ],
+      "image": [
+        "2",
+        0
+      ]
+    },
+    "class_type": "ControlNetApplyAdvanced",
+    "_meta": {
+      "title": "Apply ControlNet"
+    }
+  },
+  "10": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": false,
+      "mode": "reduce-overhead",
+      "controlnet": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadControlNet",
+    "_meta": {
+      "title": "TorchCompileLoadControlNet"
+    }
+  },
+  "11": {
+    "inputs": {
+      "vae_name": "taesd"
+    },
+    "class_type": "VAELoader",
+    "_meta": {
+      "title": "Load VAE"
+    }
+  },
+  "13": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": true,
+      "mode": "reduce-overhead",
+      "compile_encoder": true,
+      "compile_decoder": true,
+      "vae": [
+        "11",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadVAE",
+    "_meta": {
+      "title": "TorchCompileLoadVAE"
+    }
+  },
+  "14": {
+    "inputs": {
+      "samples": [
+        "42",
+        0
+      ],
+      "vae": [
+        "13",
+        0
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "15": {
+    "inputs": {
+      "images": [
+        "14",
+        0
+      ]
+    },
+    "class_type": "PreviewImage",
+    "_meta": {
+      "title": "Preview Image"
+    }
+  },
+  "16": {
+    "inputs": {
+      "width": 512,
+      "height": 512,
+      "batch_size": 1
+    },
+    "class_type": "EmptyLatentImage",
+    "_meta": {
+      "title": "Empty Latent Image"
+    }
+  },
+  "18": {
+    "inputs": {
+      "stop_at_clip_layer": -2,
+      "clip": [
+        "36",
+        0
+      ]
+    },
+    "class_type": "CLIPSetLastLayer",
+    "_meta": {
+      "title": "CLIP Set Last Layer"
+    }
+  },
+  "19": {
+    "inputs": {
+      "strength": 0.9,
+      "start_percent": 0,
+      "end_percent": 1,
+      "positive": [
+        "5",
+        0
+      ],
+      "negative": [
+        "6",
+        0
+      ],
+      "control_net": [
+        "21",
+        0
+      ],
+      "image": [
+        "39",
+        1
+      ]
+    },
+    "class_type": "ControlNetApplyAdvanced",
+    "_meta": {
+      "title": "Apply ControlNet"
+    }
+  },
+  "20": {
+    "inputs": {
+      "control_net_name": "control_v2p_sd15_mediapipe_face.safetensors"
+    },
+    "class_type": "ControlNetLoader",
+    "_meta": {
+      "title": "Load ControlNet Model"
+    }
+  },
+  "21": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": false,
+      "mode": "reduce-overhead",
+      "controlnet": [
+        "20",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadControlNet",
+    "_meta": {
+      "title": "TorchCompileLoadControlNet"
+    }
+  },
+  "26": {
+    "inputs": {
+      "image": "example.png",
+      "upload": "image"
+    },
+    "class_type": "LoadImage",
+    "_meta": {
+      "title": "Load Image"
+    }
+  },
+  "36": {
+    "inputs": {
+      "clip_name": "CLIPText/model.fp16.safetensors",
+      "type": "stable_diffusion",
+      "device": "default"
+    },
+    "class_type": "CLIPLoader",
+    "_meta": {
+      "title": "Load CLIP"
+    }
+  },
+  "37": {
+    "inputs": {
+      "conditioning_to": [
+        "9",
+        1
+      ],
+      "conditioning_from": [
+        "19",
+        1
+      ]
+    },
+    "class_type": "ConditioningConcat",
+    "_meta": {
+      "title": "Conditioning (Concat)"
+    }
+  },
+  "38": {
+    "inputs": {
+      "conditioning_to": [
+        "9",
+        0
+      ],
+      "conditioning_from": [
+        "19",
+        0
+      ]
+    },
+    "class_type": "ConditioningConcat",
+    "_meta": {
+      "title": "Conditioning (Concat)"
+    }
+  },
+  "39": {
+    "inputs": {
+      "static_image_mode": false,
+      "max_num_faces": 1,
+      "refine_landmarks": false,
+      "min_detection_confidence": 0.35000000000000003,
+      "min_tracking_confidence": 0.35000000000000003,
+      "draw_mesh": true,
+      "draw_contours": true,
+      "mesh_color": "green",
+      "mesh_thickness": 2,
+      "image": [
+        "26",
+        0
+      ]
+    },
+    "class_type": "FaceMeshNode",
+    "_meta": {
+      "title": "Face Mesh"
+    }
+  },
+  "41": {
+    "inputs": {
+      "num_steps": 3
+    },
+    "class_type": "StreamBatchSampler",
+    "_meta": {
+      "title": "Stream Batch Sampler"
+    }
+  },
+  "42": {
+    "inputs": {
+      "add_noise": true,
+      "noise_seed": 977567268563782,
+      "cfg": 1,
+      "model": [
+        "3",
+        0
+      ],
+      "positive": [
+        "38",
+        0
+      ],
+      "negative": [
+        "37",
+        0
+      ],
+      "sampler": [
+        "41",
+        0
+      ],
+      "sigmas": [
+        "44",
+        0
+      ],
+      "latent_image": [
+        "43",
+        0
+      ]
+    },
+    "class_type": "SamplerCustom",
+    "_meta": {
+      "title": "SamplerCustom"
+    }
+  },
+  "43": {
+    "inputs": {
+      "buffer_size": 3,
+      "latent": [
+        "16",
+        0
+      ]
+    },
+    "class_type": "StreamFrameBuffer",
+    "_meta": {
+      "title": "Stream Frame Buffer"
+    }
+  },
+  "44": {
+    "inputs": {
+      "t_index_list": "0,16,32",
+      "num_inference_steps": 50,
+      "model": [
+        "3",
+        0
+      ]
+    },
+    "class_type": "StreamScheduler",
+    "_meta": {
+      "title": "Stream Scheduler"
+    }
+  }
+}

examples/workflows/sd15-tensorrt-feature-bank-api.json

@@ -0,0 +1,234 @@
+{
+  "1": {
+    "inputs": {
+      "image": "example.png"
+    },
+    "class_type": "LoadImage",
+    "_meta": {
+      "title": "Load Image"
+    }
+  },
+  "2": {
+    "inputs": {
+      "engine": "depth_anything_vitl14-fp16.engine",
+      "images": [
+        "1",
+        0
+      ]
+    },
+    "class_type": "DepthAnythingTensorrt",
+    "_meta": {
+      "title": "Depth Anything Tensorrt"
+    }
+  },
+  "3": {
+    "inputs": {
+      "unet_name": "static-dreamshaper8_SD15_$stat-b-1-h-512-w-512_00001_.engine",
+      "model_type": "SD15"
+    },
+    "class_type": "TensorRTLoader",
+    "_meta": {
+      "title": "TensorRT Loader"
+    }
+  },
+  "5": {
+    "inputs": {
+      "text": "the hulk",
+      "clip": [
+        "23",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "6": {
+    "inputs": {
+      "text": "",
+      "clip": [
+        "23",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "7": {
+    "inputs": {
+      "seed": 785664736216738,
+      "steps": 1,
+      "cfg": 1,
+      "sampler_name": "lcm",
+      "scheduler": "normal",
+      "denoise": 1,
+      "model": [
+        "24",
+        0
+      ],
+      "positive": [
+        "9",
+        0
+      ],
+      "negative": [
+        "9",
+        1
+      ],
+      "latent_image": [
+        "16",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "KSampler"
+    }
+  },
+  "8": {
+    "inputs": {
+      "control_net_name": "control_v11f1p_sd15_depth_fp16.safetensors"
+    },
+    "class_type": "ControlNetLoader",
+    "_meta": {
+      "title": "Load ControlNet Model"
+    }
+  },
+  "9": {
+    "inputs": {
+      "strength": 1,
+      "start_percent": 0,
+      "end_percent": 1,
+      "positive": [
+        "5",
+        0
+      ],
+      "negative": [
+        "6",
+        0
+      ],
+      "control_net": [
+        "10",
+        0
+      ],
+      "image": [
+        "2",
+        0
+      ]
+    },
+    "class_type": "ControlNetApplyAdvanced",
+    "_meta": {
+      "title": "Apply ControlNet"
+    }
+  },
+  "10": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": false,
+      "mode": "reduce-overhead",
+      "controlnet": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadControlNet",
+    "_meta": {
+      "title": "TorchCompileLoadControlNet"
+    }
+  },
+  "11": {
+    "inputs": {
+      "vae_name": "taesd"
+    },
+    "class_type": "VAELoader",
+    "_meta": {
+      "title": "Load VAE"
+    }
+  },
+  "13": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": true,
+      "mode": "reduce-overhead",
+      "compile_encoder": true,
+      "compile_decoder": true,
+      "vae": [
+        "11",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadVAE",
+    "_meta": {
+      "title": "TorchCompileLoadVAE"
+    }
+  },
+  "14": {
+    "inputs": {
+      "samples": [
+        "7",
+        0
+      ],
+      "vae": [
+        "13",
+        0
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "15": {
+    "inputs": {
+      "images": [
+        "14",
+        0
+      ]
+    },
+    "class_type": "PreviewImage",
+    "_meta": {
+      "title": "Preview Image"
+    }
+  },
+  "16": {
+    "inputs": {
+      "width": 512,
+      "height": 512,
+      "batch_size": 1
+    },
+    "class_type": "EmptyLatentImage",
+    "_meta": {
+      "title": "Empty Latent Image"
+    }
+  },
+  "23": {
+    "inputs": {
+      "clip_name": "CLIPText/model.fp16.safetensors",
+      "type": "stable_diffusion",
+      "device": "default"
+    },
+    "class_type": "CLIPLoader",
+    "_meta": {
+      "title": "Load CLIP"
+    }
+  },
+  "24": {
+    "inputs": {
+      "use_feature_injection": false,
+      "feature_injection_strength": 0.8,
+      "feature_similarity_threshold": 0.98,
+      "feature_cache_interval": 4,
+      "feature_bank_max_frames": 4,
+      "model": [
+        "3",
+        0
+      ]
+    },
+    "class_type": "FeatureBankAttentionProcessor",
+    "_meta": {
+      "title": "Feature Bank Attention Processor"
+    }
+  }
+}

node_wrappers/batched_denoising.py

@@ -0,0 +1,15 @@
+from ..src.batched_denoising.batched_denoising import StreamBatchSampler, StreamScheduler, StreamFrameBuffer
+
+NODE_CLASS_MAPPINGS = {
+    "StreamBatchSampler": StreamBatchSampler,
+    "StreamScheduler": StreamScheduler,
+    "StreamFrameBuffer": StreamFrameBuffer,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "StreamBatchSampler": "Stream Batch Sampler",
+    "StreamScheduler": "Stream Scheduler",
+    "StreamFrameBuffer": "Stream Frame Buffer",
+}
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

node_wrappers/facemesh.py

@@ -0,0 +1,13 @@
+from ..src.media_pipe.facemesh import *
+
+NODE_CLASS_MAPPINGS = {
+    "FaceMeshNode": FaceMeshNode,
+    "FaceMeshDrawNode": FaceMeshDrawNode,
+    "FaceMeshMaskNode": FaceMeshMaskNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FaceMeshNode": "Face Mesh",
+    "FaceMeshDrawNode": "Face Mesh Draw",
+    "FaceMeshMaskNode": "Face Region Mask"
+}

node_wrappers/feature_bank.py

@@ -0,0 +1,21 @@
+"""A ComfyUI native node implementation of the Feature Bank of the StreamV2V paper.
+
+
+References:
+    StreamV2V: A Deep Generative Model for High-fidelity View Synthesis - Liang et al.,
+    2025, URL: https://arxiv.org/abs/2405.15757
+"""
+
+from ..src.feature_bank.feature_bank import (
+    FeatureBankAttentionProcessorNode,
+)
+
+NODE_CLASS_MAPPINGS = {
+    "FeatureBankAttentionProcessor": FeatureBankAttentionProcessorNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FeatureBankAttentionProcessor": "Feature Bank Attention Processor"
+}
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

requirements.txt

@@ -0,0 +1,3 @@
+mediapipe
+torch
+diffusers

src/__init__.py

@@ -0,0 +1 @@
+# Dummy file ensuring this package will be recognized.

src/batched_denoising/batched_denoising.py

@@ -0,0 +1,842 @@
+import torch
+from .control_base import ControlNodeBase
+import comfy.model_management
+import comfy.samplers
+import math
+
+
+class StreamBatchSampler(ControlNodeBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        inputs = super().INPUT_TYPES()
+        inputs["required"].update(
+            {
+                "latent": ("LATENT",),
+                "num_timesteps": (
+                    "INT",
+                    {
+                        "default": 2,
+                        "min": 1,
+                        "max": 10,
+                        "step": 1,
+                        "tooltip": "Number of denoising steps to use. More steps = better quality but MUCH slower. StreamDiffusion's speed comes from using just 2 steps for img2img (32,45) or 4 for txt2img (0,16,32,45)",
+                    },
+                ),
+                "frame_buffer_size": (
+                    "INT",
+                    {
+                        "default": 1,
+                        "min": 1,
+                        "max": 10,
+                        "step": 1,
+                        "tooltip": "How many frames to process together for temporal consistency. Each frame is processed at each timestep (total batch = frame_buffer_size * num_timesteps). Higher values reduce flickering but use more VRAM and don't improve speed",
+                    },
+                ),
+                "cfg_type": (
+                    ["none", "full", "self", "initialize"],
+                    {
+                        "default": "self",
+                        "tooltip": "'self' is fastest and most memory efficient, 'full' is standard SD behavior but slower, 'initialize' is a middle ground, 'none' disables guidance",
+                    },
+                ),
+                "delta": (
+                    "FLOAT",
+                    {
+                        "default": 1.0,
+                        "min": 0.0,
+                        "max": 5.0,
+                        "step": 0.1,
+                        "tooltip": "Only used with 'self' CFG. Controls strength of self-guidance. Higher values = stronger guidance but more artifacts. Default 1.0 works well",
+                    },
+                ),
+                "similarity_threshold": (
+                    "FLOAT",
+                    {
+                        "default": 0.98,
+                        "min": 0.0,
+                        "max": 1.0,
+                        "step": 0.01,
+                        "tooltip": "Skip frames that are too similar to save compute. 0 = process all frames. Try 0.98 if you want to skip static scenes. Higher = more skipping",
+                        # TODO: see if this is backwards from original
+                    },
+                ),
+                "max_skip_frames": (
+                    "INT",
+                    {
+                        "default": 10,
+                        "min": 1,
+                        "max": 100,
+                        "step": 1,
+                        "tooltip": "Maximum frames to skip when using similarity filter. Prevents getting stuck on static scenes. Lower = more responsive to changes",
+                    },
+                ),
+            }
+        )
+        inputs["optional"] = {
+            "prompt": (
+                "STRING",
+                {
+                    "multiline": True,
+                    "default": "",
+                    "tooltip": "Optional text prompt. Enables KV caching for faster processing. Leave empty to disable. Same prompt across frames = faster",
+                },
+            ),
+        }
+        return inputs
+
+    RETURN_TYPES = (
+        "SAMPLER",
+        "LATENT",
+    )
+    FUNCTION = "update"
+    CATEGORY = "StreamPack/sampling"
+    DESCRIPTION = "Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps (StreamDiffusion)"
+
+    def __init__(self):
+        super().__init__()
+        self.frame_buffer_size = None
+        self.frame_buffer = None  # Will be a tensor instead of list
+        self.x_t_latent_buffer = None
+        self.stock_noise = None
+        self.cached_noise = {}  # Cache for different expanded versions
+        self.debug = False  # Debug flag for print statements
+        self.last_batch_shape = None  # Track last batch shape
+        # Add CFG state
+        self.cfg_type = "self"
+        self.guidance_scale = 1.2
+        self.delta = 1.0
+        # Add cross attention state
+        self.prompt = ""
+        self.last_prompt = None
+        self.kv_cache = None
+        self.cross_attention_hook = None
+        self.last_model_hash = None
+        # Add similarity filter state
+        self.similarity_threshold = 0.0
+        self.max_skip_frames = 10
+        self.similarity_filter = None
+        # Add buffer state
+        self.num_timesteps = None
+        # Add caching for alpha/beta and their square roots
+        self.cached_sigmas = None
+        self.alpha_prod_t = None
+        self.beta_prod_t = None
+        self.alpha_prod_t_sqrt = None
+        self.beta_prod_t_sqrt = None
+        # Add noise sequence and step counter
+        self.noise_sequence = None
+        self.step_counter = 0
+
+    def get_expanded_noise(self, batch_size, scale=1.0):
+        """Get expanded noise with caching based on batch size and scale"""
+        cache_key = (batch_size, scale)
+        if self.stock_noise is None:
+            return None
+
+        if cache_key not in self.cached_noise:
+            self.cached_noise[cache_key] = (
+                self.stock_noise.unsqueeze(0).expand(batch_size, -1, -1, -1) * scale
+            )
+            if self.debug:
+                print(
+                    f"[StreamBatchSampler] Cached expanded noise for batch_size={batch_size}, scale={scale}"
+                )
+
+        return self.cached_noise[cache_key]
+
+    def invalidate_noise_cache(self):
+        """Clear all cached noise when stock noise changes"""
+        self.cached_noise.clear()
+
+    def process_with_cfg(self, model, x, sigma_batch, extra_args):
+        """Apply classifier-free guidance based on current cfg_type"""
+        cond_scale = extra_args["model_options"].get("cond_scale", 1.0)
+
+        if self.cfg_type == "none" or cond_scale <= 1.0:
+            if self.debug:
+                print("[CFG] No guidance - running model directly")
+            return model(x, sigma_batch, **extra_args)
+
+        if self.cfg_type == "full":
+            if self.debug:
+                print("[CFG] Using full CFG")
+            # Double the batch for cond/uncond
+            x_double = torch.cat([x, x], dim=0)
+            sigma_double = torch.cat([sigma_batch, sigma_batch], dim=0)
+            model_output = model(x_double, sigma_double, **extra_args)
+            noise_pred_uncond, noise_pred_text = model_output.chunk(2)
+            # Combine predictions
+            eps_cfg = noise_pred_uncond + cond_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+            return eps_cfg
+
+        elif self.cfg_type == "initialize":
+            if self.debug:
+                print("[CFG] Using initialize CFG")
+            # Add single uncond at start
+            x_plus_uc = torch.cat([x[0:1], x], dim=0)
+            sigma_plus_uc = torch.cat([sigma_batch[0:1], sigma_batch], dim=0)
+            model_output = model(x_plus_uc, sigma_plus_uc, **extra_args)
+            noise_pred_uncond = model_output[0:1]
+            noise_pred_text = model_output[1:]
+            # Update stock noise with uncond prediction
+            self.stock_noise = noise_pred_uncond[0].clone()
+            self.invalidate_noise_cache()
+            return noise_pred_text  # Returning noise_pred_text as per original logic
+
+        else:  # self
+            if self.debug:
+                print("\n[CFG] Using RCFG (self)")
+                print(f"[CFG] Input shape: {x.shape}, Sigma shape: {sigma_batch.shape}")
+
+            # Get current model prediction
+            eps_c = model(x, sigma_batch, **extra_args)
+            if self.debug:
+                print(
+                    f"[CFG] Current prediction (eps_c) stats: min={eps_c.min():.3f}, max={eps_c.max():.3f}, mean={eps_c.mean():.3f}"
+                )
+
+            # Initialize stock noise if needed
+            if self.stock_noise is None:
+                if self.debug:
+                    print("[CFG] Initializing stock noise with current prediction")
+                self.stock_noise = eps_c[0].clone().detach()
+                self.invalidate_noise_cache()
+                return eps_c
+
+            # Apply RCFG equation
+            noise_pred_uncond = self.stock_noise.to(eps_c.device) * self.delta
+            if self.debug:
+                print(f"[CFG] Using delta={self.delta}, cond_scale={cond_scale}")
+                print(
+                    f"[CFG] Stock noise stats: min={self.stock_noise.min():.3f}, max={self.stock_noise.max():.3f}, mean={self.stock_noise.mean():.3f}"
+                )
+                print(
+                    f"[CFG] Uncond prediction stats: min={noise_pred_uncond.min():.3f}, max={noise_pred_uncond.max():.3f}, mean={noise_pred_uncond.mean():.3f}"
+                )
+
+            # Combine predictions exactly as StreamDiffusion does
+            eps_cfg = noise_pred_uncond + cond_scale * (eps_c - noise_pred_uncond)
+
+            if self.debug:
+                print(
+                    f"[CFG] Final output (eps_cfg) stats: min={eps_cfg.min():.3f}, max={eps_cfg.max():.3f}, mean={eps_cfg.mean():.3f}"
+                )
+                print(
+                    f"[CFG] Difference (eps_cfg - eps_c) stats: min={(eps_cfg-eps_c).min():.3f}, max={(eps_cfg-eps_c).max():.3f}, mean={(eps_cfg-eps_c).mean():.3f}\n"
+                )
+
+            return eps_cfg
+
+    def setup_cross_attention(self, model, prompt):
+        """Set up cross attention optimization for the model"""
+        if not prompt:  # Don't use optimization if no prompt provided
+            return model
+
+        def optimized_attention_forward(module, x, context=None, mask=None):
+            """Optimized cross attention with KV caching and ComfyUI's optimizations"""
+            batch_size = x.shape[0]
+
+            # Generate query from input
+            q = module.to_q(x)
+
+            # Check if we need to recompute KV cache
+            if self.kv_cache is None or self.last_prompt != prompt:
+                if self.debug:
+                    print("[StreamBatchSampler] Computing new KV cache")
+                k = module.to_k(context)
+                v = module.to_v(context)
+                self.kv_cache = (k, v)
+                self.last_prompt = prompt
+            else:
+                if self.debug:
+                    print("[StreamBatchSampler] Using cached KV pairs")
+                k, v = self.kv_cache
+
+            # Handle batch size expansion if needed
+            if k.shape[0] == 1 and batch_size > 1:
+                k = k.expand(batch_size, -1, -1)
+                v = v.expand(batch_size, -1, -1)
+
+            # Use ComfyUI's optimized attention
+            from comfy.ldm.modules.attention import optimized_attention
+
+            out = optimized_attention(q, k, v, module.heads, mask=mask)
+            return module.to_out[0](out)
+
+        # Get the actual model from the sampler
+        if hasattr(model, "model"):
+            unet = model.model
+        else:
+            return model
+
+        # Only set up hooks if prompt has changed
+        if self.last_prompt != prompt:
+            # Remove old hooks if they exist
+            if self.cross_attention_hook is not None:
+                self.cross_attention_hook.remove()
+
+            def create_cross_attn_patch(module, input, output):
+                """Hook to replace cross attention modules with optimized version"""
+                if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"):
+                    if not hasattr(module, "_original_forward"):
+                        module._original_forward = module.forward
+                    module.forward = (
+                        lambda *args, **kwargs: optimized_attention_forward(
+                            module, *args, **kwargs
+                        )
+                    )
+                return output
+
+            # Register hooks on cross attention blocks
+            def register_hooks(module):
+                if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"):
+                    self.cross_attention_hook = module.register_forward_hook(
+                        create_cross_attn_patch
+                    )
+
+            unet.apply(register_hooks)
+            self.last_prompt = prompt
+
+        return model
+
+    def compute_alpha_beta(self, sigmas):
+        """Pre-compute alpha and beta terms for the given sigmas (LCM-like using cosine schedule)"""
+        self.cached_sigmas = sigmas.clone()
+        print(f"\n[Alpha/Beta] Input sigmas: {sigmas}")
+
+        num_train_timesteps = 1000  # Standard for many diffusion models
+        timesteps = torch.linspace(
+            1, 0, num_train_timesteps, dtype=torch.float32, device=sigmas.device
+        )
+        alphas_cumprod = torch.cos((timesteps + 0.008) / 1.008 * math.pi / 2).pow(2)
+        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+
+        # Map sigmas to timesteps
+        scaled_sigmas = (
+            sigmas / sigmas.max() if sigmas.max() > 0 else torch.zeros_like(sigmas)
+        )
+        timesteps_for_sigma = (1 - scaled_sigmas) * (num_train_timesteps - 1)
+        indices = (
+            torch.round(timesteps_for_sigma).long().clamp(0, num_train_timesteps - 1)
+        )
+
+        alpha_prod_t = alphas_cumprod[indices].view(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        self.alpha_prod_t = alpha_prod_t
+        self.beta_prod_t = beta_prod_t
+        self.alpha_prod_t_sqrt = torch.sqrt(alpha_prod_t)
+        self.beta_prod_t_sqrt = torch.sqrt(beta_prod_t)
+
+        print(f"[Alpha/Beta] Final stats:")
+        print(
+            f"  alpha_prod_t: min={self.alpha_prod_t.min():.4f}, max={self.alpha_prod_t.max():.4f}, mean={self.alpha_prod_t.mean():.4f}"
+        )
+        print(
+            f"  beta_prod_t: min={self.beta_prod_t.min():.4f}, max={self.beta_prod_t.max():.4f}, mean={self.beta_prod_t.mean():.4f}"
+        )
+        print(
+            f"  alpha_sqrt: min={self.alpha_prod_t_sqrt.min():.4f}, max={self.alpha_prod_t_sqrt.max():.4f}, mean={self.alpha_prod_t_sqrt.mean():.4f}"
+        )
+        print(
+            f"  beta_sqrt: min={self.beta_prod_t_sqrt.min():.4f}, max={self.beta_prod_t_sqrt.max():.4f}, mean={self.beta_prod_t_sqrt.mean():.4f}\n"
+        )
+
+    def scheduler_step_batch(self, denoised_batch, scaled_noise):
+        """Compute scheduler step for the entire batch (LCM-like)"""
+        print("\n[Scheduler] Starting scheduler step (LCM-like):")
+        print(f"  Denoised shape: {denoised_batch.shape}")
+        print(f"  Scaled noise shape: {scaled_noise.shape}")
+        print(
+            f"  Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}"
+        )
+        print(
+            f"  Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}"
+        )
+
+        # Assuming self.alpha_prod_t and self.beta_prod_t are for the current timestep
+        alpha_t_sqrt = self.alpha_prod_t_sqrt[: denoised_batch.shape[0]]
+        beta_t_sqrt = self.beta_prod_t_sqrt[: denoised_batch.shape[0]]
+
+        # Calculate the predicted original sample (x_0) from the noisy sample and model output
+        # This is similar to the inverse of the forward diffusion step
+        pred_original_sample = (
+            denoised_batch - beta_t_sqrt.view(-1, 1, 1, 1) * scaled_noise
+        ) / (alpha_t_sqrt.view(-1, 1, 1, 1) + 1e-7)
+        print(
+            f"  Predicted original sample stats: min={pred_original_sample.min():.4f}, max={pred_original_sample.max():.4f}, mean={pred_original_sample.mean():.4f}"
+        )
+
+        # For the reverse step, we need the alpha and beta for the previous timestep
+        # Since we are operating on a batch of timesteps, we'll use the values corresponding to the first timestep in the batch for updating stock_noise
+        # In StreamDiffusion paper, the update happens based on the result of the last denoising step
+
+        # Get alpha and beta for the next step: (which would be the previous in the denoising sequence)
+        alpha_next_sqrt = torch.cat(
+            [
+                self.alpha_prod_t_sqrt[1 : denoised_batch.shape[0]],
+                torch.ones_like(self.alpha_prod_t_sqrt[0:1]),
+            ],
+            dim=0,
+        )
+
+        # Calculate the coefficient for the predicted original sample in the next step
+        sigma_next = (1 - alpha_next_sqrt.pow(2)) / (alpha_next_sqrt + 1e-7)
+        sigma_t = (1 - alpha_t_sqrt.pow(2)) / (alpha_t_sqrt + 1e-7)
+
+        # Guidance scale (set to 1 for unconditional prediction)
+        guidance_scale = 1
+
+        # Calculate delta_x (the change in the noisy sample)
+        delta_x = (sigma_next - sigma_t) * guidance_scale * pred_original_sample
+
+        print(
+            f"  Final delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}\n"
+        )
+
+        return delta_x
+
+    def sample(
+        self, model, noise, sigmas, extra_args=None, callback=None, disable=None
+    ):
+        """Sample with staggered batch denoising steps"""
+        extra_args = {} if extra_args is None else extra_args
+
+        print(f"\n[Sample] Starting sample step {self.step_counter}:")
+        print(f"  Input noise shape: {noise.shape}")
+        print(f"  Sigmas shape: {sigmas.shape}")
+        print(f"  Sigmas values: {sigmas[:self.num_timesteps]}")
+
+        # Set up cross attention optimization if prompt provided
+        if self.prompt:
+            model = self.setup_cross_attention(model, self.prompt)
+
+        # Get number of frames in batch and available sigmas
+        batch_size = noise.shape[0]
+
+        # Calculate total batch size based on frame_buffer_size and number of timesteps
+        total_batch_size = self.frame_buffer_size * self.num_timesteps
+        if batch_size != total_batch_size:
+            raise ValueError(
+                f"Batch size ({batch_size}) must match frame_buffer_size * num_timesteps ({total_batch_size})"
+            )
+
+        # Pre-compute alpha and beta terms if not already computed or if sigmas changed
+        if self.cached_sigmas is None or not torch.equal(self.cached_sigmas, sigmas):
+            self.compute_alpha_beta(sigmas)
+
+        # Update noise sequence like StreamDiffusion
+        if self.step_counter == 0:
+            # First step - use initial noise
+            self.stock_noise = self.init_noise[0].clone()
+        else:
+            # Shift noise sequence
+            self.stock_noise = torch.cat(
+                [self.init_noise[1:], self.init_noise[0:1]], dim=0
+            )[0]
+
+        print(
+            f"  Stock noise stats: min={self.stock_noise.min():.4f}, max={self.stock_noise.max():.4f}, mean={self.stock_noise.mean():.4f}"
+        )
+
+        # Scale noise for each frame based on its sigma - vectorized with caching
+        sigma_scaled_noise = self.get_expanded_noise(batch_size)
+        print(f"  Expanded noise shape: {sigma_scaled_noise.shape}")
+        print(
+            f"  Expanded noise stats: min={sigma_scaled_noise.min():.4f}, max={sigma_scaled_noise.max():.4f}, mean={sigma_scaled_noise.mean():.4f}"
+        )
+
+        x = noise + sigma_scaled_noise * sigmas[: self.num_timesteps].view(-1, 1, 1, 1)
+        print(
+            f"  Initial x stats: min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}"
+        )
+
+        # Initialize frame buffer if needed
+        if self.x_t_latent_buffer is None and self.num_timesteps > 1:
+            self.x_t_latent_buffer = x[0].clone()
+            print("  Initialized latent buffer with first frame")
+
+        # Use buffer for first frame to maintain temporal consistency
+        if self.num_timesteps > 1:
+            x = torch.cat([self.x_t_latent_buffer.unsqueeze(0), x[1:]], dim=0)
+            print(
+                f"  After buffer concat x stats: min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}"
+            )
+
+        # Run model on entire batch at once
+        with torch.no_grad():
+            # Process all frames in parallel
+            sigma_batch = sigmas[: self.num_timesteps]
+            # Process with CFG handling inside process_with_cfg
+            denoised_batch = self.process_with_cfg(model, x, sigma_batch, extra_args)
+            print(f"\n[Sample] After denoising:")
+            print(f"  Denoised shape: {denoised_batch.shape}")
+            print(
+                f"  Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}"
+            )
+
+            # Update buffer with intermediate results
+            if self.num_timesteps > 1:
+                # Store result from first frame as buffer for next iteration
+                self.x_t_latent_buffer = denoised_batch[0].clone()
+                print(
+                    f"  Updated buffer stats: min={self.x_t_latent_buffer.min():.4f}, max={self.x_t_latent_buffer.max():.4f}, mean={self.x_t_latent_buffer.mean():.4f}"
+                )
+                # Return result from last frame
+                x_0_pred_out = denoised_batch[-1].unsqueeze(0)
+            else:
+                x_0_pred_out = denoised_batch
+                self.x_t_latent_buffer = None
+
+            print(
+                f"  Output pred stats: min={x_0_pred_out.min():.4f}, max={x_0_pred_out.max():.4f}, mean={x_0_pred_out.mean():.4f}"
+            )
+
+            # Call callback if provided
+            if callback is not None:
+                callback(
+                    {
+                        "x": x_0_pred_out,
+                        "i": 0,
+                        "sigma": sigmas[0],
+                        "sigma_hat": sigmas[0],
+                        "denoised": x_0_pred_out,
+                    }
+                )
+
+            # Update stock noise using scheduler steps
+            if self.cfg_type == "self":
+                print("\n[Sample] Updating stock noise:")
+                # Scale noise properly
+                scaled_noise = (
+                    self.beta_prod_t_sqrt[: denoised_batch.shape[0]].view(-1, 1, 1, 1)
+                    * self.stock_noise
+                )
+                print(
+                    f"  Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}"
+                )
+
+                # Get delta_x using scheduler step
+                delta_x = self.scheduler_step_batch(denoised_batch, scaled_noise)
+                print(
+                    f"  Initial delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}"
+                )
+
+                # Prepare next step's alpha/beta terms
+                alpha_next = torch.cat(
+                    [
+                        self.alpha_prod_t_sqrt[1 : denoised_batch.shape[0]],
+                        torch.ones_like(self.alpha_prod_t_sqrt[0:1]),
+                    ],
+                    dim=0,
+                )
+                beta_next = torch.cat(
+                    [
+                        self.beta_prod_t_sqrt[1 : denoised_batch.shape[0]],
+                        torch.zeros_like(self.beta_prod_t_sqrt[0:1]),
+                    ],
+                    dim=0,
+                )
+
+                print(
+                    f"  Alpha next stats: min={alpha_next.min():.4f}, max={alpha_next.max():.4f}, mean={alpha_next.mean():.4f}"
+                )
+                print(
+                    f"  Beta next stats: min={beta_next.min():.4f}, max={beta_next.max():.4f}, mean={beta_next.mean():.4f}"
+                )
+
+                # Update noise prediction
+                delta_x = (alpha_next.view(-1, 1, 1, 1) * delta_x) / (
+                    beta_next.view(-1, 1, 1, 1) + 1e-7
+                )
+                print(
+                    f"  Scaled delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}"
+                )
+
+                # Use noise sequence for init_noise
+                init_noise = torch.cat(
+                    [self.init_noise[1:], self.init_noise[0:1]], dim=0
+                )
+                print(
+                    f"  Init noise stats: min={init_noise.min():.4f}, max={init_noise.max():.4f}, mean={init_noise.mean():.4f}"
+                )
+
+                # Update stock noise with proper scaling
+                self.stock_noise = init_noise[0] + delta_x[0]
+                print(
+                    f"  Final stock noise stats: min={self.stock_noise.min():.4f}, max={self.stock_noise.max():.4f}, mean={self.stock_noise.mean():.4f}\n"
+                )
+
+            # Update step counter
+            self.step_counter += 1
+
+        return x_0_pred_out
+
+    def buffer(self, latent, frame_buffer_size=1, num_timesteps=2, always_execute=True):
+        """Add new frame to buffer and return batch when ready"""
+        self.frame_buffer_size = frame_buffer_size
+        self.num_timesteps = num_timesteps
+
+        # Extract latent tensor from input and remove batch dimension if present
+        x = latent["samples"]
+        if x.dim() == 4:  # [B,C,H,W]
+            x = x.squeeze(0)  # Remove batch dimension -> [C,H,W]
+
+        # Initialize or resize frame buffer if needed
+        if (
+            self.frame_buffer is None
+            or self.frame_buffer.shape[0] != self.frame_buffer_size
+        ):
+            # Pre-allocate tensor buffer
+            self.frame_buffer = torch.zeros(
+                (self.frame_buffer_size,) + x.shape, dtype=x.dtype, device=x.device
+            )
+            if self.debug:
+                print(
+                    f"[StreamFrameBuffer] Pre-allocated tensor buffer for {self.frame_buffer_size} frames"
+                )
+
+        # Rotate frames and add new frame using tensor operations
+        if self.frame_buffer_size > 1:
+            self.frame_buffer = torch.roll(self.frame_buffer, -1, dims=0)
+        self.frame_buffer[-1] = x
+
+        # Create batch by repeating frames for each timestep
+        batch = self.frame_buffer.repeat_interleave(
+            self.num_timesteps, dim=0
+        )  # [buffer_size*num_timesteps,C,H,W]
+
+        # Return as latent dict
+        return {"samples": batch}
+
+    def update(
+        self,
+        latent,
+        num_timesteps,
+        frame_buffer_size=1,
+        cfg_type="self",
+        delta=1.0,
+        prompt="",
+        similarity_threshold=0.0,
+        max_skip_frames=10,
+        always_execute=True,
+    ):
+        """Create sampler with specified settings and process frame"""
+        self.frame_buffer_size = frame_buffer_size
+        self.cfg_type = cfg_type
+        self.delta = delta
+        self.prompt = prompt
+        self.num_timesteps = num_timesteps
+
+        # Update similarity filter settings
+        self.similarity_threshold = similarity_threshold
+        self.max_skip_frames = max_skip_frames
+        if self.similarity_filter is not None:
+            self.similarity_filter.set_threshold(similarity_threshold)
+            self.similarity_filter.set_max_skip_frame(max_skip_frames)
+
+        # Create sampler with custom sample function
+        sampler = comfy.samplers.KSAMPLER(self.sample)
+
+        # Process frame
+        buffered = self.buffer(latent)
+
+        # Get latent dimensions from input
+        _, _, self.latent_height, self.latent_width = latent["samples"].shape
+
+        # Add noise sequence initialization
+        if self.noise_sequence is None or self.noise_sequence.shape[0] != num_timesteps:
+            # Initialize with random noise like StreamDiffusion
+            self.init_noise = torch.randn(
+                (num_timesteps, 4, self.latent_height, self.latent_width),
+                device=comfy.model_management.get_torch_device(),
+                dtype=latent["samples"].dtype,
+            )
+
+            # Initialize stock noise as zeros like StreamDiffusion
+            self.stock_noise = torch.zeros_like(self.init_noise[0])
+
+            # Create noise sequence from init_noise
+            self.noise_sequence = self.init_noise.clone()
+
+            # No need for manual shifting as it's handled in sample() method
+
+        return (sampler, buffered)
+
+
+class StreamScheduler(ControlNodeBase):
+    """Implements StreamDiffusion's efficient timestep selection"""
+
+    RETURN_TYPES = ("SIGMAS",)
+    FUNCTION = "update"
+    CATEGORY = "real-time/sampling"
+
+    def __init__(self):
+        super().__init__()
+        self.cached_indices = None
+        self.last_t_index_list = None
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        inputs = super().INPUT_TYPES()
+        inputs["required"].update(
+            {
+                "model": ("MODEL",),
+                "t_index_list": (
+                    "STRING",
+                    {
+                        "default": "32,45",
+                        "tooltip": "Comma-separated list of timesteps to actually use for denoising. For LCM: use '32,45' for img2img or '0,16,32,45' for txt2img. For SDXL Turbo: use '45,49' for 2 steps or '49' for 1 step (Turbo works best with steps near the end of denoising)",
+                    },
+                ),
+                "num_inference_steps": (
+                    "INT",
+                    {
+                        "default": 50,
+                        "min": 1,
+                        "max": 1000,
+                        "step": 1,
+                        "tooltip": "Total number of timesteps in schedule. StreamDiffusion uses 50 by default. Only timesteps specified in t_index_list are actually used.",
+                    },
+                ),
+            }
+        )
+        return inputs
+
+    def parse_timesteps(self, t_index_list):
+        """Parse and sort timesteps, with caching"""
+        if t_index_list == self.last_t_index_list and self.cached_indices is not None:
+            return self.cached_indices
+
+        try:
+            # Parse and sort in reverse order once
+            indices = sorted(
+                [int(t.strip()) for t in t_index_list.split(",")], reverse=True
+            )
+            self.last_t_index_list = t_index_list
+            self.cached_indices = indices
+            return indices
+        except ValueError as e:
+            print(f"Error parsing timesteps: {e}. Using default [32,45]")
+            self.last_t_index_list = "32,45"
+            self.cached_indices = [45, 32]  # Already sorted in reverse
+            return self.cached_indices
+
+    def update(
+        self, model, t_index_list="32,45", num_inference_steps=50, always_execute=True
+    ):
+        # Get model's sampling parameters
+        model_sampling = model.get_model_object("model_sampling")
+
+        # Get parsed and sorted timesteps
+        t_indices = self.parse_timesteps(t_index_list)
+
+        # Create full schedule using normal scheduler
+        full_sigmas = comfy.samplers.normal_scheduler(
+            model_sampling, num_inference_steps
+        )
+
+        # Select only the sigmas at our desired indices
+        selected_sigmas = []
+        for t in t_indices:
+            if t < 0 or t >= num_inference_steps:
+                print(
+                    f"Warning: timestep {t} out of range [0,{num_inference_steps}), skipping"
+                )
+                continue
+            selected_sigmas.append(float(full_sigmas[t]))
+
+        # Add final sigma
+        selected_sigmas.append(0.0)
+
+        # Create tensor directly on target device
+        device = comfy.model_management.get_torch_device()
+        selected_sigmas = torch.tensor(
+            selected_sigmas, dtype=torch.float32, device=device
+        )
+        return (selected_sigmas,)
+
+
+class StreamFrameBuffer:
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "latent": ("LATENT",),
+                "buffer_size": ("INT", {
+                    "default": 4,
+                    "min": 1,
+                    "max": 10,
+                    "step": 1,
+                    "tooltip": "Number of frames to buffer before starting batch processing. Should match number of denoising steps."
+                }),
+            },
+        }
+    
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "update"
+    CATEGORY = "StreamPack/sampling"
+    DESCRIPTION = "Accumulates frames to enable staggered batch denoising like StreamDiffusion. Use in conjunction with StreamBatchSampler"
+    
+    
+    def __init__(self):
+        self.frame_buffer = None  # Tensor of shape [buffer_size, C, H, W]
+        self.buffer_size = None
+        self.buffer_pos = 0  # Current position
+        self.is_initialized = False  # Track buffer initialization
+        self.is_txt2img_mode = False
+    
+    def update(self, latent, buffer_size=4):
+        """Add new frame to buffer and return batch when ready"""
+        self.buffer_size = buffer_size
+        
+        # Extract latent tensor from input and remove batch dimension if present
+        x = latent["samples"]
+        
+        # Check if this is a txt2img (zeros tensor) or img2img mode
+        # In ComfyUI, EmptyLatentImage returns a zeros tensor with shape [batch_size, 4, h//8, w//8]
+        # We consider it txt2img mode if the tensor contains all zeros
+        is_txt2img_mode = torch.sum(torch.abs(x)) < 1e-6
+        self.is_txt2img_mode = is_txt2img_mode
+        
+        # If it's a batch with size 1, remove the batch dimension for our buffer
+        if x.dim() == 4 and x.shape[0] == 1:  # [1,C,H,W]
+            x = x.squeeze(0)  # Remove batch dimension -> [C,H,W]
+        
+        # Initialize or resize frame_buffer as a tensor
+        if not self.is_initialized or self.frame_buffer.shape[0] != self.buffer_size or \
+           self.frame_buffer.shape[1:] != x.shape:
+            # Pre-allocate buffer with correct shape
+            self.frame_buffer = torch.zeros(
+                (self.buffer_size, *x.shape),
+                device=x.device,
+                dtype=x.dtype
+            )
+            if self.is_txt2img_mode or not self.is_initialized:
+                # Optimization: Use broadcasting to fill buffer with copies
+                self.frame_buffer[:] = x.unsqueeze(0)  # Broadcast x to [buffer_size, C, H, W]
+
+            self.is_initialized = True
+            self.buffer_pos = 0
+        else:
+            # Add new frame to buffer using ring buffer logic
+            self.frame_buffer[self.buffer_pos] = x  # In-place update
+            self.buffer_pos = (self.buffer_pos + 1) % self.buffer_size  # Circular increment
+        
+        # Optimization: frame_buffer is already a tensor batch, no need to stack
+        batch = self.frame_buffer
+
+        
+        # Return as latent dict with preserved dimensions
+        result = {"samples": batch}
+        
+        # Preserve height and width if present in input
+        if "height" in latent:
+            result["height"] = latent["height"]
+        if "width" in latent:
+            result["width"] = latent["width"]
+            
+        return (result,)

src/batched_denoising/control_base.py

@@ -0,0 +1,92 @@
+from abc import ABC, abstractmethod
+import time
+
+class Pattern(ABC):
+    """Base class for all patterns (movement, timing, etc)"""
+    @abstractmethod
+    def calculate(self, phase: float, min_val: float, max_val: float) -> float:
+        """Calculate pattern value at given phase (0-1)"""
+        pass
+
+    @abstractmethod
+    def get_name(self) -> str:
+        """Pattern name for UI"""
+        pass
+
+class StateManager:
+    """Handles persistent state between executions"""
+    def __init__(self):
+            self._states = {}
+    
+    def get_state(self, node_id: str, default=None):
+        return self._states.get(node_id, default)
+    
+    def set_state(self, node_id: str, state):
+        self._states[node_id] = state
+        
+    def clear_state(self, node_id: str) -> bool:
+        """Clear state for a specific node. Returns True if state was found and cleared."""
+        if node_id in self._states:
+            del self._states[node_id]
+            return True
+        return False
+        
+    def clear_all_states(self):
+        """Clear all states"""
+        self._states.clear()
+
+    def get_all_node_ids(self):
+        """Get list of all node IDs with state"""
+        return list(self._states.keys())
+        
+    def cleanup_orphaned_states(self, active_node_ids: set[str]):
+        """Remove states for nodes that no longer exist"""
+        orphaned = [nid for nid in self._states.keys() if nid not in active_node_ids]
+        for nid in orphaned:
+            del self._states[nid]
+        return len(orphaned)
+
+class ControlNodeBase(ABC):
+    """Base class for all control nodes"""
+    
+    # Shared state manager
+    state_manager = StateManager()
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        """Base input types common to all control nodes"""
+        return {
+            "required": {
+                "always_execute": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "When enabled, the node updates every execution"
+                }),
+            },
+        }
+    
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        if kwargs.get('always_execute', True):
+            return float(time.time())
+        return False
+    
+    def __init__(self):
+        self.node_id = str(id(self))
+        
+    def get_state(self, default=None):
+        return self.state_manager.get_state(self.node_id, default)
+        
+    def set_state(self, state):
+        self.state_manager.set_state(self.node_id, state)
+
+    def cleanup(self):
+        """Clean up node state when node is destroyed"""
+        self.state_manager.clear_state(self.node_id)
+
+    def __del__(self):
+        self.cleanup()
+
+    @abstractmethod
+    def update(self, *args, **kwargs):
+        """Main update method to be implemented by control nodes"""
+        pass

src/feature_bank/__init__.py

@@ -0,0 +1 @@
+# Dummy file ensuring this package will be recognized.

src/feature_bank/feature_bank.py

@@ -0,0 +1,536 @@
+"""Source code for the Feature Bank of the StreamV2V paper.
+
+References:
+    StreamV2V: A Deep Generative Model for High-fidelity View Synthesis - Liang et al.,
+    2025, URL: https://arxiv.org/abs/2405.15757
+"""
+
+from typing import Optional, Deque, List, Tuple, Dict, Any
+from collections import deque
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from diffusers.utils import USE_PEFT_BACKEND
+
+
+def get_nn_feats(
+    x: torch.Tensor, y: torch.Tensor, threshold: float = 0.9
+) -> torch.Tensor:
+    """
+    Selects features from x or nearest neighbors from y based on cosine similarity.
+
+    Args:
+        x (torch.Tensor): Input features tensor.
+        y (torch.Tensor): Neighbor features tensor.
+         threshold (float): Cosine similarity threshold for neighbor selection.
+
+    Returns:
+        torch.Tensor: Selected features tensor.
+    """
+    if isinstance(x, deque):
+        x = torch.cat(list(x), dim=1)
+    if isinstance(y, deque):
+        y = torch.cat(list(y), dim=1)
+
+    x_norm = F.normalize(x, p=2, dim=-1)
+    y_norm = F.normalize(y, p=2, dim=-1)
+
+    cosine_similarity = torch.matmul(x_norm, y_norm.transpose(1, 2))
+    max_cosine_values, nearest_neighbors_indices = torch.max(cosine_similarity, dim=-1)
+    mask = max_cosine_values < threshold
+
+    indices_expanded = nearest_neighbors_indices.unsqueeze(-1).expand(
+        -1, -1, x_norm.size(-1)
+    )
+    nearest_neighbor_tensor = torch.gather(y, 1, indices_expanded)
+    selected_tensor = torch.where(mask.unsqueeze(-1), x, nearest_neighbor_tensor)
+
+    return selected_tensor
+
+
+class FeatureBank:
+    """
+    Feature bank to cache and retrieve attention keys, values, and outputs.
+
+    Args:
+        max_frames (int): Maximum number of frames to store in the cache.
+    """
+
+    def __init__(self, max_frames: int = 4):
+        self.cached_key: Deque[torch.FloatTensor] = deque(maxlen=max_frames)
+        self.cached_value: Deque[torch.FloatTensor] = deque(maxlen=max_frames)
+        self.cached_output: Deque[torch.FloatTensor] = deque(maxlen=max_frames)
+
+    def add(
+        self,
+        key: torch.FloatTensor,
+        value: torch.FloatTensor,
+        output: torch.FloatTensor,
+    ):
+        """Add key, value, and output tensors to the feature bank."""
+        self.cached_key.append(key.clone())
+        self.cached_value.append(value.clone())
+        self.cached_output.append(output.clone())
+
+    def get(
+        self,
+    ) -> Tuple[
+        List[torch.FloatTensor], List[torch.FloatTensor], List[torch.FloatTensor]
+    ]:
+        """Retrieve cached keys, values, and outputs as lists."""
+        return list(self.cached_key), list(self.cached_value), list(self.cached_output)
+
+    def clear(self):
+        """Clear all cached features."""
+        self.cached_key.clear()
+        self.cached_value.clear()
+        self.cached_output.clear()
+
+
+class FeatureBankAttentionProcessor:
+    """
+    Attention processor with feature injection.
+    """
+
+    DEBUG_PRINTS = (
+        False  # Class-level debug flag, turn on/off all debug prints from this class
+    )
+
+    def __init__(
+        self,
+        use_feature_injection=True,
+        feature_injection_strength=0.8,
+        feature_similarity_threshold=0.98,
+        interval=4,
+        max_frames=4,
+    ):
+        self.use_feature_injection = use_feature_injection
+        self.fi_strength = feature_injection_strength
+        self.threshold = feature_similarity_threshold
+        self.frame_id = 0
+        self.interval = interval
+        self.feature_bank = FeatureBank(max_frames=max_frames)
+        self.key_projection: Optional[nn.Linear] = (
+            None  # Linear layer for key projection
+        )
+        self.value_projection: Optional[nn.Linear] = (
+            None  # Linear layer for value projection
+        )
+        self.output_projection: Optional[nn.Linear] = (
+            None  # Linear layer for output projection
+        )
+
+    def _debug_print(self, *args, **kwargs):
+        """Conditional debug print based on class-level DEBUG_PRINTS flag."""
+        if self.DEBUG_PRINTS:
+            print("[FeatureBankAttentionProcessor DEBUG]", *args, **kwargs)
+
+    def _project_tensor(
+        self,
+        tensor: torch.Tensor,
+        projection_layer: Optional[nn.Linear],
+        expected_dim: int,
+        layer_type: str,
+    ) -> torch.Tensor:
+        """
+        Handles tensor projection if dimension mismatch is detected.
+        Creates projection layer if needed and applies it.
+        """
+        if (
+            projection_layer is None
+            or projection_layer.in_features != tensor.shape[-1]
+            or projection_layer.out_features != expected_dim
+        ):
+            self._debug_print(
+                f"\n--- DEBUG: {layer_type} dimension mismatch! Projecting cached {layer_type.lower()}s. ---"
+            )
+            self._debug_print(
+                f"Original {layer_type.lower()} dimension: {tensor.shape[-1]}, Expected: {expected_dim}"
+            )
+            new_projection_layer = (
+                nn.Linear(tensor.shape[-1], expected_dim)
+                .to(tensor.device)
+                .to(tensor.dtype)
+            )
+            self._debug_print(
+                f"Creating new {layer_type.lower()} projection layer: In Features={tensor.shape[-1]}, Out Features={expected_dim}"
+            )
+            if layer_type == "Key":
+                self.key_projection = new_projection_layer
+            elif layer_type == "Value":
+                self.value_projection = new_projection_layer
+            elif layer_type == "Output":
+                self.output_projection = new_projection_layer
+            projection_layer = (
+                new_projection_layer  # Use the newly created or updated layer
+            )
+
+        projected_tensor = projection_layer(tensor)
+        self._debug_print(
+            f"Shape of cached {layer_type.lower()}s after projection: {projected_tensor.shape}"
+        )
+        return projected_tensor
+
+    def __call__(
+        self,
+        attn: nn.Module,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        """
+        Forward pass of feature bank attention with feature injection.
+        """
+        residual = hidden_states
+        args = () if USE_PEFT_BACKEND else (scale,)
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, key_tokens, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+
+        query = attn.to_q(hidden_states)
+
+        # Self-attention detection (heuristic based on input feature dimensions)
+        is_self_attn = attn.to_k.in_features == attn.to_q.in_features
+        current_encoder_hidden_states = (
+            hidden_states if is_self_attn else encoder_hidden_states
+        )
+
+        self._debug_print("\n--- Attention Module Info ---")
+        self._debug_print("Module type:", type(attn))
+        self._debug_print("Module:", attn)
+        self._debug_print("Is self-attention (heuristic):", is_self_attn)
+        self._debug_print("Shape of hidden_states:", hidden_states.shape)
+        self._debug_print(
+            "Shape of current_encoder_hidden_states before attn.to_k:",
+            (
+                current_encoder_hidden_states.shape
+                if current_encoder_hidden_states is not None
+                else "None"
+            ),
+        )
+        self._debug_print("Shape of attn.to_k.weight:", attn.to_k.weight.shape)
+
+        key = attn.to_k(current_encoder_hidden_states)
+        value = attn.to_v(current_encoder_hidden_states)
+
+        cached_key_tensor = key.clone()
+        cached_value_tensor = value.clone()
+
+        num_heads = attn.heads  # Assuming attn.heads exists in CrossAttention
+        head_dim = query.shape[-1] // num_heads
+        query = (
+            query.view(batch_size, -1, num_heads, head_dim).transpose(1, 2).contiguous()
+        )  # [batch_size, num_heads, seq_len_q, head_dim]
+
+        if is_self_attn:
+            (
+                cached_keys_list,
+                cached_values_list,
+                cached_outputs,
+            ) = self.feature_bank.get()
+
+            self._debug_print(
+                "\n--- Feature Bank Get Output ---"
+            )  # Separator for clarity
+            self._debug_print(
+                f"Type of cached_keys_list from feature_bank.get(): {type(cached_keys_list)}"
+            )
+            self._debug_print(
+                f"Content of cached_keys_list from feature_bank.get(): {cached_keys_list}"
+            )
+
+            if (
+                cached_keys_list
+                and cached_keys_list[0] is not None
+                and torch.is_tensor(cached_keys_list[0])
+            ):
+                cached_key_tensor_from_list = cached_keys_list[0]
+                cached_key_tensor_from_list = self._project_tensor(
+                    cached_key_tensor_from_list,
+                    self.key_projection,
+                    key.shape[-1],
+                    "Key",
+                )
+
+                self._debug_print("\n--- DEBUG: Before torch.cat (Keys) ---")
+                self._debug_print(
+                    "Module:", attn
+                )  # Print the module itself for identification
+                self._debug_print("Type of key:", type(key))
+                self._debug_print("Shape of key:", key.shape)
+                self._debug_print(
+                    "Type of cached_key_tensor_from_list:",
+                    type(cached_key_tensor_from_list),
+                )
+                self._debug_print(
+                    "Shape of cached_key_tensor_from_list:",
+                    cached_key_tensor_from_list.shape,
+                )
+
+                key = torch.cat([key, cached_key_tensor_from_list], dim=1)
+
+            if (
+                cached_values_list
+                and cached_values_list[0] is not None
+                and torch.is_tensor(cached_values_list[0])
+            ):
+                cached_value_tensor_from_list = cached_values_list[0]
+                cached_value_tensor_from_list = self._project_tensor(
+                    cached_value_tensor_from_list,
+                    self.value_projection,
+                    value.shape[-1],
+                    "Value",
+                )
+
+                self._debug_print("\n--- DEBUG: Before torch.cat (Values) ---")
+                self._debug_print(
+                    "Module:", attn
+                )  # Print the module itself for identification
+                self._debug_print("Type of value:", type(value))
+                self._debug_print("Shape of value:", value.shape)
+                self._debug_print(
+                    "Type of cached_value_tensor_from_list:",
+                    type(cached_value_tensor_from_list),
+                )
+                self._debug_print(
+                    "Shape of cached_value_tensor_from_list:",
+                    cached_value_tensor_from_list.shape,
+                )
+
+                value = torch.cat([value, cached_value_tensor_from_list], dim=1)
+
+            key = (
+                key.view(-1, key.shape[1], num_heads, head_dim)
+                .transpose(1, 2)
+                .contiguous()
+            )  # [batch_size, num_heads, seq_len_kv, head_dim] (seq_len_kv == key_tokens)
+            value = (
+                value.view(-1, value.shape[1], num_heads, head_dim)
+                .transpose(1, 2)
+                .contiguous()
+            )  # [batch_size, num_heads, seq_len_kv, head_dim] (seq_len_kv == key_tokens)
+        else:
+            key = (
+                key.view(batch_size, -1, num_heads, head_dim)
+                .transpose(1, 2)
+                .contiguous()
+            )  # [batch_size, num_heads, seq_len_kv, head_dim] (seq_len_kv == key_tokens)
+            value = (
+                value.view(batch_size, -1, num_heads, head_dim)
+                .transpose(1, 2)
+                .contiguous()
+            )  # [batch_size, num_heads, seq_len_kv, head_dim] (seq_len_kv == key_tokens)
+
+        # --- Manual attention implementation ---
+        scale_factor = 1.0 / query.shape[-1] ** 0.5
+        query = query * scale_factor
+        attn_weights = torch.bmm(
+            query.view(batch_size * num_heads, -1, head_dim),
+            key.view(batch_size * num_heads, -1, head_dim).transpose(-2, -1),
+        )  # [batch_size * num_heads, seq_len_q, seq_len_kv]
+        attn_weights = torch.softmax(attn_weights, dim=-1)
+        attn_output = torch.bmm(
+            attn_weights, value.view(batch_size * num_heads, -1, head_dim)
+        )  # [batch_size * num_heads, seq_len_q, head_dim]
+        hidden_states = (
+            attn_output.view(batch_size, num_heads, -1, head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )  # [batch_size, seq_len_q, num_heads * head_dim]
+        # --- End of manual attention implementation ---
+
+        self._debug_print(f"Type of attn before to_out: {type(attn)}")  # Debug print
+
+        hidden_states = hidden_states.to(
+            query.dtype
+        )  # Keep this line to ensure correct dtype
+        hidden_states = hidden_states.reshape(
+            batch_size, -1, num_heads * head_dim
+        ).contiguous()  # Modified reshape - removed transpose after reshape
+
+        hidden_states = attn.to_out[0](hidden_states, *args)
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        rescale_output_factor = 1.0
+        cached_output_tensor = hidden_states.clone()
+
+        if self.use_feature_injection:
+            _, _, cached_outputs = self.feature_bank.get()
+            self._debug_print(
+                f"Type of cached_outputs: {type(cached_outputs)}"
+            )  # Debug print
+
+            if (
+                cached_outputs
+                and cached_outputs[0] is not None
+                and cached_outputs[0].numel() > 0
+            ):  # Corrected condition
+                cached_output_tensor_from_list = cached_outputs[
+                    0
+                ]  # Get the cached output tensor
+                cached_output_tensor_from_list = self._project_tensor(
+                    cached_output_tensor_from_list,
+                    self.output_projection,
+                    hidden_states.shape[-1],
+                    "Output",
+                )
+
+                self._debug_print("\n--- DEBUG: Before get_nn_feats ---")
+                self._debug_print("Shape of hidden_states:", hidden_states.shape)
+                self._debug_print(
+                    "Shape of cached_output_tensor_from_list:",
+                    cached_output_tensor_from_list.shape,
+                )
+
+                nn_hidden_states = get_nn_feats(
+                    hidden_states,
+                    cached_output_tensor_from_list,
+                    threshold=self.threshold,
+                )  # Pass PROJECTED tensor
+                hidden_states = (
+                    hidden_states * (1 - self.fi_strength)
+                    + self.fi_strength * nn_hidden_states
+                )
+
+        hidden_states = hidden_states / rescale_output_factor
+
+        if self.frame_id % self.interval == 0 and is_self_attn:
+            self.feature_bank.add(
+                cached_key_tensor, cached_value_tensor, cached_output_tensor
+            )
+
+        self.frame_id += 1
+        return hidden_states
+
+
+class FeatureBankAttentionProcessorNode:
+    """
+    ComfyUI Node for Feature Bank Attention Processor.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.cross_attention_hook = None
+        self.feature_bank_processor = FeatureBankAttentionProcessor()
+
+    @classmethod
+    def INPUT_TYPES(s) -> Dict[str, Dict[str, Tuple[str, Dict[str, Any]]]]:
+        """Input types for FeatureBankAttentionProcessorNode node."""
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "use_feature_injection": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),  # Configurable FI
+                "feature_injection_strength": (
+                    "FLOAT",
+                    {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01},
+                ),  # Configurable FI strength
+                "feature_similarity_threshold": (
+                    "FLOAT",
+                    {"default": 0.98, "min": 0.0, "max": 1.0, "step": 0.01},
+                ),  # Configurable threshold
+                "feature_cache_interval": (
+                    "INT",
+                    {"default": 4, "min": 1, "max": 64},
+                ),  # Configurable interval
+                "feature_bank_max_frames": (
+                    "INT",
+                    {"default": 4, "min": 1, "max": 16},
+                ),  # Configurable max frames in bank
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "update_node"
+    CATEGORY = "feature-bank"
+
+    def update_node(
+        self,
+        model: Any,
+        use_feature_injection: bool,
+        feature_injection_strength: float,
+        feature_similarity_threshold: float,
+        feature_cache_interval: int,
+        feature_bank_max_frames: int,
+    ) -> Tuple[Any]:
+        """
+        Update method for FeatureBankAttentionProcessorNode. Sets up feature injection and attention hooks.
+        """
+        print(f"[FeatureBankAttentionProcessorNode] Initializing ")
+
+        # Update FeatureBankProcessor configuration from node inputs on each update call
+        self.feature_bank_processor.use_feature_injection = use_feature_injection
+        self.feature_bank_processor.fi_strength = feature_injection_strength
+        self.feature_bank_processor.threshold = feature_similarity_threshold
+        self.feature_bank_processor.interval = feature_cache_interval
+        self.feature_bank_processor.feature_bank = FeatureBank(
+            max_frames=feature_bank_max_frames
+        )  # Re-instantiate feature bank with new max_frames if it changes
+
+        def cross_attention_forward(
+            module: nn.Module,
+            x: torch.Tensor,
+            context: Optional[torch.Tensor] = None,
+            mask: Optional[torch.Tensor] = None,
+            value: Optional[torch.Tensor] = None,
+            **kwargs,
+        ) -> (
+            torch.Tensor
+        ):  # Added **kwargs to accept potential extra args from original forward
+            """Optimized cross attention using FeatureBankAttentionProcessor."""
+            context = x if context is None else context
+            encoder_hidden_states = context  # Assuming 'context' in cross_attention_forward corresponds to encoder_hidden_states
+
+            return self.feature_bank_processor(
+                attn=module,  # Pass the current module as 'attn'
+                hidden_states=x,
+                encoder_hidden_states=encoder_hidden_states,
+                scale=1.0,  # Or pass scale if it's relevant, currently hardcoded to 1.0 in processor
+            )
+
+        def hook_cross_attention(
+            module: nn.Module, input: Tuple[torch.Tensor], output: torch.Tensor
+        ) -> torch.Tensor:
+            """Forward hook to replace cross-attention forward method."""
+            if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"):
+                if not hasattr(module, "_original_forward"):
+                    module._original_forward = module.forward
+                module.forward = lambda *args, **kwargs: cross_attention_forward(
+                    module, *args, **kwargs
+                )
+            return output
+
+        # Remove old hooks if they exist
+        if self.cross_attention_hook is not None:
+            self.cross_attention_hook.remove()
+
+        # Clone model and apply hooks
+        m = model.clone()
+
+        def register_hooks(module: nn.Module):
+            """Recursively register forward hooks for cross-attention modules."""
+            if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"):
+                self.cross_attention_hook = module.register_forward_hook(
+                    hook_cross_attention
+                )
+
+        m.model.apply(register_hooks)
+        return (m,)

src/media_pipe/facemesh.py

@@ -0,0 +1,608 @@
+import numpy as np
+import cv2
+import mediapipe as mp
+import time
+import torch
+
+class FaceMeshUtils:
+    """
+    Utility class for shared face mesh functionality.
+    Contains drawing methods and specifications used by multiple nodes.
+    """
+    def __init__(self):
+        # Initialize MediaPipe drawing utilities
+        self.mp_drawing = mp.solutions.drawing_utils
+        self.mp_drawing_styles = mp.solutions.drawing_styles
+        self.mp_face_mesh = mp.solutions.face_mesh
+        
+        # Define drawing specifications for different facial features
+        self.f_thick = 2
+        self.f_rad = 1
+        self.right_iris_draw = self.mp_drawing.DrawingSpec(color=(10, 200, 250), thickness=self.f_thick, circle_radius=self.f_rad)
+        self.right_eye_draw = self.mp_drawing.DrawingSpec(color=(10, 200, 180), thickness=self.f_thick, circle_radius=self.f_rad)
+        self.right_eyebrow_draw = self.mp_drawing.DrawingSpec(color=(10, 220, 180), thickness=self.f_thick, circle_radius=self.f_rad)
+        self.left_iris_draw = self.mp_drawing.DrawingSpec(color=(250, 200, 10), thickness=self.f_thick, circle_radius=self.f_rad)
+        self.left_eye_draw = self.mp_drawing.DrawingSpec(color=(180, 200, 10), thickness=self.f_thick, circle_radius=self.f_rad)
+        self.left_eyebrow_draw = self.mp_drawing.DrawingSpec(color=(180, 220, 10), thickness=self.f_thick, circle_radius=self.f_rad)
+        self.mouth_draw = self.mp_drawing.DrawingSpec(color=(10, 180, 10), thickness=self.f_thick, circle_radius=self.f_rad)
+        self.head_draw = self.mp_drawing.DrawingSpec(color=(10, 200, 10), thickness=self.f_thick, circle_radius=self.f_rad)
+        
+        # Create face connection specifications
+        self.face_connection_spec = {}
+        for edge in self.mp_face_mesh.FACEMESH_FACE_OVAL:
+            self.face_connection_spec[edge] = self.head_draw
+        for edge in self.mp_face_mesh.FACEMESH_LEFT_EYE:
+            self.face_connection_spec[edge] = self.left_eye_draw
+        for edge in self.mp_face_mesh.FACEMESH_LEFT_EYEBROW:
+            self.face_connection_spec[edge] = self.left_eyebrow_draw
+        for edge in self.mp_face_mesh.FACEMESH_RIGHT_EYE:
+            self.face_connection_spec[edge] = self.right_eye_draw
+        for edge in self.mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
+            self.face_connection_spec[edge] = self.right_eyebrow_draw
+        for edge in self.mp_face_mesh.FACEMESH_LIPS:
+            self.face_connection_spec[edge] = self.mouth_draw
+        
+        # Iris landmark specification
+        self.iris_landmark_spec = {468: self.right_iris_draw, 473: self.left_iris_draw}
+    
+    def reverse_channels(self, image):
+        """Given a numpy array in RGB form, convert to BGR. Will also convert from BGR to RGB."""
+        # im[:,:,::-1] is a neat hack to convert BGR to RGB by reversing the indexing order.
+        return image[:, :, ::-1]
+    
+    def draw_pupils(self, image, landmark_list, drawing_spec, halfwidth=2):
+        """Draw pupils with custom function for better control"""
+        if len(image.shape) != 3:
+            raise ValueError("Input image must be H,W,C.")
+        image_rows, image_cols, image_channels = image.shape
+        if image_channels != 3:  # BGR channels
+            raise ValueError('Input image must contain three channel bgr data.')
+            
+        # Ensure image is contiguous in memory for OpenCV compatibility
+        image = np.ascontiguousarray(image, dtype=np.uint8)
+            
+        for idx, landmark in enumerate(landmark_list.landmark):
+            if ((landmark.HasField('visibility') and landmark.visibility < 0.9) or
+                (landmark.HasField('presence') and landmark.presence < 0.5)):
+                continue
+            if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:
+                continue
+            image_x = int(image_cols*landmark.x)
+            image_y = int(image_rows*landmark.y)
+            draw_color = None
+            if isinstance(drawing_spec, dict):
+                if drawing_spec.get(idx) is None:
+                    continue
+                else:
+                    draw_color = drawing_spec[idx].color
+            elif isinstance(drawing_spec, self.mp_drawing.DrawingSpec):
+                draw_color = drawing_spec.color
+            image[image_y-halfwidth:image_y+halfwidth, image_x-halfwidth:image_x+halfwidth, :] = draw_color
+    
+    def draw_facemesh(self, img_rgb, results, thickness=1):
+        """Draw face mesh visualization using the ControlNet Aux methodology on a black background"""
+        # Create an empty image for drawing
+        empty = np.zeros_like(img_rgb)
+        
+        # Ensure the image is contiguous in memory and has the correct data type for OpenCV
+        empty = np.ascontiguousarray(empty, dtype=np.uint8)
+        
+        # Update thickness based on input parameter
+        for spec in self.face_connection_spec.values():
+            spec.thickness = thickness
+        
+        # Draw detected faces
+        if results and results.multi_face_landmarks:
+            for face_landmarks in results.multi_face_landmarks:
+                self.mp_drawing.draw_landmarks(
+                    empty,
+                    face_landmarks,
+                    connections=self.face_connection_spec.keys(),
+                    landmark_drawing_spec=None,
+                    connection_drawing_spec=self.face_connection_spec
+                )
+                self.draw_pupils(empty, face_landmarks, self.iris_landmark_spec, 2)
+        
+        # Flip BGR back to RGB
+        empty = self.reverse_channels(empty).copy()
+        
+        return empty
+        
+    def draw_facemesh_on_image(self, img_rgb, results, thickness=1):
+        """Draw face mesh visualization on the provided image"""
+        # Make a copy of the input image to avoid modifying the original
+        img_copy = img_rgb.copy()
+        
+        # Convert RGB to BGR for OpenCV and MediaPipe drawing functions
+        img_copy = self.reverse_channels(img_copy)
+        
+        # Ensure the image is contiguous in memory and has the correct data type for OpenCV
+        img_copy = np.ascontiguousarray(img_copy, dtype=np.uint8)
+        
+        # Update thickness based on input parameter
+        for spec in self.face_connection_spec.values():
+            spec.thickness = thickness
+        
+        # Draw detected faces
+        if results and results.multi_face_landmarks:
+            for face_landmarks in results.multi_face_landmarks:
+                self.mp_drawing.draw_landmarks(
+                    img_copy,
+                    face_landmarks,
+                    connections=self.face_connection_spec.keys(),
+                    landmark_drawing_spec=None,
+                    connection_drawing_spec=self.face_connection_spec
+                )
+                self.draw_pupils(img_copy, face_landmarks, self.iris_landmark_spec, 2)
+        
+        # Convert back from BGR to RGB
+        img_copy = self.reverse_channels(img_copy).copy()
+        
+        return img_copy
+
+# Create a global instance of the utils class
+face_mesh_utils = FaceMeshUtils()
+
+class FaceMeshDrawNode:
+    """
+    A node that draws MediaPipe face mesh on images.
+    Takes face mesh data from FaceMeshDetectNode and visualizes it.
+    Optimized for visualization only, requires an input image.
+    Uses the same drawing methodology as ControlNet Aux for optimal performance.
+    """
+    def __init__(self):
+        # Use the shared utils instance
+        self.utils = face_mesh_utils
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "facemesh_data": ("FACEMESH_DATA",),
+                "draw_contours": ("BOOLEAN", {
+                    "default": True,
+                    "description": "Draw contour lines connecting mesh points"
+                }),
+                "mesh_thickness": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.1,
+                    "max": 5.0,
+                    "step": 0.1,
+                    "display": "slider",
+                    "description": "Thickness of mesh lines"
+                }),
+                "mesh_color": (["green", "red", "blue", "yellow", "white"], {
+                    "default": "green",
+                    "description": "Color of the face mesh lines"
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "draw_facemesh"
+    CATEGORY = "ComfyUI_FaceMesh"
+    
+    def draw_facemesh(self, image, facemesh_data, draw_contours, mesh_thickness, mesh_color):
+        """Draw face mesh visualization on the input image"""
+        # Get batch size 
+        batch_size = image.shape[0]
+        output_images = []
+        
+        # Process each image in the batch
+        for b in range(batch_size):
+            # Get current image
+            img = image[b].numpy()
+            
+            # Get corresponding facemesh data
+            current_facemesh_data = facemesh_data[b]
+            
+            # If no face detected, add original image unchanged
+            if not current_facemesh_data["face_detected"]:
+                output_images.append(img)
+                continue
+            
+            # Convert from float [0,1] to uint8 [0,255] for drawing
+            img_rgb = (img * 255).astype(np.uint8)
+            
+            # Ensure the image is contiguous in memory for OpenCV compatibility
+            img_rgb = np.ascontiguousarray(img_rgb, dtype=np.uint8)
+            
+            # Use original results for drawing
+            results = current_facemesh_data["original_results"]
+            
+            # Draw the face mesh on the input image using the shared utility
+            output_img_rgb = self.utils.draw_facemesh_on_image(img_rgb, results, thickness=int(mesh_thickness))
+            
+            # Convert back to float [0,1] for ComfyUI
+            output_img = output_img_rgb.astype(np.float32) / 255.0
+            
+            # Add to output list
+            output_images.append(output_img)
+        
+        # Convert output images list to tensor with batch dimension
+        output_batch = torch.from_numpy(np.stack(output_images, axis=0))
+        
+        return (output_batch,)
+
+class FaceMeshMaskNode:
+    """
+    Creates binary masks for different regions of the face detected by FaceMeshDetectNode.
+    """
+    # Define facial region indices
+    FACIAL_REGIONS = {
+        # These are approximate regions based on MediaPipe Face Mesh indices
+        "whole_face": [
+            # Face contour/oval - same as face_oval for proper face filling
+            10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136, 172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109, 10
+        ],
+        "mouth": [
+            61, 185, 40, 39, 37, 0, 267, 269, 270, 409, 291, 375, 321, 405, 314, 17, 84, 181, 91, 146
+        ],
+        "left_eye": [
+            # Left eye outline (from the person's perspective)
+            263, 249, 390, 373, 374, 380, 381, 382, 362, 263, 466, 388, 387, 386, 385, 384, 398, 362
+        ],
+        "right_eye": [
+            # Right eye outline
+            33, 7, 163, 144, 145, 153, 154, 155, 133, 33, 246, 161, 160, 159, 158, 157, 173, 133
+        ],
+        "nose": [
+            # Nose area
+            168, 6, 197, 195, 5, 4, 1, 19, 94, 2, 164, 0, 11, 12, 13, 14, 15, 16, 17, 18, 200, 199, 175
+        ],
+        "left_eyebrow": [
+            # Left eyebrow
+            276, 283, 282, 295, 285, 300, 293, 334, 296, 336
+        ],
+        "right_eyebrow": [
+            # Right eyebrow
+            46, 53, 52, 65, 55, 70, 63, 105, 66, 107
+        ],
+        "face_oval": [
+            # Face contour/oval
+            10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136, 172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109, 10
+        ]
+    }
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        regions = list(s.FACIAL_REGIONS.keys())
+        
+        return {
+            "required": {
+                "facemesh_data": ("FACEMESH_DATA",),
+                "face_index": ("INT", {
+                    "default": 0,
+                    "min": 0,
+                    "max": 9,
+                    "step": 1,
+                    "display": "number",
+                    "description": "Which face to mask (if multiple detected)"
+                }),
+                "region": (regions, {
+                    "default": "mouth",
+                    "description": "Facial region to mask"
+                }),
+                "fill_region": ("BOOLEAN", {
+                    "default": True,
+                    "description": "Fill the region or just draw contour"
+                }),
+                "mask_value": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "display": "slider",
+                    "description": "Value/intensity of the mask (0.0-1.0)"
+                }),
+                "dilation_pixels": ("INT", {
+                    "default": 0,
+                    "min": 0,
+                    "max": 50,
+                    "step": 1,
+                    "display": "slider",
+                    "description": "Dilate mask by this many pixels"
+                }),
+                "blur_radius": ("INT", {
+                    "default": 0,
+                    "min": 0,
+                    "max": 20,
+                    "step": 1,
+                    "display": "slider",
+                    "description": "Apply gaussian blur with this radius"
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("MASK",)
+    RETURN_NAMES = ("mask",)
+    FUNCTION = "create_mask"
+    CATEGORY = "ComfyUI_FaceMesh"
+
+    def create_mask(self, facemesh_data, face_index, region, fill_region, mask_value, dilation_pixels, blur_radius):
+        """Create a mask for the specified facial region with custom intensity value"""
+        # Get number of images in batch
+        batch_size = len(facemesh_data)
+        output_masks = []
+        
+        # Process each image in the batch
+        for b in range(batch_size):
+            # Get current facemesh data
+            current_facemesh_data = facemesh_data[b]
+            
+            # If no face detected, return blank image
+            if not current_facemesh_data["face_detected"] or face_index >= current_facemesh_data["num_faces_detected"]:
+                # Create a blank mask with original image dimensions
+                h, w = current_facemesh_data["image_shape"]
+                mask_float = np.zeros((h, w), dtype=np.float32)
+                output_masks.append(mask_float)
+                continue
+            
+            # Get image shape
+            h, w = current_facemesh_data["image_shape"]
+            
+            # Create blank mask
+            mask = np.zeros((h, w), dtype=np.uint8)
+            
+            # Get the landmarks for the specified face
+            landmarks = current_facemesh_data["landmarks"][face_index]
+            
+            # Get indices for the specified region
+            region_indices = self.FACIAL_REGIONS[region]
+            
+            # Extract points for the region - directly use numpy for better performance
+            landmarks_array = np.array(landmarks)
+            if len(region_indices) > 0:
+                # Get subset of landmarks for the region
+                region_landmarks = np.array([landmarks_array[i] for i in region_indices if i < len(landmarks_array)])
+                
+                if len(region_landmarks) > 2:  # Need at least 3 points to create a contour
+                    # Convert normalized coordinates to pixel coordinates
+                    points = np.zeros((len(region_landmarks), 2), dtype=np.int32)
+                    points[:, 0] = (region_landmarks[:, 0] * w).astype(np.int32)  # x coordinates
+                    points[:, 1] = (region_landmarks[:, 1] * h).astype(np.int32)  # y coordinates
+                    
+                    # Ensure mask is contiguous in memory for OpenCV compatibility
+                    mask = np.ascontiguousarray(mask, dtype=np.uint8)
+                    
+                    # Use 255 for drawing, we'll scale to mask_value later
+                    fill_value = 255
+                    
+                    # Special handling for mouth region - create a convex hull to ensure proper filling
+                    if region == "mouth" and fill_region:
+                        # Use convex hull to get a proper fillable shape
+                        hull = cv2.convexHull(points)
+                        cv2.fillConvexPoly(mask, hull, fill_value)
+                    # For whole_face and face_oval, always fill regardless of fill_region setting
+                    elif (region == "whole_face" or region == "face_oval") and fill_region:
+                        # Use convex hull for better filling of the face
+                        hull = cv2.convexHull(points)
+                        cv2.fillConvexPoly(mask, hull, fill_value)
+                    else:
+                        # Create the mask based on fill_region option
+                        if fill_region:
+                            # Fill the polygon with white (255)
+                            cv2.fillPoly(mask, [points], fill_value)
+                        else:
+                            # Just draw the outline with white (255)
+                            cv2.polylines(mask, [points], True, fill_value, 2)
+                    
+                    # Apply dilation if specified - use optimized kernel
+                    if dilation_pixels > 0:
+                        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_pixels, dilation_pixels))
+                        mask = cv2.dilate(mask, kernel, iterations=1)
+                    
+                    # Apply gaussian blur if specified
+                    if blur_radius > 0:
+                        mask = cv2.GaussianBlur(mask, (blur_radius * 2 + 1, blur_radius * 2 + 1), 0)
+            
+            # Convert mask to float [0,1] for ComfyUI and apply mask_value
+            mask_float = mask.astype(np.float32) / 255.0 * mask_value
+            
+            # Add to output list
+            output_masks.append(mask_float)
+        
+        # Convert to batch tensor - BHW format
+        mask_batch = torch.from_numpy(np.stack(output_masks, axis=0))
+        
+        return (mask_batch,)
+
+class FaceMeshNode:
+    """
+    A node that uses MediaPipe to detect face mesh from images and draws it.
+    Optimized for real-time video processing.
+    Uses the same drawing methodology as ControlNet Aux for optimal performance.
+    """
+    def __init__(self):
+        # Initialize face mesh detector once during node creation
+        self.mp_face_mesh = mp.solutions.face_mesh
+        self.face_mesh = self.mp_face_mesh.FaceMesh(
+            static_image_mode=False,
+            max_num_faces=1,
+            refine_landmarks=False,
+            min_detection_confidence=0.5,
+            min_tracking_confidence=0.5
+        )
+        self.last_results = None
+        # Track current settings to know when to reinitialize
+        self.current_settings = {
+            "static_image_mode": False,
+            "max_num_faces": 1,
+            "refine_landmarks": False,
+            "min_detection_confidence": 0.5,
+            "min_tracking_confidence": 0.5
+        }
+        # Use the shared utils instance
+        self.utils = face_mesh_utils
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "static_image_mode": ("BOOLEAN", {
+                    "default": False,
+                    "description": "Set to true for static images, false for video streams (faster)"
+                }),
+                "max_num_faces": ("INT", {
+                    "default": 1,
+                    "min": 1,
+                    "max": 10,
+                    "step": 1,
+                    "display": "number",
+                    "description": "Maximum number of faces to detect (lower is faster)"
+                }),
+                "refine_landmarks": ("BOOLEAN", {
+                    "default": False,
+                    "description": "Refine landmarks around eyes and lips (slower but more accurate)"
+                }),
+                "min_detection_confidence": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.1,
+                    "max": 1.0,
+                    "step": 0.05,
+                    "display": "slider",
+                    "description": "Minimum confidence for face detection"
+                }),
+                "min_tracking_confidence": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.1,
+                    "max": 1.0,
+                    "step": 0.05,
+                    "display": "slider",
+                    "description": "Minimum confidence for face tracking"
+                }),
+                "draw_mesh": ("BOOLEAN", {
+                    "default": True,
+                    "description": "Draw face mesh on output image"
+                }),
+                "draw_contours": ("BOOLEAN", {
+                    "default": True,
+                    "description": "Draw contour lines connecting mesh points"
+                }),
+                "mesh_color": (["green", "red", "blue", "yellow", "white"], {
+                    "default": "green",
+                    "description": "Color of the face mesh lines"
+                }),
+                "mesh_thickness": ("FLOAT", {
+                    "default": 2.0,
+                    "min": 0.1,
+                    "max": 5.0,
+                    "step": 0.1,
+                    "display": "slider",
+                    "description": "Thickness of mesh lines"
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("FACEMESH_DATA", "IMAGE")
+    RETURN_NAMES = ("facemesh_data", "annotations")
+    FUNCTION = "detect_facemesh"
+    CATEGORY = "ComfyUI_FaceMesh"
+
+    def detect_facemesh(self, image, static_image_mode, max_num_faces, refine_landmarks, 
+                       min_detection_confidence, min_tracking_confidence,
+                       draw_mesh, draw_contours, mesh_color, mesh_thickness):
+        """Detect face mesh in the input image and optionally draw it using ControlNet Aux methodology"""
+        # Check if any parameters have changed and we need to reinitialize
+        new_settings = {
+            "static_image_mode": static_image_mode,
+            "max_num_faces": max_num_faces,
+            "refine_landmarks": refine_landmarks,
+            "min_detection_confidence": min_detection_confidence,
+            "min_tracking_confidence": min_tracking_confidence
+        }
+        
+        # Reinitialize if settings changed
+        if new_settings != self.current_settings:
+            self.face_mesh = self.mp_face_mesh.FaceMesh(
+                static_image_mode=static_image_mode,
+                max_num_faces=max_num_faces,
+                refine_landmarks=refine_landmarks,
+                min_detection_confidence=min_detection_confidence,
+                min_tracking_confidence=min_tracking_confidence
+            )
+            self.current_settings = new_settings
+        
+        # Get batch size
+        batch_size = image.shape[0]
+        
+        # Initialize empty lists to store batch results
+        facemesh_data_batch = []
+        output_images = []
+        
+        # Process each image in the batch
+        for b in range(batch_size):
+            # Get current image from batch
+            img = image[b].numpy()
+            
+            # Convert from float [0,1] to uint8 [0,255] for mediapipe
+            img_rgb = (img * 255).astype(np.uint8)
+            
+            # Ensure the image is contiguous in memory for OpenCV compatibility
+            img_rgb = np.ascontiguousarray(img_rgb, dtype=np.uint8)
+            
+            # Process the image
+            results = self.face_mesh.process(img_rgb)
+            
+            # Prepare facemesh data structure for this image
+            facemesh_data = {
+                "landmarks": [],
+                "face_detected": False,
+                "num_faces_detected": 0,
+                "tesselation": self.mp_face_mesh.FACEMESH_TESSELATION,
+                "contours": self.mp_face_mesh.FACEMESH_CONTOURS,
+                "image_shape": img.shape[:2],
+                "original_results": results  # Store original results for native drawing
+            }
+            
+            # Process detection results
+            if results and results.multi_face_landmarks:
+                facemesh_data["face_detected"] = True
+                facemesh_data["num_faces_detected"] = len(results.multi_face_landmarks)
+                
+                all_landmarks = []
+                
+                for face_idx, face_landmarks in enumerate(results.multi_face_landmarks):
+                    # Extract all landmarks at once for more efficiency
+                    face_points = [(lm.x, lm.y, lm.z) for lm in face_landmarks.landmark]
+                    all_landmarks.append(face_points)
+                
+                facemesh_data["landmarks"] = all_landmarks
+            
+            # Draw annotations if requested - on a black background for FaceMeshNode
+            if draw_mesh and facemesh_data["face_detected"]:
+                # Draw the face mesh on a black background using the shared utility
+                empty = self.utils.draw_facemesh(img_rgb, results, thickness=int(mesh_thickness))
+                
+                # Convert back to float [0,1] for ComfyUI
+                output_img = empty.astype(np.float32) / 255.0
+                
+                # Add image to output list
+                output_images.append(output_img)
+            else:
+                # Use original image if no drawing requested or no face detected
+                output_images.append(img)
+            
+            # Add facemesh data to batch results
+            facemesh_data_batch.append(facemesh_data)
+        
+        # Convert output images list to tensor with batch dimension
+        output_batch = torch.from_numpy(np.stack(output_images, axis=0))
+        
+        return (facemesh_data_batch, output_batch)
+
+# Node mapping
+NODE_CLASS_MAPPINGS = {
+    "FaceMeshNode": FaceMeshNode,
+    "FaceMeshDrawNode": FaceMeshDrawNode,
+    "FaceMeshMaskNode": FaceMeshMaskNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FaceMeshNode": "Face Mesh",
+    "FaceMeshDrawNode": "Face Mesh Draw",
+    "FaceMeshMaskNode": "Face Region Mask"
+}