add batched denoising workflow

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,287 @@
+import torch
+import comfy.model_management
+import comfy.samplers
+
+class StreamBatchSampler:
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "num_steps": ("INT", {
+                    "default": 4,
+                    "min": 1,
+                    "max": 10,
+                    "step": 1,
+                    "tooltip": "Number of denoising steps. Should match the frame buffer size."
+                }),
+            },
+        }
+    
+    RETURN_TYPES = ("SAMPLER",)
+    FUNCTION = "update"
+    CATEGORY = "StreamPack/sampling"
+    DESCRIPTION = "Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps. Also adds temportal consistency to the denoising process."
+    
+    def __init__(self):
+        self.num_steps = None
+        self.frame_buffer = []
+        self.x_t_latent_buffer = None
+        self.stock_noise = None
+        self.is_txt2img_mode = False
+        
+        # Initialize all buffers
+        self.zeros_reference = None
+        self.random_noise_buffer = None
+        self.sigmas_view_buffer = None
+        self.expanded_stock_noise = None
+        self.working_buffer = None
+        self.output_buffer = None
+    
+    def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None):
+        """Sample with staggered batch denoising steps - Optimized version"""
+        extra_args = {} if extra_args is None else extra_args
+        
+        # Get number of frames in batch and available sigmas
+        batch_size = noise.shape[0]
+        num_sigmas = len(sigmas) - 1  # Subtract 1 because last sigma is the target (0.0)
+        
+        # Reuse zeros buffer for txt2img detection
+        if self.zeros_reference is None:
+            self.zeros_reference = torch.zeros(1, device=noise.device, dtype=noise.dtype)
+        
+        # Check if noise tensor is all zeros
+        self.is_txt2img_mode = torch.abs(noise).sum() < 1e-5
+        
+        if self.is_txt2img_mode:
+            # If txt2img mode, reuse the noise tensor directly
+            if self.random_noise_buffer is None or self.random_noise_buffer.shape != noise.shape:
+                self.random_noise_buffer = torch.empty_like(noise)
+            
+            # Generate random noise in-place
+            self.random_noise_buffer.normal_()
+            x = self.random_noise_buffer  # Use pre-allocated buffer
+        else:
+            # If not txt2img, we'll still need to add noise later
+            x = noise
+        
+        # Verify batch size matches number of timesteps
+        if batch_size != num_sigmas:
+            raise ValueError(f"Batch size ({batch_size}) must match number of timesteps ({num_sigmas})")
+        
+        # Pre-compute alpha and beta terms
+        alpha_prod_t = (sigmas[:-1] / sigmas[0]).view(-1, 1, 1, 1)  # [B,1,1,1]
+        beta_prod_t = (1 - alpha_prod_t)
+        
+        # Initialize stock noise with reuse
+        if self.stock_noise is None or self.stock_noise.shape != noise[0].shape:
+            self.stock_noise = torch.empty_like(noise[0])
+            self.stock_noise.normal_()  # Generate random noise in-place
+        
+        # Pre-allocate and reuse view buffer for sigmas
+        if self.sigmas_view_buffer is None or self.sigmas_view_buffer.shape[0] != len(sigmas)-1:
+            self.sigmas_view_buffer = torch.empty((len(sigmas)-1, 1, 1, 1), 
+                                               device=sigmas.device, 
+                                               dtype=sigmas.dtype)
+        # In-place copy of sigmas view
+        self.sigmas_view_buffer.copy_(sigmas[:-1].view(-1, 1, 1, 1))
+        
+        # Eliminate unsqueeze allocation by pre-expanding stock noise
+        if self.expanded_stock_noise is None or self.expanded_stock_noise.shape[0] != batch_size:
+            self.expanded_stock_noise = self.stock_noise.expand(batch_size, *self.stock_noise.shape)
+        
+        # Apply noise with pre-allocated buffers - no new memory allocation
+        if not self.is_txt2img_mode:  # Already handled txt2img case above
+            # If we need a working buffer separate from noise input:
+            if id(x) == id(noise):  # They're the same object, need a separate buffer
+                if self.working_buffer is None or self.working_buffer.shape != noise.shape:
+                    self.working_buffer = torch.empty_like(noise)
+                x = self.working_buffer
+                # Add noise to input
+                torch.add(noise, self.expanded_stock_noise * self.sigmas_view_buffer, out=x)
+        
+        # Initialize and manage latent buffer with memory optimization
+        if (self.x_t_latent_buffer is None or self.is_txt2img_mode) and num_sigmas > 1:
+            # Pre-allocate or resize as needed
+            if self.x_t_latent_buffer is None or self.x_t_latent_buffer.shape != x[0].shape:
+                self.x_t_latent_buffer = torch.empty_like(x[0])
+            self.x_t_latent_buffer.copy_(x[0])
+        
+        # Use buffer for first frame to maintain temporal consistency
+        if num_sigmas > 1:
+            x[0].copy_(self.x_t_latent_buffer)
+        
+        with torch.no_grad():
+            # Process all frames in parallel
+            sigma_batch = sigmas[:-1]
+            
+            denoised_batch = model(x, sigma_batch, **extra_args)
+            
+            # Update buffer with intermediate results
+            if num_sigmas > 1:
+                # Store result from first frame as buffer for next iteration
+                self.x_t_latent_buffer.copy_(denoised_batch[0])  # In-place update
+                
+                # Pre-allocate output buffer
+                if self.output_buffer is None or self.output_buffer.shape != (1, *denoised_batch[-1].shape):
+                    self.output_buffer = torch.empty(1, *denoised_batch[-1].shape, 
+                                                  device=denoised_batch.device,
+                                                  dtype=denoised_batch.dtype)
+                # Copy the result directly to pre-allocated buffer
+                self.output_buffer[0].copy_(denoised_batch[-1])
+                x_0_pred_out = self.output_buffer
+            else:
+                x_0_pred_out = denoised_batch
+                self.x_t_latent_buffer = None
+            
+            # 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': denoised_batch[-1:]})
+        
+        return x_0_pred_out
+    
+    def update(self, num_steps=4):
+        """Create sampler with specified settings"""
+        self.num_steps = num_steps
+        sampler = comfy.samplers.KSAMPLER(self.sample)
+        return (sampler,)
+
+
+class StreamScheduler:
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "t_index_list": ("STRING", {
+                    "default": "0,16,32,49",
+                    "tooltip": "Comma-separated list of timesteps to actually use for denoising. Examples: '32,45' for img2img or '0,16,32,45' for txt2img"
+                }),
+                "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_TYPES = ("SIGMAS",)
+    FUNCTION = "update"
+    CATEGORY = "StreamPack/sampling"
+    DESCRIPTION = "Implements StreamDiffusion's efficient timestep selection. Use in conjunction with StreamBatchSampler."
+    
+    def update(self, model, t_index_list="32,45", num_inference_steps=50):
+        # Get model's sampling parameters
+        model_sampling = model.get_model_object("model_sampling")
+        
+        # Parse timestep list
+        try:
+            t_index_list = [int(t.strip()) for t in t_index_list.split(",")]
+        except ValueError as e:
+            
+            t_index_list = [32, 45]
+            
+        # 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, but in reverse order
+        # This ensures we go from high noise to low noise
+        selected_sigmas = []
+        for t in sorted(t_index_list, reverse=True):  # Sort in reverse to go from high noise to low
+            if t < 0 or t >= num_inference_steps:
+                
+                continue
+            selected_sigmas.append(float(full_sigmas[t]))
+            
+        # Add final sigma
+        selected_sigmas.append(0.0)
+        
+        # Convert to tensor and move to appropriate device
+        selected_sigmas = torch.FloatTensor(selected_sigmas).to(comfy.model_management.get_torch_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/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"
+}