|
1 |
| -from ..base.control_base import ControlNodeBase |
| 1 | +import os |
| 2 | +import sys |
2 | 3 | import time
|
3 | 4 | import numpy as np
|
4 | 5 | import cv2
|
5 | 6 | import torch
|
| 7 | +import random |
| 8 | + |
| 9 | +# Add package root to Python path |
| 10 | +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| 11 | +from base.control_base import ControlNodeBase |
| 12 | +from controls.similar_image_filter import SimilarImageFilter |
6 | 13 |
|
7 | 14 | class FPSMonitor(ControlNodeBase):
|
8 | 15 | """Generates an FPS overlay as an image and mask"""
|
@@ -121,3 +128,154 @@ def update(self, width: int, height: int, text_color: int, text_size: float, win
|
121 | 128 | self.set_state(state)
|
122 | 129 |
|
123 | 130 | return (state["cached_image"], state["cached_mask"])
|
| 131 | + |
| 132 | +class SimilarityFilter(ControlNodeBase): |
| 133 | + DESCRIPTION = "Filters out similar consecutive images and outputs a signal to control downstream execution." |
| 134 | + |
| 135 | + @classmethod |
| 136 | + def INPUT_TYPES(s): |
| 137 | + return { |
| 138 | + "required": { |
| 139 | + "image": ("IMAGE", { |
| 140 | + "tooltip": "Input image to compare with previous frame" |
| 141 | + }), |
| 142 | + "always_execute": ("BOOLEAN", { |
| 143 | + "default": False, |
| 144 | + }), |
| 145 | + "threshold": ("FLOAT", { |
| 146 | + "default": 0.98, |
| 147 | + "min": 0.0, |
| 148 | + "max": 1.0, |
| 149 | + "step": 0.01, |
| 150 | + "tooltip": "Similarity threshold (0-1). Higher values mean more frames are considered similar" |
| 151 | + }), |
| 152 | + "max_skip_frames": ("INT", { |
| 153 | + "default": 10, |
| 154 | + "min": 1, |
| 155 | + "max": 100, |
| 156 | + "step": 1, |
| 157 | + "tooltip": "Maximum number of consecutive frames to skip before forcing execution" |
| 158 | + }), |
| 159 | + } |
| 160 | + } |
| 161 | + |
| 162 | + RETURN_TYPES = ("IMAGE", "BOOLEAN") |
| 163 | + RETURN_NAMES = ("image", "should_execute") |
| 164 | + FUNCTION = "update" |
| 165 | + CATEGORY = "real-time/control/utility" |
| 166 | + |
| 167 | + def __init__(self): |
| 168 | + super().__init__() |
| 169 | + self._similarity_filter = SimilarImageFilter() |
| 170 | + |
| 171 | + def update(self, image, threshold=0.98, max_skip_frames=10, always_execute=False): |
| 172 | + print(f"[DEBUG] Input image object: {hex(id(image))}, shape: {image.shape}, device: {image.device}") |
| 173 | + |
| 174 | + # Get state with defaults |
| 175 | + state = self.get_state({ |
| 176 | + "prev_image": None, |
| 177 | + "skip_count": 0 |
| 178 | + }) |
| 179 | + |
| 180 | + # First frame case |
| 181 | + if state["prev_image"] is None: |
| 182 | + state["prev_image"] = image.detach().clone() |
| 183 | + state["skip_count"] = 0 |
| 184 | + self.set_state(state) |
| 185 | + return (image, True) # Always execute first frame |
| 186 | + |
| 187 | + # Update filter parameters |
| 188 | + self._similarity_filter.set_threshold(threshold) |
| 189 | + self._similarity_filter.set_max_skip_frame(max_skip_frames) |
| 190 | + |
| 191 | + # Use filter to check similarity |
| 192 | + result = self._similarity_filter(image) |
| 193 | + should_execute = result is not None |
| 194 | + |
| 195 | + # If we should skip (probability check) |
| 196 | + if not should_execute: |
| 197 | + # Check if we've skipped too many frames |
| 198 | + if state["skip_count"] >= max_skip_frames: |
| 199 | + state["prev_image"] = image.detach().clone() |
| 200 | + state["skip_count"] = 0 |
| 201 | + self.set_state(state) |
| 202 | + return (image, True) # Force execution after max skips |
| 203 | + |
| 204 | + # Skip frame - return previous image and False for execution |
| 205 | + state["skip_count"] += 1 |
| 206 | + self.set_state(state) |
| 207 | + return (state["prev_image"], False) |
| 208 | + |
| 209 | + # Frame is different enough - process it |
| 210 | + state["prev_image"] = image.detach().clone() |
| 211 | + state["skip_count"] = 0 |
| 212 | + self.set_state(state) |
| 213 | + return (image, True) |
| 214 | + |
| 215 | +class AlwaysEqualProxy(str): |
| 216 | + #borrowed from https://github.com/theUpsider/ComfyUI-Logic |
| 217 | + def __eq__(self, _): |
| 218 | + return True |
| 219 | + |
| 220 | + def __ne__(self, _): |
| 221 | + return False |
| 222 | + |
| 223 | +class LazyCondition(ControlNodeBase): |
| 224 | + DESCRIPTION = "Uses lazy evaluation to truly skip execution of unused paths. Maintains state of the last value to circumvent feedback loops." |
| 225 | + |
| 226 | + @classmethod |
| 227 | + def INPUT_TYPES(s): |
| 228 | + return { |
| 229 | + "required": { |
| 230 | + "condition": (AlwaysEqualProxy("*"), { |
| 231 | + "tooltip": "When truthy (non-zero, non-empty, non-None), evaluates and returns if_true path. When falsy, returns either fallback or previous state of if_true.", |
| 232 | + "forceInput": True, |
| 233 | + }), |
| 234 | + "if_true": (AlwaysEqualProxy("*"), { |
| 235 | + "lazy": True, |
| 236 | + "tooltip": "The path that should only be evaluated when condition is truthy" |
| 237 | + }), |
| 238 | + "fallback": (AlwaysEqualProxy("*"), { |
| 239 | + "tooltip": "Alternative value to use when condition is falsy or no previous state of if_true" |
| 240 | + }), |
| 241 | + "use_fallback": ("BOOLEAN", { |
| 242 | + "default": False, |
| 243 | + "tooltip": "When False, uses last successful if_true result (if one exists). When True, uses fallback value" |
| 244 | + }), |
| 245 | + } |
| 246 | + } |
| 247 | + |
| 248 | + RETURN_TYPES = (AlwaysEqualProxy("*"),) |
| 249 | + FUNCTION = "update" |
| 250 | + CATEGORY = "real-time/control/utility" |
| 251 | + |
| 252 | + def check_lazy_status(self, condition, if_true, fallback, use_fallback): |
| 253 | + """Only evaluate the if_true path if condition is truthy.""" |
| 254 | + needed = ["fallback"] # Always need the fallback value |
| 255 | + if condition: |
| 256 | + needed.append("if_true") |
| 257 | + return needed |
| 258 | + |
| 259 | + def update(self, condition, if_true, fallback, use_fallback): |
| 260 | + """Route to either if_true output or fallback value, maintaining state of last if_true.""" |
| 261 | + state = self.get_state({ |
| 262 | + "prev_output": None |
| 263 | + }) |
| 264 | + |
| 265 | + if condition: # Let Python handle truthiness |
| 266 | + # Update last state when we run if_true path |
| 267 | + state["prev_output"] = if_true if if_true is not None else fallback |
| 268 | + if hasattr(if_true, 'detach'): |
| 269 | + state["prev_output"] = if_true.detach().clone() |
| 270 | + self.set_state(state) |
| 271 | + return (if_true,) |
| 272 | + else: |
| 273 | + if use_fallback or state["prev_output"] is None: |
| 274 | + return (fallback,) |
| 275 | + else: |
| 276 | + return (state["prev_output"],) |
| 277 | + |
| 278 | + |
| 279 | + |
| 280 | + |
| 281 | + |
0 commit comments