chore: update .gitignore to include Python packaging artifacts (#46)

* chore(pyproject): Add pyproject.toml for Custom Node Registry
* use dynamic dependencies

---------

Co-authored-by: snomiao <snomiao+comfy-pr@gmail.com>
Co-authored-by: John | Elite Encoder <john@eliteencoder.net>

.github/workflows/publish.yml

@@ -0,0 +1,28 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+      - master
+    paths:
+      - "pyproject.toml"
+
+permissions:
+  issues: write
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'livepeer' }}
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+        with:
+          submodules: true
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@v1
+        with:
+          ## Add your own personal access token to your Github Repository secrets and reference it here.
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

@@ -0,0 +1,22 @@
+# Misc
+__pycache__
+*.pyc
+
+# Python packaging
+*.egg-info/
+*.egg-info
+build/
+dist/
+
+# IDEs
+.vscode
+*.code-workspace
+
+# Model weights
+*.pb
+*.pth
+*.safetensors
+*.pt
+*.pth.tar
+*.pth.tar.gz
+*.pth.tar.xz

__init__.py

@@ -0,0 +1,44 @@
+"""Expose all nodes in the Stream Pack."""
+
+import logging
+import pathlib
+import importlib
+import traceback
+
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger(__name__)
+
+
+def load_nodes() -> tuple[dict[str, type], dict[str, str]]:
+    """Load all nodes in the Stream Pack."""
+    nodes_dir = pathlib.Path(__file__).parent / "node_wrappers"
+    node_class_mappings, node_display_name_mappings = {}, {}
+
+    # Dynamically import all Python modules in the node_wrappers directory.
+    for module_path in nodes_dir.iterdir():
+        if (
+            module_path.is_file()
+            and module_path.suffix == ".py"
+            and module_path.stem != "__init__"
+        ):
+            try:
+                module_name = f"{__package__}.node_wrappers.{module_path.stem}"
+                module = importlib.import_module(module_name)
+
+                # Update mappings if defined in the module.
+                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)
+            except Exception:
+                format_exc = traceback.format_exc()
+                log.error(f"Failed to load module {module_name}:\n{format_exc}")
+
+    return node_class_mappings, node_display_name_mappings
+
+
+# Dynamically collect mappings from submodules.
+NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = load_nodes()
+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-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"
+    }
+  }
+}

examples/workflows/sd15-tensorrt-feature-bank-multi-cnet-depthanything-face-mesh-api.json

@@ -0,0 +1,352 @@
+{
+  "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)"
+    }
+  },
+  "7": {
+    "inputs": {
+      "seed": 951969040799634,
+      "steps": 2,
+      "cfg": 1,
+      "sampler_name": "lcm",
+      "scheduler": "normal",
+      "denoise": 1,
+      "model": [
+        "41",
+        0
+      ],
+      "positive": [
+        "38",
+        0
+      ],
+      "negative": [
+        "37",
+        0
+      ],
+      "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": 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": [
+        "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"
+    }
+  },
+  "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": {
+      "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"
+    }
+  }
+}

examples/workflows/sd15-tensorrt-multi-cnet-depthanything-face-mesh-api.json

@@ -0,0 +1,334 @@
+{
+  "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)"
+    }
+  },
+  "7": {
+    "inputs": {
+      "seed": 951969040799634,
+      "steps": 2,
+      "cfg": 1,
+      "sampler_name": "lcm",
+      "scheduler": "normal",
+      "denoise": 1,
+      "model": [
+        "3",
+        0
+      ],
+      "positive": [
+        "38",
+        0
+      ],
+      "negative": [
+        "37",
+        0
+      ],
+      "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": 0.9000000000000001,
+      "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"
+    }
+  },
+  "18": {
+    "inputs": {
+      "stop_at_clip_layer": -2,
+      "clip": [
+        "36",
+        0
+      ]
+    },
+    "class_type": "CLIPSetLastLayer",
+    "_meta": {
+      "title": "CLIP Set Last Layer"
+    }
+  },
+  "19": {
+    "inputs": {
+      "strength": 0.9000000000000001,
+      "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"
+    },
+    "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.3500000000000001,
+      "min_tracking_confidence": 0.3500000000000001,
+      "draw_mesh": true,
+      "draw_contours": true,
+      "mesh_color": "green",
+      "mesh_thickness": 2,
+      "image": [
+        "26",
+        0
+      ]
+    },
+    "class_type": "FaceMeshNode",
+    "_meta": {
+      "title": "Face Mesh"
+    }
+  }
+}

examples/workflows/super-res-fscnn-upscale-2x-gpu-api.json

@@ -0,0 +1,50 @@
+{
+  "1": {
+    "inputs": {
+      "model_type": "FSRCNN",
+      "scale_factor": "2"
+    },
+    "class_type": "SuperResolutionModelLoader",
+    "_meta": {
+      "title": "SR Model Loader"
+    }
+  },
+  "2": {
+    "inputs": {
+      "use_cuda": "True",
+      "image": [
+        "3",
+        0
+      ],
+      "sr_model": [
+        "1",
+        0
+      ]
+    },
+    "class_type": "SuperResolutionUpscale",
+    "_meta": {
+      "title": "SR Upscale"
+    }
+  },
+  "3": {
+    "inputs": {
+      "image": "example-512x512.png"
+    },
+    "class_type": "LoadImage",
+    "_meta": {
+      "title": "Load Image"
+    }
+  },
+  "4": {
+    "inputs": {
+      "images": [
+        "2",
+        0
+      ]
+    },
+    "class_type": "PreviewImage",
+    "_meta": {
+      "title": "Preview Image"
+    }
+  }
+}

node_wrappers/audio_transcription.py

@@ -0,0 +1,23 @@
+"""
+Real-time Audio Transcription Node for ComfyUI.
+
+This node buffers audio segments and performs transcription using faster-whisper,
+with controlled output timing to prevent message flooding.
+"""
+
+# Import node classes
+from ..src.audio_transcription.transcription_nodes import AudioTranscriptionNode
+from ..src.audio_transcription.srt_generator_node import SRTGeneratorNode
+
+# Register the node
+NODE_CLASS_MAPPINGS = {
+    "AudioTranscriptionNode": AudioTranscriptionNode, 
+    "SRTGeneratorNode": SRTGeneratorNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "AudioTranscriptionNode": "Audio Transcription (Real-time)", 
+    "SRTGeneratorNode": "SRT Generator"
+}
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

node_wrappers/face_mesh.py

@@ -0,0 +1,21 @@
+"""Provides nodes for facial landmark detection and utilities using MediaPipe FaceMesh.
+Inspired by the ComfyUI ControlNet Auxiliary MediaPipe Face node.
+
+References:
+    https://github.com/Fannovel16/comfyui_controlnet_aux/blob/main/src/custom_controlnet_aux/mediapipe_face/mediapipe_face_common.py
+"""
+
+from ..src.media_pipe.face_mesh 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",
+}
+
+__all__ = ["FaceMeshNode", "FaceMeshDrawNode", "FaceMeshMaskNode"]

node_wrappers/feature_bank.py

@@ -0,0 +1,20 @@
+"""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"]

node_wrappers/super_resolution.py

@@ -0,0 +1,16 @@
+# Import node classes
+from ..src.super_resolution.nodes.sr_nodes import SuperResolutionModelLoader, SuperResolutionUpscale
+
+# Define NODE_CLASS_MAPPINGS for ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "SuperResolutionModelLoader": SuperResolutionModelLoader,
+    "SuperResolutionUpscale": SuperResolutionUpscale,
+}
+
+# Define NODE_DISPLAY_NAME_MAPPINGS for ComfyUI
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SuperResolutionModelLoader": "SR Model Loader",
+    "SuperResolutionUpscale": "SR Upscale",
+}
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

prestartup_script.py

@@ -0,0 +1,22 @@
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def check_opencv_cuda():
+    try:
+        import cv2
+        # Check if CUDA is already available
+        logger.info("ComfyUI-Stream-Pack: Checking OpenCV CUDA availability...")
+        if hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0:
+            logger.info("OpenCV CUDA is installed and working")
+            return True
+        else:
+            logger.warning("OpenCV CUDA is not available or not working. Super Resolution node may not work")
+            return False
+    except ImportError:
+        logger.warning("OpenCV is not installed")
+        return False
+
+check_opencv_cuda()

pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "stream-pack"
+description = "A suite of custom ComfyUI nodes for building real-time video and audio workflows using ComfyStream."
+version = "1.0.0"
+license = {file = "LICENSE"}
+dynamic = ["dependencies"]
+
+[project.urls]
+Repository = "https://github.com/livepeer/ComfyUI-Stream-Pack"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "livepeer-comfystream"
+DisplayName = "ComfyUI-Stream-Pack"
+Icon = "https://raw.githubusercontent.com/livepeer/comfystream-docs/main/logo/icon-light-120px.svg"
+
+[tool.setuptools.dynamic]
+dependencies = {file = ["requirements.txt"]}

requirements.txt

@@ -0,0 +1,8 @@
+mediapipe
+torch
+diffusers
+numpy
+requests
+opencv-contrib-python>=4.8.0
+faster_whisper
+scipy

src/__init__.py

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

src/audio_transcription/__init__.py

@@ -0,0 +1 @@
+"""Audio transcription package for ComfyUI real-time transcription capabilities."""

src/audio_transcription/srt_generator_node.py

@@ -0,0 +1,237 @@
+"""
+SRT subtitle generation node for ComfyStream.
+Based on project-transcript SRT generation capabilities.
+"""
+
+import logging
+from datetime import timedelta
+from typing import List, Dict, Any, Union
+import json
+
+logger = logging.getLogger(__name__)
+
+
+class SRTGeneratorNode:
+    """
+    Generate SRT subtitle files from transcription with timing information.
+    Compatible with AudioTranscriptionNode output when word timestamps are enabled.
+    """
+    
+    CATEGORY = "text_utils"
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("srt_content",)
+    
+    def __init__(self):
+        self.subtitle_counter = 1
+        
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "transcription_data": ("STRING", {
+                    "tooltip": "Transcription text or JSON with timing data from AudioTranscriptionNode"
+                }),
+                "segment_start_time": ("FLOAT", {
+                    "default": 0.0,
+                    "min": 0.0,
+                    "max": 86400.0,  # 24 hours max
+                    "step": 0.001,
+                    "tooltip": "Start time of this segment in seconds (for absolute timing)"
+                }),
+                "segment_duration": ("FLOAT", {
+                    "default": 3.0,
+                    "min": 0.1,
+                    "max": 60.0,
+                    "step": 0.1,
+                    "tooltip": "Duration of this segment in seconds"
+                }),
+                "use_absolute_time": ("BOOLEAN", {
+                    "default": False,
+                    "tooltip": "Use absolute timing (True) or segment-relative timing (False)"
+                }),
+                "minimum_duration": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.1,
+                    "max": 10.0,
+                    "step": 0.1,
+                    "tooltip": "Minimum subtitle duration in seconds"
+                })
+            }
+        }
+    
+    FUNCTION = "generate_srt"
+    
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        return float("nan")
+    
+    def generate_srt(self, 
+                    transcription_data: str,
+                    segment_start_time: float = 0.0,
+                    segment_duration: float = 3.0,
+                    use_absolute_time: bool = False,
+                    minimum_duration: float = 1.0) -> tuple:
+        """
+        Generate SRT subtitle content from transcription data.
+        
+        Args:
+            transcription_data: Text transcription or JSON with timing data
+            segment_start_time: Start time of this segment (for absolute timing)
+            segment_duration: Duration of this segment
+            use_absolute_time: Whether to use absolute or segment-relative timing
+            minimum_duration: Minimum subtitle duration
+            
+        Returns:
+            Tuple containing SRT formatted content
+        """
+        try:
+            if not transcription_data or not transcription_data.strip() or len(transcription_data.strip()) < 5:
+                logger.debug(f"Empty or minimal transcription data: '{transcription_data}'")
+                return ("",)
+            
+            # Pass through warmup sentinel values (they'll be filtered later in the pipeline)
+            if "__WARMUP_SENTINEL__" in transcription_data:
+                logger.debug("Passing through warmup sentinel for pipeline detection")
+                return (transcription_data,)
+            
+            # Try to parse as JSON first (if AudioTranscriptionNode supports structured output)
+            segments = self._parse_transcription_data(transcription_data)
+            
+            if not segments:
+                logger.debug("No valid segments found in transcription data")
+                return ("",)
+            
+            # Generate SRT content
+            srt_content = self._generate_srt_from_segments(
+                segments, 
+                segment_start_time,
+                segment_duration,
+                use_absolute_time,
+                minimum_duration
+            )
+            
+            logger.debug(f"Generated SRT with {len(segments)} segments")
+            return (srt_content,)
+            
+        except Exception as e:
+            logger.error(f"Error generating SRT: {e}")
+            return ("",)
+    
+    def _parse_transcription_data(self, data: str) -> List[Dict[str, Any]]:
+        """
+        Parse transcription data into segments.
+        Enhanced to handle AudioTranscriptionNode JSON output formats.
+        """
+        segments = []
+        
+        try:
+            # Try parsing as JSON first
+            parsed_data = json.loads(data)
+            
+            if isinstance(parsed_data, list):
+                # Handle both segment-level and word-level JSON arrays
+                for item in parsed_data:
+                    if isinstance(item, dict):
+                        # Check if it's word-level data (has 'word' field) or segment-level (has 'text' field)
+                        if 'word' in item:
+                            # Word-level JSON from AudioTranscriptionNode (json_words format)
+                            segments.append({
+                                'start': item.get('start', 0.0),
+                                'end': item.get('end', 1.0),
+                                'text': item.get('word', '').strip()
+                            })
+                        elif 'text' in item:
+                            # Segment-level JSON from AudioTranscriptionNode (json_segments format)
+                            segments.append({
+                                'start': item.get('start', 0.0),
+                                'end': item.get('end', 1.0),
+                                'text': item.get('text', '').strip()
+                            })
+                            
+            elif isinstance(parsed_data, dict):
+                # Single segment or word
+                if 'word' in parsed_data:
+                    segments.append({
+                        'start': parsed_data.get('start', 0.0),
+                        'end': parsed_data.get('end', 1.0),
+                        'text': parsed_data.get('word', '').strip()
+                    })
+                elif 'text' in parsed_data:
+                    segments.append({
+                        'start': parsed_data.get('start', 0.0),
+                        'end': parsed_data.get('end', 1.0),
+                        'text': parsed_data.get('text', '').strip()
+                    })
+                    
+        except (json.JSONDecodeError, TypeError):
+            # Fallback to plain text - create a single segment
+            if data.strip():
+                segments.append({
+                    'start': 0.0,
+                    'end': 3.0,  # Default 3-second duration
+                    'text': data.strip()
+                })
+        
+        return [seg for seg in segments if seg['text']]  # Filter empty text
+    
+    def _generate_srt_from_segments(self, 
+                                   segments: List[Dict[str, Any]],
+                                   segment_start_time: float,
+                                   segment_duration: float,
+                                   use_absolute_time: bool,
+                                   minimum_duration: float) -> str:
+        """
+        Generate SRT content from parsed segments and wrap it in a JSON format.
+        """
+        subtitles = []
+        
+        for i, segment in enumerate(segments):
+            text = segment['text'].strip()
+            if not text:
+                continue
+            
+            # Calculate timing
+            if use_absolute_time:
+                # Absolute timing: add segment start time
+                start_time = segment_start_time + segment['start']
+                end_time = segment_start_time + segment['end']
+            else:
+                # Segment-relative timing
+                start_time = segment['start']
+                end_time = segment['end']
+            
+            # Ensure minimum duration
+            if end_time - start_time < minimum_duration:
+                end_time = start_time + minimum_duration
+            
+            # Format timing
+            start_srt = self._seconds_to_srt_time(start_time)
+            end_srt = self._seconds_to_srt_time(end_time)
+            
+            # Add subtitle entry to list
+            subtitles.append({
+                "id": self.subtitle_counter,
+                "start": start_srt,
+                "end": end_srt,
+                "text": text
+            })
+            
+            self.subtitle_counter += 1
+        
+        # Return as a JSON string
+        return json.dumps(subtitles, indent=2)
+    
+    def _seconds_to_srt_time(self, seconds: float) -> str:
+        """
+        Convert seconds to SRT time format (HH:MM:SS,mmm).
+        Compatible with project-transcript format.
+        """
+        td = timedelta(seconds=seconds)
+        total_seconds = int(td.total_seconds())
+        milliseconds = int((seconds - total_seconds) * 1000)
+        
+        hours = total_seconds // 3600
+        minutes = (total_seconds % 3600) // 60
+        secs = total_seconds % 60
+        
+        return f"{hours:02d}:{minutes:02d}:{secs:02d},{milliseconds:03d}"

src/audio_transcription/transcription_nodes.py

@@ -0,0 +1,596 @@
+"""
+Real-time Audio Transcription Node for ComfyUI.
+
+This node buffers audio segments and performs transcription using faster-whisper,
+with controlled output timing to prevent message flooding.
+"""
+
+import asyncio
+import json
+import logging
+import time
+import tempfile
+import os
+import numpy as np
+import torch
+from typing import Optional, List, Deque, Dict, Any
+from collections import deque
+from dataclasses import dataclass
+from pathlib import Path
+from queue import Queue, Empty
+from faster_whisper import WhisperModel
+from scipy import signal
+
+logger = logging.getLogger(__name__)
+
+
+class CriticalTranscriptionError(Exception):
+    """Raised for unrecoverable errors where the stream should stop."""
+    pass
+
+
+@dataclass
+class TranscriptionSegment:
+    """Represents a transcribed segment with timing information."""
+    start: float  # Start time in seconds (buffer-relative)
+    end: float    # End time in seconds (buffer-relative)
+    text: str     # Transcribed text
+    confidence: float = 0.0  # Confidence score
+
+
+class AudioTranscriptionNode:
+    """
+    Real-time audio transcription node that buffers audio and outputs 
+    transcribed text on a controlled schedule to prevent message flooding.
+    """
+    
+    # ------------------------------------------------------------------
+    # Configurable class attributes (easy future exposure via UI):
+    #   TRANSCRIPTION_QUEUE_MAXLEN: Max queued outputs waiting for timing window
+    #   WORD_TIMESTAMPS / BEAM_SIZE / BEST_OF : Whisper inference params
+    #   DEFAULT_VAD              : Default Voice Activity Detection toggle
+    #   MAX_BUFFER_DURATION      : Safety limit to prevent memory exhaustion
+    #   MAX_BUFFER_SAMPLES       : Hard limit on buffer size in samples
+    #   
+    #   RECOMMENDED CONFIGURATIONS:
+    #   - Fast response: accumulation_duration=2.0, audio_chunk_size_ms=1000.0
+    #   - Balanced: accumulation_duration=4.0, audio_chunk_size_ms=2000.0  
+    #   - High quality: accumulation_duration=8.0, audio_chunk_size_ms=4000.0
+    # ------------------------------------------------------------------
+    TRANSCRIPTION_QUEUE_MAXLEN = 100
+    WORD_TIMESTAMPS = True
+    BEAM_SIZE = 1
+    BEST_OF = 1
+    DEFAULT_VAD = True
+    
+    # Safety limits to prevent unbounded buffer growth
+    MAX_BUFFER_DURATION = 60.0  # Maximum 60 seconds of audio in buffer
+    MAX_BUFFER_SAMPLES = 16000 * 60  # 60 seconds at 16kHz (960k samples)
+
+    CATEGORY = "audio_utils"
+    RETURN_TYPES = ("STRING",)
+    
+    def __init__(self):
+        # Smart accumulation buffer for better transcription quality (optimized for real-time)
+        self.sample_rate = None
+        self.accumulation_buffer = np.empty(0, dtype=np.float32)
+        self.accumulation_duration = 0.0  # Current buffer duration in seconds
+        self.target_accumulation_duration = 3.0  # Target duration for optimal real-time Whisper results
+        
+        # Whisper model (per-instance lazy load)
+        self._whisper_model = None
+        
+        # Transcription queue for controlled output
+        self.transcription_queue: Deque[str] = deque(maxlen=self.TRANSCRIPTION_QUEUE_MAXLEN)
+        self.last_output_time = 0.0
+        
+        # Processing state
+        self.total_audio_processed = 0.0  # Total audio time processed
+        self.transcription_count = 0
+        
+
+        
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "transcription_interval": ("FLOAT", {
+                    "default": 2.0, 
+                    "min": 1.0, 
+                    "max": 10.0,
+                    "step": 0.5,
+                    "tooltip": "Minimum seconds between transcription outputs (optimized for real-time)"
+                }),
+                "accumulation_duration": ("FLOAT", {
+                    "default": 3.0,
+                    "min": 2.0,
+                    "max": 10.0,
+                    "step": 0.5,
+                    "tooltip": "Audio accumulation duration for optimal Whisper transcription (reduced for real-time: shorter = faster output, longer = better quality)"
+                }),
+                "whisper_model": (["tiny", "base", "small", "medium", "large-v2"], {
+                    "default": "base",
+                    "tooltip": "Whisper model size (larger = more accurate but slower)"
+                }),
+                "language": (["auto", "en", "es", "fr", "de", "it", "pt", "ru", "ja", "ko", "zh"], {
+                    "default": "auto",
+                    "tooltip": "Language for transcription (auto = auto-detect)"
+                }),
+                "enable_vad": ("BOOLEAN", {
+                    "default": cls.DEFAULT_VAD,
+                    "tooltip": "Enable Voice Activity Detection to filter silence"
+                })
+            },
+            "optional": {
+                "output_format": (["text", "json_segments", "json_words"], {
+                    "default": "json_segments",
+                    "tooltip": "Output format: text (simple), json_segments (with timing), json_words (word-level timing)"
+                })
+            }
+        }
+    
+    FUNCTION = "execute"
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        return float("nan")
+    
+
+
+    def _load_whisper_model_now(self, model_size: str):
+        """Load the Whisper model using faster-whisper's automatic download."""
+        # Use CPU for compatibility, can be changed to CUDA if available
+        device = "cpu"
+        try:
+            # Try CUDA first if available
+            import torch  # type: ignore
+            if torch.cuda.is_available():
+                device = "cuda"
+                compute_type = "float16"
+            else:
+                compute_type = "int8"
+        except Exception:
+            compute_type = "int8"
+
+        logger.info(f"Loading Whisper model via faster-whisper: {model_size}")
+        try:
+            self._whisper_model = WhisperModel(
+                model_size,
+                device=device,
+                compute_type=compute_type
+            )
+            # Attach size name for change detection
+            self._whisper_model.model_size_name = model_size  # type: ignore[attr-defined]
+            logger.info(f"Whisper model '{model_size}' loaded successfully on {device}")
+        except (MemoryError, RuntimeError) as e:
+            logger.error(f"Critical: Failed to load Whisper model '{model_size}' due to resource issue: {e}")
+            raise CriticalTranscriptionError(f"Model load failed (resources): {e}") from e
+        except OSError as e:
+            logger.error(f"Critical: Filesystem/model download issue for '{model_size}': {e}")
+            raise CriticalTranscriptionError(f"Filesystem/model download failure: {e}") from e
+        except Exception as e:
+            logger.error(f"Critical: Unexpected error loading model '{model_size}': {e}")
+            raise CriticalTranscriptionError(f"Unexpected model load error: {e}") from e
+
+    def _ensure_model_loaded(self, model_size: str) -> bool:
+        """Ensure the Whisper model is loaded. Returns True if a fresh load occurred."""
+        if (self._whisper_model is None or 
+            getattr(self._whisper_model, 'model_size_name', '') != model_size):
+            logger.info(f"Loading Whisper model: {model_size}")
+            self._load_whisper_model_now(model_size)
+            return True
+        return False
+    
+    def _process_audio_chunk(self, audio_input: np.ndarray, sample_rate: int, accumulation_duration: float):
+        """
+        Process a single audio chunk and accumulate it for optimal Whisper transcription.
+        Returns True when we have enough audio for transcription.
+        """
+        # Initialize sample rate on first call
+        if self.sample_rate is None:
+            self.sample_rate = sample_rate
+            self.target_accumulation_duration = accumulation_duration
+            logger.info(f"Initialized transcription node with sample rate: {sample_rate}Hz, target accumulation: {accumulation_duration}s")
+        
+        # Skip empty frames
+        if audio_input is None or audio_input.size == 0:
+            return False
+        
+        # Normalize audio for Whisper
+        normalized_audio = self._normalize_audio_for_whisper(audio_input)
+        
+        # Add to accumulation buffer
+        if self.accumulation_buffer.size == 0:
+            self.accumulation_buffer = normalized_audio.copy()
+        else:
+            self.accumulation_buffer = np.concatenate([self.accumulation_buffer, normalized_audio])
+        
+        # Update accumulation duration
+        self.accumulation_duration = len(self.accumulation_buffer) / self.sample_rate
+        
+        # Safety check: enforce buffer limits to prevent memory exhaustion
+        self._enforce_buffer_limits()
+        
+        # Check if we have enough audio for transcription
+        ready = self.accumulation_duration >= self.target_accumulation_duration
+        
+        if ready:
+            logger.debug(f"Audio accumulation ready: {self.accumulation_duration:.2f}s >= {self.target_accumulation_duration}s ({len(self.accumulation_buffer)} samples)")
+        
+        return ready
+    
+    def _normalize_audio_for_whisper(self, audio_input: np.ndarray) -> np.ndarray:
+        """
+        Minimal audio normalization to avoid artifacts.
+        Assumes LoadAudioTensorStream has already done most processing.
+        """
+        if audio_input is None or audio_input.size == 0:
+            return np.array([], dtype=np.int16)
+        
+        # Ensure numpy array
+        if not isinstance(audio_input, np.ndarray):
+            audio_input = np.array(audio_input)
+        
+        # If it's already int16 and 1D, return as-is
+        if audio_input.dtype == np.int16 and audio_input.ndim == 1:
+            return audio_input
+        
+        # Handle multi-channel by taking first channel only
+        if audio_input.ndim > 1:
+            audio_input = audio_input.flatten() if audio_input.size < 1000 else audio_input[:, 0] if audio_input.shape[1] > audio_input.shape[0] else audio_input[0, :]
+        
+        # Convert to int16 only if necessary
+        if audio_input.dtype != np.int16:
+            if audio_input.dtype in [np.float32, np.float64]:
+                # Very conservative conversion to avoid clipping artifacts
+                audio_input = np.clip(audio_input, -0.99, 0.99)  # Leave headroom
+                audio_input = (audio_input * 32000).astype(np.int16)  # Slightly lower scale
+            else:
+                audio_input = audio_input.astype(np.int16)
+        
+        return audio_input
+    
+    def _enforce_buffer_limits(self):
+        """
+        Enforce safety limits on the accumulation buffer to prevent memory exhaustion.
+        Trims buffer from the beginning if limits are exceeded.
+        """
+        if self.sample_rate is None or self.accumulation_buffer.size == 0:
+            return
+        
+        # Check duration limit
+        if self.accumulation_duration > self.MAX_BUFFER_DURATION:
+            # Calculate how many samples to keep (keep the most recent audio)
+            samples_to_keep = int(self.MAX_BUFFER_DURATION * self.sample_rate)
+            if samples_to_keep < len(self.accumulation_buffer):
+                # Trim from the beginning, keep the end
+                self.accumulation_buffer = self.accumulation_buffer[-samples_to_keep:]
+                self.accumulation_duration = len(self.accumulation_buffer) / self.sample_rate
+                logger.warning(f"Buffer duration limit exceeded, trimmed to {self.accumulation_duration:.2f}s (kept most recent {self.MAX_BUFFER_DURATION}s)")
+        
+        # Check sample count limit (secondary safety check)
+        if len(self.accumulation_buffer) > self.MAX_BUFFER_SAMPLES:
+            self.accumulation_buffer = self.accumulation_buffer[-self.MAX_BUFFER_SAMPLES:]
+            self.accumulation_duration = len(self.accumulation_buffer) / self.sample_rate
+            logger.warning(f"Buffer sample limit exceeded, trimmed to {len(self.accumulation_buffer)} samples ({self.accumulation_duration:.2f}s)")
+
+    
+    def _transcribe_accumulated_audio(self, model_size: str, language: str, enable_vad: bool, batch_size: int, output_format: str = "json_segments") -> Optional[str]:
+        """Transcribe the accumulated audio buffer and return combined text."""
+        if self.accumulation_buffer.size == 0:
+            return None
+        try:
+            # Ensure model is loaded; if it was freshly loaded, emit a single warmup sentinel
+            send_warmup_sentinel = self._ensure_model_loaded(model_size)
+            # Emit a single warmup sentinel immediately after model loading completes
+            if send_warmup_sentinel:
+                if output_format == "text":
+                    result = "__WARMUP_SENTINEL__"
+                elif output_format == "json_segments":
+                    result = json.dumps([{ "start": 0.0, "end": 1.0, "text": "__WARMUP_SENTINEL__"}])
+                elif output_format == "json_words":
+                    result = json.dumps([{ "start": 0.0, "end": 1.0, "word": "__WARMUP_SENTINEL__", "probability": 1.0}])
+                else:
+                    result = "__WARMUP_SENTINEL__"
+                logger.debug("Warmup sentinel emitted once after model load")
+                return result
+
+            # Use the accumulated audio buffer
+            audio_chunk = self.accumulation_buffer
+            
+            # Validate audio buffer quality before transcription
+            if len(audio_chunk) == 0:
+                logger.warning("Audio buffer is empty")
+                return None
+            
+            # Check for audio artifacts (all zeros, extreme values, etc.)
+            if np.all(audio_chunk == 0):
+                return ""
+            
+            # Check for clipping or extreme values that might cause hallucinations
+            if np.abs(audio_chunk).max() < 100:  # Very quiet audio
+                return ""
+
+            # Save audio to temporary WAV file for whisper
+            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
+                temp_path = temp_file.name
+
+            try:
+                # Write clean WAV file to avoid format conversion artifacts
+                try:
+                    from scipy.io.wavfile import write  # type: ignore
+                    # Ensure audio is properly scaled for WAV format
+                    clean_audio = np.copy(audio_chunk)
+                    write(temp_path, int(self.sample_rate or 16000), clean_audio)
+                except ImportError:
+                    # Fallback to manual WAV writing
+                    self._write_wav_file(temp_path, audio_chunk, int(self.sample_rate or 16000))
+
+                language_code = None if language == "auto" else language
+                segments, info = self._whisper_model.transcribe(  # type: ignore[attr-defined]
+                    temp_path,
+                    language=language_code,
+                    word_timestamps=self.WORD_TIMESTAMPS,
+                    vad_filter=enable_vad,
+                    beam_size=self.BEAM_SIZE,
+                    best_of=self.BEST_OF
+                )
+
+                # Simple format output without complex deduplication (causes accuracy issues)
+                if output_format == "text":
+                    texts = [s.text.strip() for s in segments if s.text and s.text.strip()]
+                    result = " ".join(texts)
+                elif output_format == "json_segments":
+                    seg_list = [
+                        {"start": s.start, "end": s.end, "text": s.text.strip()}
+                        for s in segments if s.text and s.text.strip()
+                    ]
+                    result = json.dumps(seg_list, ensure_ascii=False)
+                elif output_format == "json_words":
+                    words = []
+                    for s in segments:
+                        s_words = getattr(s, 'words', None)
+                        if s_words:
+                            for w in s_words:
+                                if w.word.strip():
+                                    words.append({
+                                        "start": w.start,
+                                        "end": w.end,
+                                        "word": w.word.strip(),
+                                        "probability": getattr(w, 'probability', 1.0)
+                                    })
+                        elif s.text and s.text.strip():
+                            words.append({
+                                "start": s.start,
+                                "end": s.end,
+                                "word": s.text.strip(),
+                                "probability": 1.0
+                            })
+                    result = json.dumps(words, ensure_ascii=False)
+                else:
+                    result = None
+
+                if not result or (output_format in ["json_segments", "json_words"] and result == "[]"):
+                    result = ""  # Return empty string for ComfyUI compatibility
+                    logger.debug(f"No transcription produced, returning empty")
+
+
+                # Metrics
+                if self.sample_rate and self.accumulation_buffer.size > 0:
+                    self.total_audio_processed += len(self.accumulation_buffer) / self.sample_rate
+                self.transcription_count += 1
+
+                # Clear the accumulation buffer after processing (no overlap to prevent duplication)
+                self.accumulation_buffer = np.empty(0, dtype=np.float32)
+                self.accumulation_duration = 0.0
+                logger.debug("Cleared accumulation buffer for next cycle")
+
+                if result:
+                    logger.debug(
+                        f"Transcribed chunk {self.transcription_count} ({output_format}): '{str(result)[:50]}...' ({len(str(result))} chars)"
+                    )
+                else:
+                    logger.debug(f"Transcribed chunk {self.transcription_count} produced no result")
+                return result
+            finally:
+                try:
+                    os.unlink(temp_path)
+                except FileNotFoundError:
+                    logger.debug(f"Temp transcription file already removed: {temp_path}")
+                except PermissionError as e:
+                    logger.warning(f"Permission error removing temp file {temp_path}: {e}")
+                except OSError as e:
+                    logger.warning(f"OS error removing temp file {temp_path}: {e}")
+                except Exception as e:
+                    logger.warning(f"Unexpected error removing temp file {temp_path}: {e}")
+        except CriticalTranscriptionError:
+            # Propagate critical to stop workflow
+            raise
+        except (MemoryError, OSError) as e:
+            logger.error(f"Critical transcription failure: {e}")
+            raise CriticalTranscriptionError(f"Critical transcription failure: {e}") from e
+        except Exception as e:
+            logger.warning(f"Non-critical transcription error, skipping chunk: {e}")
+            return None
+    
+    def _resample_to_16khz(self, audio_data: np.ndarray, original_sample_rate: int) -> np.ndarray:
+        """
+        Resample audio to 16kHz for optimal Whisper performance.
+        
+        Args:
+            audio_data: Input audio data as numpy array
+            original_sample_rate: Original sample rate of the audio
+            
+        Returns:
+            Resampled audio data at 16kHz
+        """
+        target_sample_rate = 16000
+        
+        if original_sample_rate == target_sample_rate:
+            return audio_data
+        
+        if audio_data.size == 0:
+            return audio_data
+        # Use scipy.signal.resample for high-quality resampling
+        num_samples = int(len(audio_data) * target_sample_rate / original_sample_rate)
+        resampled_audio = signal.resample(audio_data, num_samples)
+        
+        logger.debug(f"Resampled audio from {original_sample_rate}Hz to {target_sample_rate}Hz "
+                    f"({len(audio_data)} -> {len(resampled_audio)} samples)")
+        
+        return resampled_audio.astype(np.float32)
+
+    def _write_wav_file(self, filename: str, audio_data: np.ndarray, sample_rate: int):
+        """Write WAV file manually if scipy is not available."""
+        import struct
+        import wave
+        
+        with wave.open(filename, 'wb') as wav_file:
+            wav_file.setnchannels(1)  # Mono
+            wav_file.setsampwidth(2)  # 16-bit
+            wav_file.setframerate(sample_rate)
+            wav_file.writeframes(audio_data.tobytes())
+    
+    def _should_output_transcription(self, transcription_interval: float) -> bool:
+        """Check if we should output transcription based on queue or timing."""
+        # For real-time streaming, prioritize queue content
+        if len(self.transcription_queue) > 0:
+            # logger.debug(f"Queue has {len(self.transcription_queue)} items - outputting immediately")  # Too frequent
+            return True
+            
+        # Fallback to timing interval only if queue is empty
+        current_time = time.time()
+        
+        effective_interval = transcription_interval
+            
+        return (current_time - self.last_output_time) >= effective_interval
+    
+    def _get_queued_transcription(self) -> Optional[str]:
+        """Get transcription from queue if available."""
+        try:
+            return self.transcription_queue.popleft()
+        except IndexError:
+            return None
+    
+    def execute(self, audio, transcription_interval, accumulation_duration,
+                whisper_model="base", language="auto", enable_vad=True, output_format="json_segments"):
+        """
+        Execute transcription on streaming audio input.
+        
+        Args:
+            audio: Audio input from LoadAudioTensor - AUDIO dictionary with waveform and sample_rate
+            transcription_interval: Minimum seconds between outputs (prevents message flooding) - optimized for real-time
+            accumulation_duration: Audio accumulation duration for optimal Whisper results - reduced for real-time
+            whisper_model: Whisper model size (tiny/base/small/medium/large-v2)
+            language: Language for transcription (auto for auto-detection)
+            enable_vad: Enable voice activity detection to filter silence
+            output_format: Output format (text/json_segments/json_words)
+            
+        Returns:
+            Tuple containing transcribed text (empty string if not ready to output)
+        """
+        try:
+            # Extract audio data and sample rate from AUDIO dictionary
+            if not isinstance(audio, dict) or 'waveform' not in audio or 'sample_rate' not in audio:
+                logger.warning(f"Invalid audio format: expected AUDIO dict with waveform and sample_rate, got {type(audio)}")
+                return ("",)
+            
+            waveform = audio['waveform']
+            sample_rate = audio['sample_rate']
+            
+            # Convert torch tensor to numpy if needed
+            if isinstance(waveform, torch.Tensor):
+                audio_data = waveform.detach().cpu().numpy()
+            else:
+                audio_data = waveform
+            
+            # Handle tensor dimensions (batch, channels, samples) -> samples
+            if audio_data.ndim == 3:
+                # (batch, channels, samples) -> take first batch, average channels
+                audio_data = audio_data[0]  # Remove batch dimension
+                if audio_data.shape[0] > 1:  # Multiple channels
+                    audio_data = audio_data.mean(axis=0)  # Average channels to mono
+                else:
+                    audio_data = audio_data[0]  # Single channel
+            elif audio_data.ndim == 2:
+                # (channels, samples) -> average channels to mono if needed
+                if audio_data.shape[0] > 1:
+                    audio_data = audio_data.mean(axis=0)
+                else:
+                    audio_data = audio_data[0]
+            # If already 1D, use as-is
+            
+            # Validate inputs
+            if audio_data is None or audio_data.size == 0:
+                logger.warning("Received empty audio data, returning empty")
+                return ("",)
+            
+            # Resample to 16kHz for optimal Whisper performance
+            try:
+                audio_data = self._resample_to_16khz(audio_data, sample_rate)
+            except Exception as e:
+                logger.error(f"Audio resampling to 16Khz failed, returning empty: {e}")
+                return ("",)
+            
+            sample_rate = 16000  # Update sample rate after resampling
+            
+            # Ensure audio_data is a numpy array with proper dtype
+            if not isinstance(audio_data, np.ndarray):
+                logger.debug(f"Converting audio_data from {type(audio_data)} to numpy array.")
+                audio_data = np.array(audio_data, dtype=np.float32)
+
+            
+            # Process the audio chunk and accumulate for optimal transcription
+            ready_for_transcription = self._process_audio_chunk(audio_data, sample_rate, accumulation_duration)
+            
+            # Transcribe if we have accumulated enough audio
+            if ready_for_transcription:
+                transcription = self._transcribe_accumulated_audio(whisper_model, language, enable_vad, 16, output_format)
+                if transcription and transcription.strip() and len(transcription.strip()) > 5:
+                    # Add to queue for controlled output timing (only substantial content or warmup sentinels)
+                    self._queue_transcription(transcription)
+                    is_sentinel = "__WARMUP_SENTINEL__" in transcription
+                    if is_sentinel:
+                        logger.debug(f"Queued warmup sentinel ({output_format})")
+                    else:
+                        logger.debug(f"Queued transcription ({output_format}): '{transcription[:50]}...' (length: {len(transcription)})")
+                elif transcription and transcription != "":
+                    logger.debug(f"Skipping minimal transcription: '{transcription}' (too short)")
+                elif transcription == "":
+                    logger.debug(f"Empty transcription result (post-warmup silence)")
+                else:
+                    logger.debug(f"No transcription result returned")
+            
+            # Output transcriptions immediately when available (real-time streaming)
+            if self._should_output_transcription(transcription_interval):
+                output_text = self._get_queued_transcription()
+                if output_text and output_text.strip():
+                    self.last_output_time = time.time()
+                    
+                    # Log differently for sentinel vs regular transcription
+                    if "__WARMUP_SENTINEL__" in output_text:
+                        logger.info(f"Outputting warmup sentinel immediately ({output_format}, {len(output_text)} chars): '{output_text[:100]}...'")
+                    else:
+                        logger.info(f"Outputting transcription immediately ({output_format}, {len(output_text)} chars): '{output_text[:100]}...' (queue: {len(self.transcription_queue)})")
+                    
+                    return (output_text,)
+            
+            # Return empty string if not ready to output (SaveTextTensor will filter empty content)
+            return ("",)
+            
+        except CriticalTranscriptionError as e:
+            logger.error(f"Critical error in transcription node: {e}")
+            raise
+        except Exception as e:
+            logger.warning(f"Recoverable execute error: {e}")
+            return ("",)
+    
+    def _queue_transcription(self, transcription: str):
+        """Safely queue transcription for controlled output."""
+        if len(self.transcription_queue) == self.transcription_queue.maxlen:
+            logger.debug("Transcription queue is full, oldest entry will be replaced.")
+        self.transcription_queue.append(transcription)
+    
+
+

src/feature_bank/__init__.py

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

src/feature_bank/feature_bank.py

@@ -0,0 +1,535 @@
+"""Source code for the Feature Bank of the StreamV2V paper.
+
+References:
+    StreamV2V: Looking Backward: Streaming Video-to-Video Translation with Feature Banks - 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 extract_nn_features(
+    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.
+        """
+        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, _, _ = (
+            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 extract_nn_features ---")
+                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 = extract_nn_features(
+                    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 = "StreamPack/model_patches/unet"
+
+    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/__init__.py

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

src/media_pipe/face_mesh.py

@@ -0,0 +1,889 @@
+"""Source code for facial landmark detection and utilities using MediaPipe FaceMesh. 
+Inspired by the ComfyUI ControlNet Auxiliary MediaPipe Face node.
+
+References:
+    ComfyUI ControlNet Auxiliary MediaPipe Face node - https://github.com/Fannovel16/comfyui_controlnet_aux
+"""
+
+import numpy as np
+import cv2
+import mediapipe as mp
+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 = "StreamPack/preprocessors/controlnet"
+
+    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 = "StreamPack/preprocessors/controlnet"
+
+    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 = "StreamPack/preprocessors/controlnet"
+
+    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)

src/super_resolution/README.md

@@ -0,0 +1,88 @@
+# ComfyUI Super Resolution
+
+A collection of high-performance neural network-based Super Resolution models for ComfyUI.
+
+## Features
+
+- **Multiple SR Models**: Choose from FSRCNN, ESPCN, LapSRN, and EDSR with different quality/speed tradeoffs
+- **Modular Design**: Load models once and reuse them across multiple upscaling operations
+- **CUDA Acceleration**: Automatic GPU acceleration when available
+- **Multiple Scale Factors**: Support for 2x, 3x, and 4x upscaling
+
+## Installation
+
+1. Clone this repository to your ComfyUI custom_nodes directory:
+   ```
+   cd ComfyUI/custom_nodes
+   git clone https://github.com/yourusername/ComfyUI_SuperResolution
+   ```
+
+2. Restart ComfyUI - the necessary models will be automatically downloaded on first use
+
+## Usage
+
+1. Add a **SR Model Loader** node to your workflow
+2. Select your desired model type and scale factor
+3. Connect the model output to a **SR Upscale** node
+4. Connect an image to the upscale node
+5. Run your workflow to get high-quality upscaled images!
+
+## Model Comparison
+
+| Model | Architecture | Features | Best For | Speed | Quality |
+|-------|-------------|----------|----------|-------|---------|
+| **FSRCNN-small** | Lightweight CNN | Fast, minimal memory use | Real-time processing, mobile | ★★★★★ | ★★ |
+| **FSRCNN** | CNN with larger features | Good balance of speed/quality | General purpose | ★★★★ | ★★★ |
+| **ESPCN** | Sub-pixel convolutions | Efficient upscaling at end | Text/line drawings | ★★★★ | ★★★ |
+| **VDSR** | Very deep CNN | Better edge reconstruction | Detailed images with edges | ★★★ | ★★★★ |
+| **LapSRN** | Laplacian pyramid | Progressive upscaling | Sharp edges, details | ★★★ | ★★★★ |
+| **EDSR** | Deep residual network | Most parameters, best quality | Maximum detail | ★★ | ★★★★★ |
+
+## Technical Details
+
+Each model has unique architectural characteristics:
+
+1. **FSRCNN / FSRCNN-small**: Direct mapping from low to high resolution, lightweight with fewer parameters.
+
+2. **ESPCN**: Uses the efficient sub-pixel convolution technique that processes at low resolution and only expands dimensions at the final layer.
+
+3. **VDSR**: Very Deep Super Resolution network with 20 convolutional layers, uses global residual learning.
+
+4. **LapSRN**: Uses a Laplacian pyramid structure to progressively upscale images, which preserves edges better than the others.
+
+5. **EDSR**: Enhanced Deep Super-Resolution, a residual network with significantly more parameters, offering the highest quality but slower processing.
+
+## Performance Expectations (RTX 4090)
+
+| Model | 2x Upscaling | 4x Upscaling | Memory Usage |
+|-------|--------------|--------------|--------------|
+| FSRCNN-small | 200+ FPS | 150+ FPS | < 100MB |
+| FSRCNN | 100+ FPS | 80+ FPS | < 150MB |
+| ESPCN | 90+ FPS | 70+ FPS | < 150MB |
+| LapSRN | 60+ FPS | 40+ FPS | < 200MB |
+| EDSR | 40+ FPS | 25+ FPS | < 500MB |
+
+## When to Use Each Model
+
+- **FSRCNN-small**: When maximum speed is required
+- **FSRCNN**: For a good balance of quality and performance
+- **ESPCN**: For text and line art upscaling
+- **LapSRN**: When you need better edge preservation than FSRCNN
+- **EDSR**: When maximum quality is desired and performance is secondary
+
+## Comparison to Traditional Methods
+
+These neural network-based upscaling methods offer significantly better quality than traditional algorithms like Lanczos, Bicubic, or Bilinear interpolation, while often being just as fast thanks to CUDA acceleration.
+
+## Credits
+
+This nodepack implements models originally created by:
+
+- FSRCNN: Dong et al. (https://github.com/ryanontheinside/FSRCNN_Tensorflow)
+- EDSR: Lim et al. (https://github.com/ryanontheinside/EDSR_Tensorflow)
+- ESPCN: Shi et al. (https://github.com/ryanontheinside/TF-ESPCN)
+- LapSRN: Lai et al. (https://github.com/ryanontheinside/TF-LapSRN)
+
+## License
+
+MIT 

src/super_resolution/__init__.py

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

src/super_resolution/models/__init__.py

@@ -0,0 +1,5 @@
+# Makes models directory a proper Python package 
+
+from .sr_model import SuperResolutionModel
+
+__all__ = ["SuperResolutionModel"] 

src/super_resolution/models/sr_model.py

@@ -0,0 +1,251 @@
+import os
+import cv2
+import torch
+import numpy as np
+
+class SuperResolutionModel:
+    """Unified Super Resolution Model for all algorithm types"""
+    
+    # Configuration registry for all model types
+    MODEL_CONFIGS = {
+        "fsrcnn": {
+            "algorithm": "fsrcnn",
+            "repo_subdir": "FSRCNN_Tensorflow/models",
+            "filename_pattern": "FSRCNN_x{scale}.pb",
+            "supported_scales": [2, 3, 4],
+        },
+        "fsrcnn-small": {
+            "algorithm": "fsrcnn",
+            "repo_subdir": "FSRCNN_Tensorflow/models",
+            "filename_pattern": "FSRCNN-small_x{scale}.pb",
+            "supported_scales": [2, 3, 4],
+        },
+        "edsr": {
+            "algorithm": "edsr",
+            "repo_subdir": "EDSR_Tensorflow/models",
+            "filename_pattern": "EDSR_x{scale}.pb",
+            "supported_scales": [2, 3, 4],
+        },
+        "espcn": {
+            "algorithm": "espcn",
+            "repo_subdir": "TF-ESPCN/export",
+            "filename_pattern": "ESPCN_x{scale}.pb",
+            "supported_scales": [2, 3, 4],
+        },
+        "lapsrn": {
+            "algorithm": "lapsrn",
+            "repo_subdir": "TF-LapSRN/export",
+            "filename_pattern": "LapSRN_x{scale}.pb",
+            "supported_scales": [2, 3, 4],
+        },
+        "vdsr": {
+            "algorithm": "vdsr",
+            "repo_subdir": "TF-VDSR/export",
+            "filename_pattern": "VDSR_x{scale}.pb",
+            "supported_scales": [2, 3, 4],
+        }
+    }
+    
+    @classmethod
+    def is_cuda_available(cls):
+        """Check if CUDA is available in OpenCV"""
+        # Check if CUDA module exists in OpenCV
+        if not hasattr(cv2, "cuda") or cv2.cuda.getCudaEnabledDeviceCount() == 0:
+            return False
+        
+        # Check if DNN module has CUDA support
+        try:
+            # Try to get available backends if the method exists
+            if hasattr(cv2.dnn, "getAvailableBackends"):
+                backends = cv2.dnn.getAvailableBackends()
+                if cv2.dnn.DNN_BACKEND_CUDA not in backends:
+                    return False
+                    
+            # Check for CUDA targets if method exists
+            if hasattr(cv2.dnn, "getAvailableTargets"):
+                targets = cv2.dnn.getAvailableTargets(cv2.dnn.DNN_BACKEND_CUDA)
+                if len(targets) == 0:
+                    return False
+        except Exception:
+            return False
+            
+        return True
+    
+    @classmethod
+    def create(cls, model_type, scale_factor, **kwargs):
+        """Factory method to create a super resolution model by type"""
+        model_type = model_type.lower()
+        
+        if model_type not in cls.MODEL_CONFIGS:
+            raise ValueError(f"Unknown model type: {model_type}")
+            
+        return cls(model_type, scale_factor, **kwargs)
+    
+    def __init__(self, model_name, scale_factor, device="cuda" if torch.cuda.is_available() else "cpu"):
+        self.model_name = model_name.lower()
+        self.scale_factor = int(scale_factor)
+        self.device = device
+        self.model = None
+        self.loaded = False
+        self.using_cuda = False
+        
+        # Get configuration for this model
+        self.config = self.MODEL_CONFIGS.get(self.model_name)
+        if not self.config:
+            raise ValueError(f"No configuration found for model: {model_name}")
+            
+        # Validate scale factor
+        if self.scale_factor not in self.config["supported_scales"]:
+            raise ValueError(f"Scale factor {scale_factor} not supported for {model_name}. "
+                           f"Supported scales: {self.config['supported_scales']}")
+    
+    def load_model(self, model_path=None):
+        """Load the super resolution model"""
+        if model_path is None:
+            # Try to find model in standard locations
+            model_dir = self.model_name.upper()
+            
+            # Search paths from highest to lowest priority
+            search_paths = [
+                os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../downloaded_models", model_dir),
+                os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../models"),
+                os.path.expanduser("~/.cache/comfyui/custom_nodes/ComfyUI_SuperResolution/models"),
+            ]
+            
+            # Add repo-specific path if available
+            if self.config["repo_subdir"]:
+                search_paths.append(
+                    os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../", self.config["repo_subdir"])
+                )
+            
+            for path in search_paths:
+                filename = self.config["filename_pattern"].format(scale=self.scale_factor)
+                candidate = os.path.join(path, filename)
+                if os.path.exists(candidate):
+                    model_path = candidate
+                    break
+                    
+        if not model_path or not os.path.exists(model_path):
+            raise FileNotFoundError(f"Could not find {self.model_name} x{self.scale_factor} model file")
+            
+        # Load model using OpenCV DNN Super Resolution with the correct method
+        try:
+            # Try different methods to create the DNN SuperRes object for cross-platform compatibility
+            # This handles differences between Windows and Linux OpenCV builds
+            try:
+                # Method 1: Windows-style API (most common)
+                self.model = cv2.dnn_superres.DnnSuperResImpl_create()
+            except AttributeError:
+                try:
+                    # Method 2: More generic constructor if available
+                    self.model = cv2.dnn_superres.SuperResolution_create()
+                except AttributeError:
+                    # Method 3: Direct class instantiation if available
+                    self.model = cv2.dnn_superres.DnnSuperResImpl()
+            
+            self.model.readModel(model_path)
+            self.model.setModel(self.config["algorithm"], self.scale_factor)
+            
+            # Check CUDA availability during model loading
+            # Note: We don't raise errors here, as that will be handled during upscale if use_cuda=True
+            system_has_cuda = torch.cuda.is_available()
+            opencv_has_cuda = self.is_cuda_available()
+            
+            if system_has_cuda and opencv_has_cuda:
+                self.model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+                self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+                print(f"Model loaded with CUDA backend for {self.model_name}")
+                self.using_cuda = True
+            else:
+                self.model.setPreferableBackend(cv2.dnn.DNN_BACKEND_DEFAULT)
+                self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
+                print(f"Model loaded with CPU backend for {self.model_name}")
+                self.using_cuda = False
+                
+            self.loaded = True
+        except Exception as e:
+            raise RuntimeError(f"Error with OpenCV DNN Super Resolution: {e}. Check that opencv-contrib-python is installed.")
+            
+        return self
+    
+    def upscale(self, image, use_cuda=True):
+        """Upscale an image using the model"""
+        if not self.loaded:
+            self.load_model()
+            
+        # Check CUDA requirements if use_cuda is enabled
+        if use_cuda:
+            system_has_cuda = torch.cuda.is_available()
+            opencv_has_cuda = self.is_cuda_available()
+            
+            if not system_has_cuda:
+                raise RuntimeError("CUDA acceleration requested but CUDA is not available on your system. "
+                                 "Either install CUDA or disable CUDA acceleration in the node settings.")
+                                 
+            if not opencv_has_cuda:
+                raise RuntimeError("CUDA acceleration requested but your OpenCV build doesn't support CUDA. "
+                                 "Install a CUDA-enabled OpenCV build (via conda: 'conda install -c conda-forge opencv cudatoolkit') "
+                                 "or disable CUDA acceleration in the node settings.")
+        
+        # Convert to cv2 format if needed
+        cv2_image = self.convert_comfy_to_cv2(image)
+        
+        # Process the image
+        upscaled = self.model.upsample(cv2_image)
+        
+        # Convert back to original format
+        return self.convert_cv2_to_comfy(upscaled, image)
+    
+    def convert_comfy_to_cv2(self, tensor):
+        """Convert ComfyUI tensor (BHWC) to cv2 image (HWC)"""
+        # Handle both tensor and numpy array inputs
+        if isinstance(tensor, torch.Tensor):
+            # Ensure tensor is on CPU and convert to numpy
+            image = tensor.detach().cpu().numpy()
+        else:
+            image = tensor
+            
+        # Take first image if batch dimension exists
+        if len(image.shape) == 4:
+            image = image[0]
+            
+        # Convert float tensor [0,1] to uint8 [0,255]
+        if image.dtype == np.float32 or image.dtype == np.float64:
+            image = (image * 255).astype(np.uint8)
+            
+        # Convert RGB to BGR for OpenCV
+        if image.shape[2] == 3:  # RGB
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+            
+        return image
+    
+    def convert_cv2_to_comfy(self, image, original_tensor=None):
+        """Convert cv2 image (HWC) back to ComfyUI tensor format (BHWC)"""
+        # Convert BGR to RGB
+        if image.shape[2] == 3:  # BGR
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            
+        # Convert to float [0,1] if original was float
+        if original_tensor is not None and (
+            isinstance(original_tensor, torch.Tensor) and original_tensor.dtype == torch.float32 or
+            isinstance(original_tensor, np.ndarray) and original_tensor.dtype == np.float32
+        ):
+            image = image.astype(np.float32) / 255.0
+            
+        # Add batch dimension if needed
+        if original_tensor is not None and len(original_tensor.shape) == 4:
+            image = np.expand_dims(image, axis=0)
+            
+        # Convert to torch tensor if original was tensor
+        if isinstance(original_tensor, torch.Tensor):
+            image = torch.from_numpy(image).to(original_tensor.device)
+            
+        return image
+    
+    def release(self):
+        """Release resources"""
+        self.model = None
+        self.loaded = False
+        
+    def __del__(self):
+        self.release() 

src/super_resolution/models/test_opencv_sr.py

@@ -0,0 +1,110 @@
+import cv2
+import sys
+import torch
+
+def test_opencv_version():
+    print(f"OpenCV Version: {cv2.__version__}")
+    print(f"Python Version: {sys.version}")
+    print(f"CUDA available (torch): {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"CUDA version (torch): {torch.version.cuda}")
+
+def test_cuda_availability():
+    print("\nTesting OpenCV CUDA Support:")
+    
+    # Check if CUDA module exists
+    print(f"cv2.cuda exists: {hasattr(cv2, 'cuda')}")
+    if hasattr(cv2, 'cuda'):
+        print(f"CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}")
+    
+    # Check DNN CUDA backend availability
+    print("\nTesting DNN CUDA Backend:")
+    if hasattr(cv2.dnn, "getAvailableBackends"):
+        backends = cv2.dnn.getAvailableBackends()
+        print(f"Available backends: {backends}")
+        print(f"Has CUDA backend: {cv2.dnn.DNN_BACKEND_CUDA in backends}")
+    else:
+        print("cv2.dnn.getAvailableBackends not available")
+        
+    if hasattr(cv2.dnn, "getAvailableTargets"):
+        try:
+            targets = cv2.dnn.getAvailableTargets(cv2.dnn.DNN_BACKEND_CUDA)
+            print(f"Available CUDA targets: {targets}")
+        except Exception as e:
+            print(f"Error getting CUDA targets: {e}")
+    else:
+        print("cv2.dnn.getAvailableTargets not available")
+
+def test_dnn_superres_availability():
+    print("\nTesting DNN SuperRes availability:")
+    print(f"cv2.dnn_superres exists: {hasattr(cv2, 'dnn_superres')}")
+    
+    if hasattr(cv2, 'dnn_superres'):
+        print("\nAvailable attributes in cv2.dnn_superres:")
+        for attr in dir(cv2.dnn_superres):
+            if not attr.startswith('__'):
+                print(f"- {attr}")
+
+def test_create_methods():
+    print("\nTesting creation methods:")
+    
+    methods = [
+        ("DnnSuperResImpl()", lambda: cv2.dnn_superres.DnnSuperResImpl()),
+        ("DnnSuperResImpl_create()", lambda: cv2.dnn_superres.DnnSuperResImpl_create())
+    ]
+    
+    working_methods = []
+    
+    for method_name, method_func in methods:
+        try:
+            model = method_func()
+            print(f"✓ {method_name} - SUCCESS")
+            working_methods.append((method_name, model))
+        except Exception as e:
+            print(f"✗ {method_name} - FAILED: {str(e)}")
+    
+    return working_methods
+
+def test_model_methods(working_methods):
+    if not working_methods:
+        print("\nNo working creation methods found!")
+        return
+        
+    print("\nTesting model methods and CUDA setup on working instances:")
+    for method_name, model in working_methods:
+        print(f"\nTesting {method_name}:")
+        try:
+            print("Basic methods:")
+            print("- hasattr readModel:", hasattr(model, 'readModel'))
+            print("- hasattr setModel:", hasattr(model, 'setModel'))
+            print("- hasattr upsample:", hasattr(model, 'upsample'))
+            
+            print("\nCUDA-related methods:")
+            print("- hasattr setPreferableBackend:", hasattr(model, 'setPreferableBackend'))
+            if hasattr(model, 'setPreferableBackend'):
+                try:
+                    model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+                    print("  ✓ Successfully set CUDA backend")
+                except Exception as e:
+                    print(f"  ✗ Failed to set CUDA backend: {e}")
+                    
+            print("- hasattr setPreferableTarget:", hasattr(model, 'setPreferableTarget'))
+            if hasattr(model, 'setPreferableTarget'):
+                try:
+                    model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+                    print("  ✓ Successfully set CUDA target")
+                except Exception as e:
+                    print(f"  ✗ Failed to set CUDA target: {e}")
+                    
+        except Exception as e:
+            print(f"Error testing methods: {str(e)}")
+
+if __name__ == "__main__":
+    print("OpenCV Super Resolution CUDA Test Script")
+    print("=" * 40)
+    
+    test_opencv_version()
+    test_cuda_availability()
+    test_dnn_superres_availability()
+    working_methods = test_create_methods()
+    test_model_methods(working_methods) 

src/super_resolution/nodes/__init__.py

@@ -0,0 +1 @@
+# Makes nodes directory a proper Python package 

src/super_resolution/nodes/sr_nodes.py

@@ -0,0 +1,95 @@
+class SuperResolutionModelLoader:
+    """Node that loads a super resolution model and outputs a reference to it"""
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_type": (["FSRCNN", "FSRCNN-small", "EDSR", "ESPCN", "LapSRN", "VDSR"], {
+                    "default": "FSRCNN",
+                    "tooltip": "Select super resolution model type:\n"
+                              "• FSRCNN: Good balance of quality and speed\n"
+                              "• FSRCNN-small: Fastest, lower quality option\n"
+                              "• ESPCN: Efficient for text and line art\n"
+                              "• LapSRN: Better edge preservation\n"
+                              "• EDSR: Highest quality, slower processing\n"
+                              "• VDSR: Very Deep SR with sharp edge reconstruction\n"
+                }),
+                "scale_factor": (["2", "3", "4"], {
+                    "default": "2",
+                    "tooltip": "Upscaling factor:\n"
+                              "• 2: Double resolution (e.g., 512x512 → 1024x1024)\n"
+                              "• 3: Triple resolution (e.g., 512x512 → 1536x1536)\n"
+                              "• 4: Quadruple resolution (e.g., 512x512 → 2048x2048)"
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("SR_MODEL",)
+    RETURN_NAMES = ("sr_model",)
+    FUNCTION = "load_model"
+    CATEGORY = "StreamPack/SuperResolution"
+
+    def load_model(self, model_type, scale_factor):
+        """Load the selected super resolution model and return a reference"""
+        
+        # Import unified model class
+        from ..models import SuperResolutionModel
+        
+        # Import model download utilities
+        from ..utils.model_utils import download_model
+        
+        # On-demand download the model file if it doesn't exist
+        scale_key = f"x{scale_factor}"
+        model_path = None
+        
+        # Download the model if needed
+        model_path = download_model(model_type, scale_key)
+        if not model_path:
+            print(f"Warning: Could not download model {model_type} {scale_key}")
+        
+        # Create the model using our factory method
+        model = SuperResolutionModel.create(model_type, int(scale_factor))
+            
+        # Load the model with the downloaded path if available
+        if model_path:
+            model.load_model(model_path)
+        else:
+            # Fall back to default paths
+            model.load_model()
+            
+        return (model,)
+
+
+class SuperResolutionUpscale:
+    """Node that applies a super resolution model to an image"""
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "sr_model": ("SR_MODEL",),
+                "use_cuda": (["True", "False"], {
+                    "default": "True",
+                    "tooltip": "Whether to use CUDA for acceleration (requires CUDA-enabled OpenCV build)"
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("upscaled_image",)
+    FUNCTION = "upscale_image"
+    OUTPUT_NODE = True
+    CATEGORY = "StreamPack/SuperResolution"
+
+    def upscale_image(self, image, sr_model, use_cuda):
+        """Upscale image using the provided super resolution model"""
+        
+        # Convert use_cuda string to boolean
+        use_cuda = use_cuda == "True"
+        
+        # Upscale the image with the use_cuda parameter
+        result = sr_model.upscale(image, use_cuda=use_cuda)
+        
+        return (result,) 

src/super_resolution/utils/__init__.py

@@ -0,0 +1,2 @@
+# Utils module for ComfyUI SuperResolution
+from .model_utils import download_model, get_model_path 

src/super_resolution/utils/model_utils.py

@@ -0,0 +1,149 @@
+import os
+import sys
+import requests
+import json
+from pathlib import Path
+import itertools
+import shutil
+import time
+
+# Fix the import path to use relative import
+from ..models.sr_model import SuperResolutionModel
+
+def get_models_root_dir():
+    """Get the root directory for downloaded SR models"""
+    # Find the ComfyUI_SuperResolution directory
+    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    # Create downloaded_models directory separate from code
+    return os.path.join(base_dir, "downloaded_models")
+
+def get_model_path(model_family, scale):
+    """Get the path to a model file, ensuring directories exist"""
+    # Normalize model family name to lowercase for consistency
+    model_family = model_family.lower()
+    
+    # Check if model family exists in our configurations
+    if model_family not in SuperResolutionModel.MODEL_CONFIGS:
+        print(f"Error: Unknown model family '{model_family}'")
+        return None
+    
+    # Get configuration for this model
+    config = SuperResolutionModel.MODEL_CONFIGS[model_family]
+    
+    # Convert scale to integer for comparison
+    scale_int = int(scale.replace('x', ''))
+    
+    # Check if scale is supported for this model
+    if scale_int not in config["supported_scales"]:
+        print(f"Error: Scale {scale} not supported for {model_family}")
+        return None
+    
+    # Get model filename from pattern
+    model_filename = config["filename_pattern"].format(scale=scale_int)
+    
+    # Create family-specific subdirectory
+    models_dir = os.path.join(get_models_root_dir(), model_family.upper())
+    os.makedirs(models_dir, exist_ok=True)
+    
+    # Return the full path to the model file
+    return os.path.join(models_dir, model_filename)
+
+def download_model_file(url, destination):
+    """Download a model file with progress indication"""
+    # Skip if model already exists
+    if os.path.exists(destination):
+        print(f"Model already exists: {os.path.basename(destination)}")
+        return True
+    
+    print(f"Downloading {os.path.basename(destination)}...")
+    try:
+        # Download with progress indication for large files
+        with requests.get(url, stream=True) as r:
+            r.raise_for_status()
+            total_size = int(r.headers.get('content-length', 0))
+            with open(destination, 'wb') as f:
+                if total_size > 1024*1024:  # Show progress for files > 1MB
+                    downloaded = 0
+                    start_time = time.time()
+                    for chunk in r.iter_content(chunk_size=8192):
+                        if chunk:
+                            f.write(chunk)
+                            downloaded += len(chunk)
+                            # Show progress every 1MB
+                            if downloaded % (1024*1024) < 8192:
+                                percent = 100 * downloaded / total_size
+                                elapsed = time.time() - start_time
+                                speed = downloaded / (1024 * 1024 * elapsed) if elapsed > 0 else 0
+                                print(f"  {percent:.1f}% ({downloaded/(1024*1024):.1f}MB) - {speed:.1f}MB/s")
+                else:
+                    # For small files, just download without progress
+                    shutil.copyfileobj(r.raw, f)
+        print(f"Downloaded {os.path.basename(destination)} successfully!")
+        return True
+    except Exception as e:
+        print(f"Error downloading {url}: {e}")
+        # Remove partial file if download failed
+        if os.path.exists(destination):
+            os.remove(destination)
+        return False
+
+def download_model(model_family, scale):
+    """Download a specific SR model on-demand if it doesn't exist"""
+    # Normalize model family name to lowercase for consistency
+    model_family_norm = model_family.lower()
+    
+    # Check if valid model family
+    if model_family_norm not in SuperResolutionModel.MODEL_CONFIGS:
+        print(f"Error: Unknown model family '{model_family}'")
+        return None
+    
+    config = SuperResolutionModel.MODEL_CONFIGS[model_family_norm]
+    
+    # Build model download URL map if we don't have a static one
+    model_urls = {
+        "fsrcnn": {
+            "base_url": "https://raw.githubusercontent.com/ryanontheinside/FSRCNN_Tensorflow/master/models/",
+        },
+        "fsrcnn-small": {
+            "base_url": "https://raw.githubusercontent.com/ryanontheinside/FSRCNN_Tensorflow/master/models/",
+        },
+        "edsr": {
+            "base_url": "https://raw.githubusercontent.com/ryanontheinside/EDSR_Tensorflow/master/models/",
+        },
+        "espcn": {
+            "base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-ESPCN/master/export/",
+        },
+        "lapsrn": {
+            "base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-LapSRN/master/export/",
+        },
+        "vdsr": {
+            "base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-VDSR/master/export/",
+        }
+    }
+    
+    # Check if we have a download URL for this model
+    if model_family_norm not in model_urls:
+        print(f"Error: No download URL for model family '{model_family}'")
+        return None
+    
+    # Get the model path
+    model_path = get_model_path(model_family, scale)
+    if not model_path:
+        return None
+    
+    # If model already exists, return its path
+    if os.path.exists(model_path):
+        return model_path
+    
+    # Extract the scale number from the format
+    scale_int = int(scale.replace('x', ''))
+    
+    # Get the base URL and create the download URL
+    base_url = model_urls[model_family_norm]["base_url"]
+    filename = config["filename_pattern"].format(scale=scale_int)
+    url = base_url + filename
+    
+    if download_model_file(url, model_path):
+        return model_path
+    else:
+        return None