Feat/batch nodes

docker/entrypoint.sh

@@ -104,6 +104,9 @@ if [ "$1" = "--build-engines" ]; then
   # Build Static Engine for Dreamshaper - Landscape (704x384)
   python src/comfystream/scripts/build_trt.py --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors --out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_\$stat-b-1-h-384-w-704_00001_.engine --width 704 --height 384
 
+  # Build Static Engine for Dreamshaper - Square (512x512) - Batch Size 2
+  python src/comfystream/scripts/build_trt.py --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors --out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_\$stat-b-2-h-512-w-512_00001_.engine --width 512 --height 512 --batch-size 2
+
   # Build Dynamic Engine for Dreamshaper
   python src/comfystream/scripts/build_trt.py \
                 --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors \
@@ -115,6 +118,7 @@ if [ "$1" = "--build-engines" ]; then
                 --max-width 448 \
                 --max-height 704
 
+
   # Build Engine for Depth Anything V2 (guarded by custom node)
   if [ -d "$DEPTH_ANYTHING_NODE_DIR" ] && [ -f "$DEPTH_ANYTHING_EXPORT_SCRIPT" ]; then
     if [ ! -f "$DEPTH_ANYTHING_DIR/$DEPTH_ANYTHING_ENGINE" ]; then

nodes/tensor_utils/__init__.py

@@ -1,13 +1,16 @@
 """Tensor utility nodes for ComfyStream"""
 
 from .load_tensor import LoadTensor
+from .performance_nodes import PerformanceTimerNode, StartPerformanceTimerNode
 from .save_tensor import SaveTensor
 from .save_text_tensor import SaveTextTensor
 
 NODE_CLASS_MAPPINGS = {
     "LoadTensor": LoadTensor,
     "SaveTensor": SaveTensor,
     "SaveTextTensor": SaveTextTensor,
+    "PerformanceTimerNode": PerformanceTimerNode,
+    "StartPerformanceTimerNode": StartPerformanceTimerNode,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {}
 

nodes/tensor_utils/load_tensor.py

@@ -1,5 +1,7 @@
 import queue
 
+import torch
+
 from comfystream import tensor_cache
 from comfystream.exceptions import ComfyStreamInputTimeoutError
 
@@ -24,17 +26,33 @@ def INPUT_TYPES(cls):
                         "tooltip": "Timeout in seconds",
                     },
                 ),
+                "batch_size": (
+                    "INT",
+                    {
+                        "default": 1,
+                        "min": 1,
+                        "max": 8,
+                        "step": 1,
+                        "tooltip": "Number of frames to stack into a single batch",
+                    },
+                ),
             }
         }
 
     @classmethod
     def IS_CHANGED(cls, **kwargs):
         return float("nan")
 
-    def execute(self, timeout_seconds: float = 1.0):
-        try:
-            frame = tensor_cache.image_inputs.get(block=True, timeout=timeout_seconds)
+    def execute(self, timeout_seconds: float = 1.0, batch_size: int = 1):
+        frames = []
+        for _ in range(batch_size):
+            try:
+                frame = tensor_cache.image_inputs.get(block=True, timeout=timeout_seconds)
+            except queue.Empty:
+                raise ComfyStreamInputTimeoutError("video", timeout_seconds)
             frame.side_data.skipped = False
-            return (frame.side_data.input,)
-        except queue.Empty:
-            raise ComfyStreamInputTimeoutError("video", timeout_seconds)
+            frames.append(frame.side_data.input)
+
+        if len(frames) == 1:
+            return (frames[0],)
+        return (torch.cat(frames, dim=0),)

nodes/tensor_utils/performance_nodes.py

@@ -0,0 +1,70 @@
+"""
+Performance measurement nodes for ComfyStream batch processing.
+These nodes integrate with the existing tensor_utils structure.
+"""
+
+from comfystream.utils import performance_timer
+
+
+class PerformanceTimerNode:
+    CATEGORY = "tensor_utils"
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("performance_summary",)
+    FUNCTION = "execute"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "operation": ("STRING", {"default": "workflow_execution"}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1}),
+                "num_images": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
+            }
+        }
+
+    @classmethod
+    def IS_CHANGED(s):
+        return float("nan")
+
+    def execute(self, operation: str, batch_size: int, num_images: int):
+        """Record performance metrics and return summary."""
+        performance_timer.record_batch_processing(batch_size, num_images)
+        performance_timer.end_timing(operation)
+
+        summary = performance_timer.get_performance_summary()
+
+        # Format summary as readable string
+        summary_str = "Performance Summary:\n"
+        summary_str += f"Total Images Processed: {summary['total_images_processed']}\n"
+        summary_str += f"Total FPS: {summary['total_fps']:.2f}\n"
+        summary_str += f"Average Batch Size: {summary['average_batch_size']:.2f}\n"
+
+        for key, value in summary.items():
+            if key not in ["total_images_processed", "total_fps", "average_batch_size"]:
+                summary_str += f"{key}: {value:.4f}\n"
+
+        return (summary_str,)
+
+
+class StartPerformanceTimerNode:
+    CATEGORY = "tensor_utils"
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("timer_started",)
+    FUNCTION = "execute"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "operation": ("STRING", {"default": "workflow_execution"}),
+            }
+        }
+
+    @classmethod
+    def IS_CHANGED(s):
+        return float("nan")
+
+    def execute(self, operation: str):
+        """Start timing an operation."""
+        performance_timer.start_timing(operation)
+        return (f"Started timing: {operation}",)

nodes/tensor_utils/save_tensor.py

@@ -14,13 +14,28 @@ def INPUT_TYPES(s):
         return {
             "required": {
                 "images": ("IMAGE",),
+            },
+            "optional": {
+                "split_batch": ("BOOLEAN", {"default": False}),
             }
         }
 
     @classmethod
     def IS_CHANGED(s):
         return float("nan")
 
-    def execute(self, images: torch.Tensor):
-        tensor_cache.image_outputs.put_nowait(images)
+    def execute(self, images: torch.Tensor, split_batch: bool = False):
+        """
+        Save tensor(s) to the tensor cache.
+        If split_batch is True and images is a batch, splits it into individual images.
+        """
+        if split_batch and images.dim() == 4 and images.shape[0] > 1:
+            # Split batch into individual images
+            for i in range(images.shape[0]):
+                single_image = images[i:i+1]  # Keep batch dimension
+                tensor_cache.image_outputs.put_nowait(single_image)
+        else:
+            # Save as single tensor (original behavior)
+            tensor_cache.image_outputs.put_nowait(images)
+
         return images

src/comfystream/utils.py

@@ -1,7 +1,9 @@
 import copy
 import importlib
 import json
-from typing import Any, Dict
+import time
+from contextlib import contextmanager
+from typing import Any, Dict, List
 
 from comfy.api.components.schema.prompt import Prompt, PromptDictInput
 
@@ -99,3 +101,78 @@ def get_default_workflow() -> dict:
         },
         "2": {"inputs": {}, "class_type": "LoadTensor", "_meta": {"title": "LoadTensor"}},
     }
+
+
+class PerformanceTimer:
+    """Utility class for measuring performance metrics in ComfyStream workflows."""
+
+    def __init__(self):
+        self.timings: Dict[str, List[float]] = {}
+        self.current_timings: Dict[str, float] = {}
+        self.batch_sizes: List[int] = []
+        self.total_images_processed = 0
+
+    def start_timing(self, operation: str):
+        """Start timing an operation."""
+        self.current_timings[operation] = time.time()
+
+    def end_timing(self, operation: str) -> float:
+        """End timing an operation and record the duration."""
+        if operation in self.current_timings:
+            duration = time.time() - self.current_timings[operation]
+            self.timings.setdefault(operation, []).append(duration)
+            del self.current_timings[operation]
+            return duration
+        return 0.0
+
+    def record_batch_processing(self, batch_size: int, num_images: int):
+        """Record a batch processing event."""
+        self.batch_sizes.append(batch_size)
+        self.total_images_processed += num_images
+
+    def get_fps(self, operation: str = "total") -> float:
+        """Calculate FPS for a specific operation."""
+        total_time = sum(self.timings.get(operation, []))
+        if total_time == 0:
+            return 0.0
+        return self.total_images_processed / total_time
+
+    def get_average_time(self, operation: str) -> float:
+        """Get average time for an operation."""
+        samples = self.timings.get(operation, [])
+        if not samples:
+            return 0.0
+        return sum(samples) / len(samples)
+
+    def get_performance_summary(self) -> Dict[str, float]:
+        """Get a comprehensive performance summary."""
+        summary = {
+            "total_images_processed": self.total_images_processed,
+            "total_fps": self.get_fps("total"),
+            "average_batch_size": (
+                sum(self.batch_sizes) / len(self.batch_sizes) if self.batch_sizes else 0
+            ),
+        }
+        for operation in self.timings:
+            summary[f"{operation}_fps"] = self.get_fps(operation)
+            summary[f"{operation}_avg_time"] = self.get_average_time(operation)
+        return summary
+
+    def reset(self):
+        """Reset all performance data."""
+        self.timings.clear()
+        self.current_timings.clear()
+        self.batch_sizes.clear()
+        self.total_images_processed = 0
+
+    @contextmanager
+    def time_operation(self, operation: str):
+        """Context manager for timing operations."""
+        self.start_timing(operation)
+        try:
+            yield
+        finally:
+            self.end_timing(operation)
+
+
+performance_timer = PerformanceTimer()