diff --git a/README.MD b/README.MD index bc67783..74c3369 100644 --- a/README.MD +++ b/README.MD @@ -26,7 +26,7 @@ The intention for this repository is to build a suite of nodes that can be used - **QuickShapeMask**: Rapidly generate shape masks (circle, square) with customizable dimensions. - **DTypeConverter**: Convert masks between different data types (float16, uint8, float32, float64). - **FastWebcamCapture**: High-performance webcam capture node with resizing capabilities. - +- **SimilarityFilter**: Filter images based on the similarity of previous frames and block execution if the similarity is too high. ### VAE Tools 🖼️ - **TAESDVaeEncode**: TAESD VAE encoding node. - **TAESDVaeDecode**: TAESD VAE decoding node. diff --git a/__init__.py b/__init__.py index c977445..6d4362c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,6 @@ from .controls.value_controls import FloatControl, IntControl, StringControl from .controls.sequence_controls import FloatSequence, IntSequence, StringSequence -from .controls.utility_controls import FPSMonitor +from .controls.utility_controls import FPSMonitor, SimilarityFilter from .controls.motion_controls import MotionController, ROINode, IntegerMotionController from .quick_shape_mask import QuickShapeMask from .tiny_vae import TAESDVaeEncode, TAESDVaeDecode @@ -14,8 +14,10 @@ NODE_CLASS_MAPPINGS = { "IntSequence": IntSequence, "StringSequence": StringSequence, "FPSMonitor": FPSMonitor, + "SimilarityFilter": SimilarityFilter, "MotionController": MotionController, "IntegerMotionController": IntegerMotionController, + "ROINode": ROINode, #"IntervalControl": IntervalCo ntrol, #"DeltaControl": DeltaControl, @@ -44,8 +46,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "TAESDVaeDecode": "TAESD VAE Decode (RyanOnTheInside)", "DTypeConverter": "DType Converter (RyanOnTheInside)", "FastWebcamCapture": "Fast Webcam Capture (RyanOnTheInside)", + "SimilarityFilter": "Similarity Filter (RyanOnTheInside)", } + WEB_DIRECTORY = "./web/js" __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/controls/utility_controls.py b/controls/utility_controls.py index 927d07d..b7513f7 100644 --- a/controls/utility_controls.py +++ b/controls/utility_controls.py @@ -3,6 +3,7 @@ import time import numpy as np import cv2 import torch +import random class FPSMonitor(ControlNodeBase): """Generates an FPS overlay as an image and mask""" @@ -121,3 +122,77 @@ class FPSMonitor(ControlNodeBase): self.set_state(state) return (state["cached_image"], state["cached_mask"]) + +class SimilarityFilter(ControlNodeBase): + """A node that filters out similar consecutive images to prevent unnecessary workflow execution using StreamDiffusion's approach.""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "always_execute": ("BOOLEAN", {"default": False}), + "threshold": ("FLOAT", {"default": 0.98, "min": 0.0, "max": 1.0, "step": 0.01}), + "max_skip_frames": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "update" + CATEGORY = "real-time/control/utility" + + def __init__(self): + super().__init__() + self.cos = torch.nn.CosineSimilarity(dim=0, eps=1e-6) + + def update(self, image, threshold=0.98, max_skip_frames=10, always_execute=False): + print(f"[DEBUG] Input image object: {hex(id(image))}, shape: {image.shape}, device: {image.device}") + + # Get state with defaults + state = self.get_state({ + "prev_image": None, + "skip_count": 0 + }) + + # First frame case + if state["prev_image"] is None: + state["prev_image"] = image.detach().clone() + state["skip_count"] = 0 + self.set_state(state) + return (image,) + + # Calculate cosine similarity + cos_sim = self.cos( + state["prev_image"].reshape(-1), + image.reshape(-1) + ).item() + + # Calculate skip probability using StreamDiffusion's approach + if threshold >= 1: + skip_prob = 0 + else: + skip_prob = max(0, 1 - (1 - cos_sim) / (1 - threshold)) + + # Generate random sample + sample = random.uniform(0, 1) + + # If we should skip (probability check) + if skip_prob >= sample: + # Check if we've skipped too many frames + if state["skip_count"] >= max_skip_frames: + state["prev_image"] = image.detach().clone() + state["skip_count"] = 0 + self.set_state(state) + return (image,) + + # Skip frame - return ExecutionBlocker to prevent downstream execution + state["skip_count"] += 1 + self.set_state(state) + from comfy_execution.graph import ExecutionBlocker + return (ExecutionBlocker(None),) + + # Frame is different enough - process it + state["prev_image"] = image.detach().clone() + state["skip_count"] = 0 + self.set_state(state) + return (image,) diff --git a/pyproject.toml b/pyproject.toml index 6b7ea0e..dfb6d44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_realtimenodes" description = "These nodes are for real-time applications of ComfyUI." -version = "1.0.0" +version = "1.1.0" license = {file = "LICENSE"} dependencies = ["opencv-python"]