project.toml: add publisherId, Icon for registry. use dynamic dependencies

* resample to 16Khz in AudioTranscriptionNode, fix missing kwargs on execute
* update to float32 for consistency
* fail fast on resampling errors

.gitignore

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

__init__.py

@@ -1,19 +1,44 @@
-"""ComfyUI Stream Pack Nodes."""
+"""Expose all nodes in the Stream Pack."""
 
-from .node_wrappers import feature_bank
+import logging
-from .node_wrappers import facemesh
+import pathlib
-from .node_wrappers import batched_denoising
+import importlib
-# Collect all NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS from submodules.
+import traceback
-NODE_CLASS_MAPPINGS = {}
-NODE_DISPLAY_NAME_MAPPINGS = {}
 
-# Import and update mappings from submodules.
+logging.basicConfig(level=logging.INFO)
-for module in [feature_bank, facemesh, batched_denoising]:
+log = logging.getLogger(__name__)
-    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)
 
+
+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']
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

examples/workflows/sd15-tensorrt-api.json

@@ -215,4 +215,4 @@
       "title": "Load CLIP"
     }
   }
-}
+}

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

@@ -231,4 +231,4 @@
       "title": "Feature Bank Attention Processor"
     }
   }
-}
+}

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

@@ -48,6 +48,36 @@
       "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"
@@ -128,7 +158,7 @@
   "14": {
     "inputs": {
       "samples": [
-        "42",
+        "7",
         0
       ],
       "vae": [
@@ -304,73 +334,19 @@
   },
   "41": {
     "inputs": {
-      "num_steps": 3
+      "use_feature_injection": false,
-    },
+      "feature_injection_strength": 0.8,
-    "class_type": "StreamBatchSampler",
+      "feature_similarity_threshold": 0.98,
-    "_meta": {
+      "feature_cache_interval": 4,
-      "title": "Stream Batch Sampler"
+      "feature_bank_max_frames": 4,
-    }
-  },
-  "42": {
-    "inputs": {
-      "add_noise": true,
-      "noise_seed": 977567268563782,
-      "cfg": 1,
-      "model": [
-        "3",
-        0
-      ],
-      "positive": [
-        "38",
-        0
-      ],
-      "negative": [
-        "37",
-        0
-      ],
-      "sampler": [
-        "41",
-        0
-      ],
-      "sigmas": [
-        "44",
-        0
-      ],
-      "latent_image": [
-        "43",
-        0
-      ]
-    },
-    "class_type": "SamplerCustom",
-    "_meta": {
-      "title": "SamplerCustom"
-    }
-  },
-  "43": {
-    "inputs": {
-      "buffer_size": 3,
-      "latent": [
-        "16",
-        0
-      ]
-    },
-    "class_type": "StreamFrameBuffer",
-    "_meta": {
-      "title": "Stream Frame Buffer"
-    }
-  },
-  "44": {
-    "inputs": {
-      "t_index_list": "0,16,32",
-      "num_inference_steps": 50,
       "model": [
         "3",
         0
       ]
     },
-    "class_type": "StreamScheduler",
+    "class_type": "FeatureBankAttentionProcessor",
     "_meta": {
-      "title": "Stream Scheduler"
+      "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/batched_denoising.py

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

node_wrappers/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/facemesh.py

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

node_wrappers/feature_bank.py

@@ -13,7 +13,6 @@ from ..src.feature_bank.feature_bank import (
 NODE_CLASS_MAPPINGS = {
     "FeatureBankAttentionProcessor": FeatureBankAttentionProcessorNode
 }
-
 NODE_DISPLAY_NAME_MAPPINGS = {
     "FeatureBankAttentionProcessor": "Feature Bank Attention Processor"
 }

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,17 @@
+[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"
+
+[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

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

src/__init__.py

@@ -1 +1 @@
-# Dummy file ensuring this package will be recognized.
+# 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/batched_denoising/batched_denoising.py

@@ -1,497 +0,0 @@
-import torch
-import comfy.model_management
-import comfy.samplers
-import math
-
-class StreamBatchSampler:
-    
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "cfg_type": (["none", "full", "self", "initialize"], {"default": "self", "tooltip": "'self' (RCFG) is fastest, 'full' is standard SD but slower, 'initialize' is a middle ground, 'none' disables guidance"}),
-                "delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1, "tooltip": "Only used with 'self' CFG. Controls strength of self-guidance. Higher = stronger guidance but more artifacts. Default 1.0", "round": 0.01}),
-            },
-        }
-    
-    RETURN_TYPES = ("SAMPLER",)
-    FUNCTION = "update"
-    CATEGORY = "StreamPack/sampling"
-    DESCRIPTION = "Implements batched denoising with different CFG options (including RCFG) for faster inference and temporal consistency."
-    
-    def __init__(self):
-        self.frame_buffer = []
-        self.x_t_latent_buffer = None
-        self.stock_noise = None
-        self.is_txt2img_mode = False
-        self.delta = 1.0 
-        self.cfg_type = "self"
-        
-        # Initialize all buffers
-        self.zeros_reference = None
-        self.random_noise_buffer = None
-        self.sigmas_view_buffer = None
-        self.expanded_stock_noise = None
-        self.working_buffer = None
-        self.output_buffer = None
-        
-        # Add RCFG state from reference
-        self.cached_sigmas = None
-        self.alpha_prod_t = None
-        self.beta_prod_t = None
-        self.alpha_prod_t_sqrt = None
-        self.beta_prod_t_sqrt = None
-        self.init_noise = None # Will be initialized in sample based on input shape
-        self.step_counter = 0
-
-    def compute_alpha_beta(self, sigmas):
-        """Pre-compute alpha and beta terms for the given sigmas (LCM-like using cosine schedule)"""
-        self.cached_sigmas = sigmas.clone()
-        print(f"\n[Alpha/Beta] Input sigmas: {sigmas}")
-
-        num_train_timesteps = 1000  # Standard for many diffusion models
-        timesteps = torch.linspace(
-            1, 0, num_train_timesteps, dtype=torch.float32, device=sigmas.device
-        )
-        alphas_cumprod = torch.cos((timesteps + 0.008) / 1.008 * math.pi / 2).pow(2)
-        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
-
-        # Map sigmas to timesteps
-        scaled_sigmas = (
-            sigmas / sigmas.max() if sigmas.max() > 0 else torch.zeros_like(sigmas)
-        )
-        timesteps_for_sigma = (1 - scaled_sigmas) * (num_train_timesteps - 1)
-        indices = (
-            torch.round(timesteps_for_sigma).long().clamp(0, num_train_timesteps - 1)
-        )
-
-        alpha_prod_t = alphas_cumprod[indices].view(-1, 1, 1, 1)
-        beta_prod_t = 1 - alpha_prod_t
-
-        self.alpha_prod_t = alpha_prod_t
-        self.beta_prod_t = beta_prod_t
-        self.alpha_prod_t_sqrt = torch.sqrt(alpha_prod_t)
-        self.beta_prod_t_sqrt = torch.sqrt(beta_prod_t)
-
-        print(f"[Alpha/Beta] Final stats:")
-        print(
-            f"  alpha_prod_t: min={self.alpha_prod_t.min():.4f}, max={self.alpha_prod_t.max():.4f}, mean={self.alpha_prod_t.mean():.4f}"
-        )
-        print(
-            f"  beta_prod_t: min={self.beta_prod_t.min():.4f}, max={self.beta_prod_t.max():.4f}, mean={self.beta_prod_t.mean():.4f}"
-        )
-        print(
-            f"  alpha_sqrt: min={self.alpha_prod_t_sqrt.min():.4f}, max={self.alpha_prod_t_sqrt.max():.4f}, mean={self.alpha_prod_t_sqrt.mean():.4f}"
-        )
-        print(
-            f"  beta_sqrt: min={self.beta_prod_t_sqrt.min():.4f}, max={self.beta_prod_t_sqrt.max():.4f}, mean={self.beta_prod_t_sqrt.mean():.4f}\n"
-        )
-
-    def scheduler_step_batch(self, denoised_batch, scaled_noise):
-            """Compute scheduler step for the entire batch (LCM-like)"""
-            print("\n[Scheduler] Starting scheduler step (LCM-like):")
-            print(f"  Denoised shape: {denoised_batch.shape}")
-            print(f"  Scaled noise shape: {scaled_noise.shape}")
-            print(
-                f"  Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}"
-            )
-            print(
-                f"  Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}"
-            )
-
-            # Assuming self.alpha_prod_t and self.beta_prod_t are for the current timestep
-            alpha_t_sqrt = self.alpha_prod_t_sqrt[: denoised_batch.shape[0]]
-            beta_t_sqrt = self.beta_prod_t_sqrt[: denoised_batch.shape[0]]
-
-            # Calculate the predicted original sample (x_0) from the noisy sample and model output
-            # This is similar to the inverse of the forward diffusion step
-            pred_original_sample = (
-                denoised_batch - beta_t_sqrt.view(-1, 1, 1, 1) * scaled_noise
-            ) / (alpha_t_sqrt.view(-1, 1, 1, 1) + 1e-7)
-            print(
-                f"  Predicted original sample stats: min={pred_original_sample.min():.4f}, max={pred_original_sample.max():.4f}, mean={pred_original_sample.mean():.4f}"
-            )
-
-            # For the reverse step, we need the alpha and beta for the previous timestep
-            # Since we are operating on a batch of timesteps, we'll use the values corresponding to the first timestep in the batch for updating stock_noise
-            # In StreamDiffusion paper, the update happens based on the result of the last denoising step
-
-            # Get alpha and beta for the next step: (which would be the previous in the denoising sequence)
-            alpha_next_sqrt = torch.cat(
-                [
-                    self.alpha_prod_t_sqrt[1 : denoised_batch.shape[0]],
-                    torch.ones_like(self.alpha_prod_t_sqrt[0:1]),
-                ],
-                dim=0,
-            )
-
-            # Calculate the coefficient for the predicted original sample in the next step
-            sigma_next = (1 - alpha_next_sqrt.pow(2)) / (alpha_next_sqrt + 1e-7)
-            sigma_t = (1 - alpha_t_sqrt.pow(2)) / (alpha_t_sqrt + 1e-7)
-
-            # Guidance scale (set to 1 for unconditional prediction)
-            guidance_scale = 1
-
-            # Calculate delta_x (the change in the noisy sample)
-            delta_x = (sigma_next - sigma_t) * guidance_scale * pred_original_sample
-
-            print(
-                f"  Final delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}\n"
-            )
-
-            return delta_x
-
-    def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None):
-        """Sample with staggered batch denoising steps - Optimized version with multiple CFG types"""
-        extra_args = {} if extra_args is None else extra_args
-        cond_scale = extra_args.get("model_options", {}).get("cfg_scale", 1.0) # Get cond_scale from model options
-        cfg_type = self.cfg_type
-        
-        # Get number of frames in batch and available sigmas
-        batch_size = noise.shape[0]
-        num_sigmas = len(sigmas) - 1  # Subtract 1 because last sigma is the target (0.0)
-        
-        # Precompute alpha/beta if needed
-        self.compute_alpha_beta(sigmas) # Sigmas include the final 0.0, compute uses all
-        
-        # Reuse zeros buffer for txt2img detection
-        if self.zeros_reference is None or self.zeros_reference.device != noise.device or self.zeros_reference.dtype != noise.dtype:
-             self.zeros_reference = torch.zeros(1, device=noise.device, dtype=noise.dtype)
-        
-        # Check if noise tensor is all zeros
-        self.is_txt2img_mode = torch.abs(noise).sum() < 1e-5
-        
-        # Initialize init_noise sequence like reference if needed (once or if shape changes)
-        noise_shape_base = noise.shape[1:] # C, H, W
-        if self.init_noise is None or self.init_noise.shape[0] != num_sigmas or self.init_noise.shape[1:] != noise_shape_base:
-            self.init_noise = torch.randn(
-                (num_sigmas, *noise_shape_base),
-                device=noise.device,
-                dtype=noise.dtype,
-            )
-            # Initialize stock_noise to zeros when init_noise is created (like reference)
-            self.stock_noise = torch.zeros_like(self.init_noise[0]) 
-            self.step_counter = 0 # Reset step counter when noise shape changes
-
-        # Shift noise sequence
-        if self.step_counter > 0:
-            self.init_noise = torch.cat([self.init_noise[1:], self.init_noise[0:1]], dim=0)
-            # Update stock noise based on the *first* element of the *shifted* init_noise
-            self.stock_noise = self.init_noise[0].clone() 
-
-
-        if self.is_txt2img_mode:
-            # If txt2img mode, use the init_noise sequence directly as the base noise for the batch
-            # This ensures consistent noise across steps/frames for txt2img
-            if self.random_noise_buffer is None or self.random_noise_buffer.shape != noise.shape:
-                self.random_noise_buffer = torch.empty_like(noise)
-            
-            # Copy from the init_noise sequence for the current batch size
-            self.random_noise_buffer.copy_(self.init_noise[:batch_size])
-            x = self.random_noise_buffer  # Use pre-allocated buffer
-        else:
-             # If not txt2img, we'll still need to add noise later based on sigmas
-            x = noise # Start with the input noise (img2img)
-        
-        # Verify batch size matches number of timesteps
-        if batch_size != num_sigmas:
-            raise ValueError(f"Batch size ({batch_size}) must match number of sigmas ({num_sigmas})")
-        
-        # Pre-allocate and reuse view buffer for sigmas
-        # Ensure buffer matches the number of steps *in this batch* (num_sigmas)
-        if self.sigmas_view_buffer is None or self.sigmas_view_buffer.shape[0] != num_sigmas:
-            self.sigmas_view_buffer = torch.empty((num_sigmas, 1, 1, 1), 
-                                               device=sigmas.device, 
-                                               dtype=sigmas.dtype)
-        # In-place copy of sigmas view (use sigmas[:-1] as these are the steps)
-        self.sigmas_view_buffer.copy_(sigmas[:-1].view(-1, 1, 1, 1))
-        
-        # Apply noise with pre-allocated buffers - no new memory allocation
-        if not self.is_txt2img_mode: 
-            # For img2img, add noise based on sigmas and init_noise sequence
-            # Expand the relevant part of init_noise for the batch
-            current_batch_init_noise = self.init_noise[:batch_size].to(x.device) # Ensure device match
-            
-            # If x is the same object as noise, use a working buffer
-            if id(x) == id(noise):  
-                if self.working_buffer is None or self.working_buffer.shape != noise.shape:
-                    self.working_buffer = torch.empty_like(noise)
-                x = self.working_buffer
-                torch.add(noise, current_batch_init_noise * self.sigmas_view_buffer, out=x)
-            else: # Otherwise, add noise directly to x
-                torch.add(noise, current_batch_init_noise * self.sigmas_view_buffer, out=x)
-
-        # Initialize and manage latent buffer with memory optimization
-        # Only buffer if more than one step
-        if num_sigmas > 1:
-            if self.x_t_latent_buffer is None or self.x_t_latent_buffer.shape != x[0].shape or self.is_txt2img_mode:
-                # Pre-allocate or resize as needed
-                if self.x_t_latent_buffer is None or self.x_t_latent_buffer.shape != x[0].shape:
-                    self.x_t_latent_buffer = torch.empty_like(x[0])
-                self.x_t_latent_buffer.copy_(x[0])
-            
-            # Use buffer for first frame to maintain temporal consistency
-            x[0].copy_(self.x_t_latent_buffer)
-        else: # Single step, no buffering needed
-             self.x_t_latent_buffer = None 
-        
-        with torch.no_grad():
-            # Process all frames in parallel
-            sigma_batch = sigmas[:-1] # Use sigmas for the steps, excluding the final 0.0
-            
-            # --- CFG Handling --- 
-            denoised_batch = None
-            stock_noise_update_value = None # Track value to update stock noise with (used in self-cfg)
-            
-            # 1. No CFG
-            if cfg_type == "none" or cond_scale <= 1.0:
-                denoised_batch = model(x, sigma_batch, **extra_args)
-                stock_noise_update_value = denoised_batch[0].clone().detach()
-
-            # 2. Full CFG
-            elif cfg_type == "full":
-                x_double = torch.cat([x, x], dim=0)
-                sigma_double = torch.cat([sigma_batch, sigma_batch], dim=0)
-                model_output = model(x_double, sigma_double, **extra_args)
-                noise_pred_uncond, noise_pred_text = model_output.chunk(2)
-                denoised_batch = noise_pred_uncond + cond_scale * (noise_pred_text - noise_pred_uncond)
-                # Update stock noise with uncond prediction (using first frame)
-                stock_noise_update_value = noise_pred_uncond[0].clone().detach()
-
-            # 3. Initialize CFG
-            elif cfg_type == "initialize":
-                x_plus_uc = torch.cat([x[0:1], x], dim=0) # Add first frame again for uncond
-                sigma_plus_uc = torch.cat([sigma_batch[0:1], sigma_batch], dim=0)
-                model_output = model(x_plus_uc, sigma_plus_uc, **extra_args)
-                noise_pred_uncond = model_output[0:1] # Uncond is the first item
-                noise_pred_cond = model_output[1:]     # Cond are the rest
-                # Calculate final denoised batch using CFG scale
-                denoised_batch = noise_pred_uncond + cond_scale * (noise_pred_cond - noise_pred_uncond) 
-                # Update stock noise with uncond prediction 
-                stock_noise_update_value = noise_pred_uncond[0].clone().detach()
-            
-            # 4. Self CFG (RCFG)
-            elif cfg_type == "self":
-                eps_c = model(x, sigma_batch, **extra_args) # Conditional prediction
-                
-                # Initialize stock noise if needed (first run after init_noise creation)
-                if self.stock_noise is None: # Should have been initialized with init_noise
-                     # Fallback: initialize with first frame prediction if somehow missed
-                     self.stock_noise = eps_c[0].clone().detach()
-                
-                # RCFG equation
-                # Use current stock_noise (which should be based on init_noise[0])
-                noise_pred_uncond = self.stock_noise.to(eps_c.device) * self.delta
-                denoised_batch = noise_pred_uncond + cond_scale * (eps_c - noise_pred_uncond)
-                
-                # Defer stock noise update until after scheduler step below
-                # Store the conditional prediction for the scheduler step
-                stock_noise_update_value = eps_c.clone().detach() # Need the whole batch for scheduler
-
-
-            # --- CFG Handling End --- 
-
-            # --- RCFG Stock Noise Update ---
-            if cfg_type == "self" and stock_noise_update_value is not None:
-                 # Get the noise sequence for the current batch
-                 current_batch_init_noise = self.init_noise[:batch_size].to(denoised_batch.device)
-                 
-             
-                 beta_sqrt_batch = self.beta_prod_t_sqrt[:batch_size].to(denoised_batch.device)
-                 scaled_noise = beta_sqrt_batch * current_batch_init_noise # Use the init_noise for this step
-
-                 delta_x = self.scheduler_step_batch(denoised_batch, scaled_noise)
-                 num_steps_in_batch = denoised_batch.shape[0]
-                 if self.alpha_prod_t_sqrt.shape[0] >= num_steps_in_batch:
-                     one_tensor = torch.ones_like(self.alpha_prod_t_sqrt[0:1])
-                     zero_tensor = torch.zeros_like(self.beta_prod_t_sqrt[0:1])
-                     alpha_next_sqrt = torch.cat([self.alpha_prod_t_sqrt[1:num_steps_in_batch], one_tensor], dim=0).to(delta_x.device)
-                     beta_next_sqrt = torch.cat([self.beta_prod_t_sqrt[1:num_steps_in_batch], zero_tensor], dim=0).to(delta_x.device)
-                 else:
-                     raise ValueError("Precomputed alpha/beta shapes mismatch for final scaling.")
-                 
-                 # Apply the scaling
-                 scaled_delta_x = (alpha_next_sqrt / (beta_next_sqrt + 1e-7)) * delta_x
-                 
-                 # And the *scaled* delta_x calculated for the *first frame* of the batch
-                 self.stock_noise += scaled_delta_x[0].to(self.stock_noise.device) # Add scaled delta for first frame
-            elif stock_noise_update_value is not None:
-                 # For non-RCFG modes, update stock noise simply with the first frame's value
-                 if self.stock_noise is None or self.stock_noise.shape != stock_noise_update_value.shape:
-                     self.stock_noise = stock_noise_update_value.clone().detach()
-                 else:
-                     self.stock_noise.copy_(stock_noise_update_value.detach())
-
-
-            # Update buffer with intermediate results
-            if num_sigmas > 1:
-                # Store result from first frame as buffer for next iteration
-                self.x_t_latent_buffer.copy_(denoised_batch[0])  # In-place update
-                
-                # Pre-allocate output buffer
-                output_shape = (1, *denoised_batch[-1].shape)
-                if self.output_buffer is None or self.output_buffer.shape != output_shape:
-                    self.output_buffer = torch.empty(output_shape, 
-                                                  device=denoised_batch.device,
-                                                  dtype=denoised_batch.dtype)
-                # Copy the result directly to pre-allocated buffer
-                self.output_buffer[0].copy_(denoised_batch[-1])
-                x_0_pred_out = self.output_buffer
-            else: # Single step
-                x_0_pred_out = denoised_batch.unsqueeze(0) # Ensure batch dim [1,C,H,W]
-                self.x_t_latent_buffer = None # No buffer needed
-            
-            # Call callback if provided
-            if callback is not None:
-                callback({'x': x_0_pred_out, 'i': self.step_counter, 'sigma': sigmas[0], 'sigma_hat': sigmas[0], 'denoised': denoised_batch[-1:]})
-        
-        # Increment step counter at the end of the sample method
-        self.step_counter += 1
-        
-        return x_0_pred_out
-    
-    def update(self, cfg_type="self", delta=1.0):
-        """Create sampler with specified settings"""
-        self.cfg_type = cfg_type
-        self.delta = delta 
-        sampler = comfy.samplers.KSAMPLER(self.sample)
-        return (sampler,)
-
-
-class StreamScheduler:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "t_index_list": ("STRING", {
-                    "default": "0,16,32,49",
-                    "tooltip": "Comma-separated list of timesteps to actually use for denoising. Examples: '32,45' for img2img or '0,16,32,45' for txt2img"
-                }),
-                "num_inference_steps": ("INT", {
-                    "default": 50,
-                    "min": 1,
-                    "max": 1000,
-                    "step": 1,
-                    "tooltip": "Total number of timesteps in schedule. StreamDiffusion uses 50 by default. Only timesteps specified in t_index_list are actually used."
-                }),
-            },
-        }
-
-    RETURN_TYPES = ("SIGMAS",)
-    FUNCTION = "update"
-    CATEGORY = "StreamPack/sampling"
-    DESCRIPTION = "Implements StreamDiffusion's efficient timestep selection. Use in conjunction with StreamBatchSampler."
-    
-    def update(self, model, t_index_list="32,45", num_inference_steps=50):
-        # Get model's sampling parameters
-        model_sampling = model.get_model_object("model_sampling")
-        
-        # Parse timestep list
-        try:
-            t_index_list = [int(t.strip()) for t in t_index_list.split(",")]
-        except ValueError as e:
-            
-            t_index_list = [32, 45]
-            
-        # Create full schedule using normal scheduler
-        full_sigmas = comfy.samplers.normal_scheduler(model_sampling, num_inference_steps)
-        
-        # Select only the sigmas at our desired indices, but in reverse order
-        # This ensures we go from high noise to low noise
-        selected_sigmas = []
-        for t in sorted(t_index_list, reverse=True):  # Sort in reverse to go from high noise to low
-            if t < 0 or t >= num_inference_steps:
-                
-                continue
-            selected_sigmas.append(float(full_sigmas[t]))
-            
-        # Add final sigma
-        selected_sigmas.append(0.0)
-        
-        # Convert to tensor and move to appropriate device
-        selected_sigmas = torch.FloatTensor(selected_sigmas).to(comfy.model_management.get_torch_device())
-        return (selected_sigmas,)
-
-
-class StreamFrameBuffer:
-    
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "latent": ("LATENT",),
-                "buffer_size": ("INT", {
-                    "default": 4,
-                    "min": 1,
-                    "max": 10,
-                    "step": 1,
-                    "tooltip": "Number of frames to buffer before starting batch processing. Should match number of denoising steps."
-                }),
-            },
-        }
-    
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "update"
-    CATEGORY = "StreamPack/sampling"
-    DESCRIPTION = "Accumulates frames to enable staggered batch denoising like StreamDiffusion. Use in conjunction with StreamBatchSampler"
-    
-    
-    def __init__(self):
-        self.frame_buffer = None  # Tensor of shape [buffer_size, C, H, W]
-        self.buffer_size = None
-        self.buffer_pos = 0  # Current position
-        self.is_initialized = False  # Track buffer initialization
-        self.is_txt2img_mode = False
-    
-    def update(self, latent, buffer_size=4):
-        """Add new frame to buffer and return batch when ready"""
-        self.buffer_size = buffer_size
-        
-        # Extract latent tensor from input and remove batch dimension if present
-        x = latent["samples"]
-        
-        # Check if this is a txt2img (zeros tensor) or img2img mode
-        # In ComfyUI, EmptyLatentImage returns a zeros tensor with shape [batch_size, 4, h//8, w//8]
-        # We consider it txt2img mode if the tensor contains all zeros
-        is_txt2img_mode = torch.sum(torch.abs(x)) < 1e-6
-        self.is_txt2img_mode = is_txt2img_mode
-        
-        # If it's a batch with size 1, remove the batch dimension for our buffer
-        if x.dim() == 4 and x.shape[0] == 1:  # [1,C,H,W]
-            x = x.squeeze(0)  # Remove batch dimension -> [C,H,W]
-        
-        # Initialize or resize frame_buffer as a tensor
-        if not self.is_initialized or self.frame_buffer.shape[0] != self.buffer_size or \
-           self.frame_buffer.shape[1:] != x.shape:
-            # Pre-allocate buffer with correct shape
-            self.frame_buffer = torch.zeros(
-                (self.buffer_size, *x.shape),
-                device=x.device,
-                dtype=x.dtype
-            )
-            if self.is_txt2img_mode or not self.is_initialized:
-                # Optimization: Use broadcasting to fill buffer with copies
-                self.frame_buffer[:] = x.unsqueeze(0)  # Broadcast x to [buffer_size, C, H, W]
-
-            self.is_initialized = True
-            self.buffer_pos = 0
-        else:
-            # Add new frame to buffer using ring buffer logic
-            self.frame_buffer[self.buffer_pos] = x  # In-place update
-            self.buffer_pos = (self.buffer_pos + 1) % self.buffer_size  # Circular increment
-        
-        # Optimization: frame_buffer is already a tensor batch, no need to stack
-        batch = self.frame_buffer
-
-        
-        # Return as latent dict with preserved dimensions
-        result = {"samples": batch}
-        
-        # Preserve height and width if present in input
-        if "height" in latent:
-            result["height"] = latent["height"]
-        if "width" in latent:
-            result["width"] = latent["width"]
-            
-        return (result,)

src/feature_bank/__init__.py

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

src/feature_bank/feature_bank.py

@@ -1,7 +1,7 @@
 """Source code for the Feature Bank of the StreamV2V paper.
 
 References:
-    StreamV2V: A Deep Generative Model for High-fidelity View Synthesis - Liang et al.,
+    StreamV2V: Looking Backward: Streaming Video-to-Video Translation with Feature Banks - Liang et al.,
     2025, URL: https://arxiv.org/abs/2405.15757
 """
 
@@ -14,7 +14,7 @@ from torch import nn
 from diffusers.utils import USE_PEFT_BACKEND
 
 
-def get_nn_feats(
+def extract_nn_features(
     x: torch.Tensor, y: torch.Tensor, threshold: float = 0.9
 ) -> torch.Tensor:
     """
@@ -182,7 +182,6 @@ class FeatureBankAttentionProcessor:
         """
         Forward pass of feature bank attention with feature injection.
         """
-        residual = hidden_states
         args = () if USE_PEFT_BACKEND else (scale,)
         input_ndim = hidden_states.ndim
 
@@ -192,7 +191,7 @@ class FeatureBankAttentionProcessor:
                 batch_size, channel, height * width
             ).transpose(1, 2)
 
-        batch_size, key_tokens, _ = (
+        batch_size, _, _ = (
             hidden_states.shape
             if encoder_hidden_states is None
             else encoder_hidden_states.shape
@@ -391,14 +390,14 @@ class FeatureBankAttentionProcessor:
                     "Output",
                 )
 
-                self._debug_print("\n--- DEBUG: Before get_nn_feats ---")
+                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 = get_nn_feats(
+                nn_hidden_states = extract_nn_features(
                     hidden_states,
                     cached_output_tensor_from_list,
                     threshold=self.threshold,
@@ -460,7 +459,7 @@ class FeatureBankAttentionProcessorNode:
 
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "update_node"
-    CATEGORY = "feature-bank"
+    CATEGORY = "StreamPack/model_patches/unet"
 
     def update_node(
         self,

src/media_pipe/__init__.py

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

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