progress

nodes/audio/audio_nodes_effects.py

@@ -69,8 +69,6 @@ class AudioGain(AudioEffect):
         amplified_waveform = apply_gain(waveform, gain_db)
         return ({"waveform": amplified_waveform, "sample_rate": sample_rate},)
 
-
-
 class AudioTimeStretch(AudioEffect):
     @classmethod
     def INPUT_TYPES(cls):

nodes/audio/audio_nodes_utility.py

@@ -189,7 +189,6 @@ class AudioSubtract(AudioUtility):
         subtracted_waveform = waveform1 - waveform2
         return ({"waveform": subtracted_waveform, "sample_rate": sample_rate},)
 
-
 class AudioDither(AudioUtility):
     @classmethod
     def INPUT_TYPES(cls):

nodes/audio/flex_audio.py

@@ -150,7 +150,6 @@ class FlexAudioBase(RyanOnTheInside, ABC):
         """Apply the effect to the audio frame. To be implemented by child classes."""
         pass
 
-
 class FlexAudioPitchShift(FlexAudioBase):
     @classmethod
     def INPUT_TYPES(cls):

nodes/audio/flex_audio_visualizer.py

@@ -1,9 +1,11 @@
 import torch
 import numpy as np
 import cv2
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from comfy.utils import ProgressBar
+from ..flex.flex_base import FlexBase
 from ... import RyanOnTheInside
+
 import matplotlib.pyplot as plt
 
 class BaseAudioProcessor:
@@ -88,116 +90,47 @@ class BaseAudioProcessor:
         # Apply smoothing
         self.spectrum = smoothing * self.spectrum + (1 - smoothing) * new_spectrum
 
-class FlexAudioVisualizerBase(RyanOnTheInside, ABC):
+class FlexAudioVisualizerBase(FlexBase, RyanOnTheInside):
     @classmethod
     def INPUT_TYPES(cls):
-        return {
+        base_inputs = super().INPUT_TYPES()
+        base_required = base_inputs.get("required", {})
+        base_optional = base_inputs.get("optional", {})
+
+        new_inputs = {
             "required": {
                 "audio": ("AUDIO",),
                 "frame_rate": ("FLOAT", {"default": 30.0, "min": 1.0, "max": 240.0, "step": 1.0}),
-                "screen_width": ("INT", {"default": 800, "min": 100, "max": 1920, "step": 1}),
-                "screen_height": ("INT", {"default": 600, "min": 100, "max": 1080, "step": 1}),
-                "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "feature_param": (cls.get_modifiable_params(), {"default": cls.get_modifiable_params()[0]}),
-                "feature_mode": (["relative", "absolute"], {"default": "relative"}),
+                "screen_width": ("INT", {"default": 768, "min": 100, "max": 1920, "step": 1}),
+                "screen_height": ("INT", {"default": 464, "min": 100, "max": 1080, "step": 1}),
                 "audio_feature_param": (cls.get_modifiable_params(), {"default": "None"}),
                 "position_x": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                 "position_y": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
             },
             "optional": {
-                "opt_feature": ("FEATURE",)
+                "opt_feature": ("FEATURE",),
+                "opt_feature_pipe": ("FEATURE_PIPE",),
             }
         }
 
+        required = {**base_required, **new_inputs["required"]}
+        optional = {**base_optional, **new_inputs["optional"]}
+
+        return {
+            "required": required,
+            "optional": optional
+        }
+
     CATEGORY = "RyanOnTheInside/FlexAudioVisualizer"
-    RETURN_TYPES = ("IMAGE","MASK")
+    RETURN_TYPES = ("IMAGE", "MASK")
     FUNCTION = "apply_effect"
 
-    def __init__(self):
-        self.progress_bar = None
-
-    def start_progress(self, total_steps, desc="Processing"):
-        self.progress_bar = ProgressBar(total_steps)
-
-    def update_progress(self):
-        if self.progress_bar:
-            self.progress_bar.update(1)
-
-    def end_progress(self):
-        self.progress_bar = None
-
     @classmethod
     @abstractmethod
     def get_modifiable_params(cls):
         """Return a list of parameter names that can be modulated."""
         pass
 
-    def modulate_param(self, param_name, param_value, feature_value, strength, mode):
-        if mode == "relative":
-            modulated_value = param_value * (1 + (feature_value - 0.5) * strength * 2)
-        else:  # absolute
-            modulated_value = param_value * feature_value * strength
-        # Ensure type consistency
-        return type(param_value)(modulated_value)
-
-    def apply_effect(self, audio, frame_rate, screen_width, screen_height, strength, feature_param, feature_mode,
-                     audio_feature_param, opt_feature=None, **kwargs):
-        # Calculate num_frames based on audio duration and frame_rate
-        audio_duration = len(audio['waveform'].squeeze(0).mean(axis=0)) / audio['sample_rate']
-        num_frames = int(audio_duration * frame_rate)
-
-        # Initialize the audio processor
-        processor = BaseAudioProcessor(audio, num_frames, screen_height, screen_width, frame_rate)
-
-        # Initialize results list
-        result = []
-
-        self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}")
-
-        for i in range(num_frames):
-            # Make a copy of kwargs for this frame to avoid altering original parameters
-            frame_kwargs = kwargs.copy()
-
-            # Get audio data for visualization (returns audio_feature_value)
-            audio_feature_value = self.get_audio_data(processor, i, **frame_kwargs)
-
-            # Modulate parameters based on opt_feature if provided
-            if opt_feature is not None and feature_param != "None":
-                # Get feature value from opt_feature
-                feature_value = opt_feature.get_value_at_frame(i)
-
-                if feature_value is not None:
-                    # Modulate the specified parameter using feature_value
-                    original_value = frame_kwargs[feature_param]
-                    modulated_value = self.modulate_param(feature_param, original_value, feature_value, strength, feature_mode)
-                    # Ensure the modulated value is within valid ranges
-                    modulated_value = self.validate_param(feature_param, modulated_value)
-                    frame_kwargs[feature_param] = modulated_value
-
-            # Modulate parameters based on audio feature value if audio_feature_param is set
-            if audio_feature_param != "None" and audio_feature_value is not None:
-                original_value = frame_kwargs[audio_feature_param]
-                modulated_value = self.modulate_param(audio_feature_param, original_value, audio_feature_value, strength, feature_mode)
-                # Ensure the modulated value is within valid ranges
-                modulated_value = self.validate_param(audio_feature_param, modulated_value)
-                frame_kwargs[audio_feature_param] = modulated_value
-
-            # Generate the image for the current frame using frame_kwargs
-            image = self.draw(processor, **frame_kwargs)
-            result.append(image)
-
-            self.update_progress()
-
-        self.end_progress()
-
-        # Convert the list of numpy arrays to a single numpy array
-        result_np = np.stack(result)  # Shape: (N, H, W, 3)
-
-        #lazy
-        result_tensor = torch.from_numpy(result_np).float()
-        mask = result_tensor[:, :, :, 0]
-        return (result_tensor,mask,)
-
     def validate_param(self, param_name, param_value):
         """
         Ensure that modulated parameter values stay within valid ranges.
@@ -255,33 +188,103 @@ class FlexAudioVisualizerBase(RyanOnTheInside, ABC):
         """
         pass
 
+    def apply_effect(self, *args, **kwargs):
+        return self.apply_effect_internal(*args, **kwargs)
+
+    def apply_effect_internal(self, audio, frame_rate, screen_width, screen_height, strength, feature_param, feature_mode,
+                         audio_feature_param, opt_feature=None, **kwargs):
+        # Calculate num_frames based on audio duration and frame_rate
+        audio_duration = len(audio['waveform'].squeeze(0).mean(axis=0)) / audio['sample_rate']
+        num_frames = int(audio_duration * frame_rate)
+
+        # Initialize the audio processor
+        processor = BaseAudioProcessor(audio, num_frames, screen_height, screen_width, frame_rate)
+
+        # Initialize results list
+        result = []
+
+        self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}")
+
+        for i in range(num_frames):
+            # Make a copy of kwargs for this frame to avoid altering original parameters
+            frame_kwargs = kwargs.copy()
+
+            # Get audio data for visualization (returns audio_feature_value)
+            audio_feature_value = self.get_audio_data(processor, i, **frame_kwargs)
+
+            # Modulate parameters based on opt_feature if provided
+            if opt_feature is not None and feature_param != "None":
+                # Get feature value from opt_feature
+                feature_value = opt_feature.get_value_at_frame(i)
+
+                if feature_value is not None:
+                    # Modulate the specified parameter using feature_value
+                    original_value = frame_kwargs[feature_param]
+                    modulated_value = self.modulate_param(feature_param, original_value, feature_value, strength, feature_mode)
+                    # Ensure the modulated value is within valid ranges
+                    modulated_value = self.validate_param(feature_param, modulated_value)
+                    frame_kwargs[feature_param] = modulated_value
+
+            # Modulate parameters based on audio feature value if audio_feature_param is set
+            if audio_feature_param != "None" and audio_feature_value is not None:
+                original_value = frame_kwargs[audio_feature_param]
+                modulated_value = self.modulate_param(audio_feature_param, original_value, audio_feature_value, strength, feature_mode)
+                # Ensure the modulated value is within valid ranges
+                modulated_value = self.validate_param(audio_feature_param, modulated_value)
+                frame_kwargs[audio_feature_param] = modulated_value
+
+            # Generate the image for the current frame using frame_kwargs
+            image = self.draw(processor, **frame_kwargs)
+            result.append(image)
+
+            self.update_progress()
+
+        self.end_progress()
+
+        # Convert the list of numpy arrays to a single numpy array
+        result_np = np.stack(result)  # Shape: (N, H, W, 3)
+
+        # Convert to tensor
+        result_tensor = torch.from_numpy(result_np).float()
+        mask = result_tensor[:, :, :, 0]
+        return (result_tensor, mask)
+
 class FlexAudioVisualizerLine(FlexAudioVisualizerBase):
     @classmethod
     def INPUT_TYPES(cls):
-        return {
-            **super().INPUT_TYPES(),
+        base_inputs = super().INPUT_TYPES()
+        base_required = base_inputs.get("required", {})
+        base_optional = base_inputs.get("optional", {})
+
+        new_inputs = {
             "required": {
-                **super().INPUT_TYPES()["required"],
                 "visualization_method": (["bar", "line"], {"default": "bar"}),
                 "visualization_feature": (["frequency", "waveform"], {"default": "frequency"}),
                 # Parameters common to both methods/features
                 "smoothing": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "rotation": ("FLOAT", {"default":0.0, "min": 0.0, "max": 360.0, "step": 1.0}),
-                # "position_y": ("FLOAT", {"default":0.5, "min":0.0, "max":1.0,"step":0.01}),
+                "rotation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0}),
                 # Additional parameters
-                "num_bars": ("INT", {"default":64, "min":1, "max":1024, "step":1}),
-                "max_height": ("FLOAT", {"default":200.0, "min":10.0, "max":2000.0, "step":10.0}),
-                "min_height": ("FLOAT", {"default":10.0, "min":0.0, "max":500.0, "step":5.0}),
-                "separation": ("FLOAT", {"default":5.0, "min":0.0, "max":100.0, "step":1.0}),
-                "curvature": ("FLOAT", {"default":0.0, "min":0.0, "max":50.0, "step":1.0}),
+                "num_bars": ("INT", {"default": 64, "min": 1, "max": 1024, "step": 1}),
+                "max_height": ("FLOAT", {"default": 200.0, "min": 10.0, "max": 2000.0, "step": 10.0}),
+                "min_height": ("FLOAT", {"default": 10.0, "min": 0.0, "max": 500.0, "step": 5.0}),
+                "separation": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step": 1.0}),
+                "curvature": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 50.0, "step": 1.0}),
                 "reflect": ("BOOLEAN", {"default": False}),
-                "curve_smoothing": ("FLOAT", {"default":0.0, "min": 0.0, "max":1.0, "step": 0.01}),
-                "fft_size": ("INT", {"default":2048, "min":256, "max":8192, "step":256}),
-                "min_frequency": ("FLOAT", {"default":20.0, "min":20.0, "max":20000.0, "step":10.0}),
-                "max_frequency": ("FLOAT", {"default":8000.0, "min":20.0, "max":20000.0, "step":10.0}),
+                "curve_smoothing": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "fft_size": ("INT", {"default": 2048, "min": 256, "max": 8192, "step": 256}),
+                "min_frequency": ("FLOAT", {"default": 20.0, "min": 20.0, "max": 20000.0, "step": 10.0}),
+                "max_frequency": ("FLOAT", {"default": 8000.0, "min": 20.0, "max": 20000.0, "step": 10.0}),
             }
         }
 
+        required = {**base_required, **new_inputs["required"]}
+        optional = base_optional
+
+        return {
+            "required": required,
+            "optional": optional
+        }
+
     FUNCTION = "apply_effect"
 
     @classmethod
@@ -409,11 +412,11 @@ class FlexAudioVisualizerLine(FlexAudioVisualizerBase):
                     mask = np.full((rect_height, rect_width), 0, dtype=np.uint8)
                     cv2.rectangle(mask, (0, 0), (rect_width - 1, rect_height - 1), 255, -1)
                     if radius > 1:
-                        mask = cv2.GaussianBlur(mask, (radius*2+1, radius*2+1), 0)
+                        mask = cv2.GaussianBlur(mask, (radius * 2 + 1, radius * 2 + 1), 0)
                     # Apply mask
                     rect[mask > 0] = color
                     # Place rect onto image
-                    image[y_start:y_end, x:x+rect_width] = rect
+                    image[y_start:y_end, x:x + rect_width] = rect
                 else:
                     cv2.rectangle(image, (x, y_start), (x_end, y_end), color, thickness=-1)
         elif visualization_method == 'line':
@@ -482,27 +485,37 @@ class FlexAudioVisualizerLine(FlexAudioVisualizerBase):
 class FlexAudioVisualizerCircular(FlexAudioVisualizerBase):
     @classmethod
     def INPUT_TYPES(cls):
-        return {
-            **super().INPUT_TYPES(),
+        base_inputs = super().INPUT_TYPES()
+        base_required = base_inputs.get("required", {})
+        base_optional = base_inputs.get("optional", {})
+
+        new_inputs = {
             "required": {
-                **super().INPUT_TYPES()["required"],
                 "visualization_method": (["bar", "line"], {"default": "bar"}),
                 "visualization_feature": (["frequency", "waveform"], {"default": "frequency"}),
                 # Parameters common to both methods/features
-                "smoothing": ("FLOAT", {"default":0.5, "min":0.0, "max":1.0, "step":0.01}),
-                "rotation": ("FLOAT", {"default": 0.0, "min":0.0, "max":360.0, "step":1.0}),
-                "num_points": ("INT", {"default":360, "min":3, "max":1000, "step":1}),
+                "smoothing": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "rotation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0}),
+                "num_points": ("INT", {"default": 360, "min": 3, "max": 1000, "step": 1}),
                 # Additional parameters
-                "fft_size": ("INT", {"default":2048, "min":256, "max":8192, "step":256}),
-                "min_frequency": ("FLOAT", {"default":20.0, "min":20.0, "max":20000.0, "step":10.0}),
-                "max_frequency": ("FLOAT", {"default":8000.0, "min":20.0, "max":20000.0, "step":10.0}),
-                "radius": ("FLOAT", {"default":200.0, "min":10.0, "max":1000.0, "step":10.0}),
-                "line_width": ("INT", {"default":2, "min":1, "max":10, "step":1}),
-                "amplitude_scale": ("FLOAT", {"default":100.0, "min":1.0, "max":1000.0, "step":10.0}),
-                "base_radius": ("FLOAT", {"default":200.0, "min":10.0, "max":1000.0, "step":10.0}),
+                "fft_size": ("INT", {"default": 2048, "min": 256, "max": 8192, "step": 256}),
+                "min_frequency": ("FLOAT", {"default": 20.0, "min": 20.0, "max": 20000.0, "step": 10.0}),
+                "max_frequency": ("FLOAT", {"default": 8000.0, "min": 20.0, "max": 20000.0, "step": 10.0}),
+                "radius": ("FLOAT", {"default": 200.0, "min": 10.0, "max": 1000.0, "step": 10.0}),
+                "line_width": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}),
+                "amplitude_scale": ("FLOAT", {"default": 100.0, "min": 1.0, "max": 1000.0, "step": 10.0}),
+                "base_radius": ("FLOAT", {"default": 200.0, "min": 10.0, "max": 1000.0, "step": 10.0}),
             }
         }
 
+        required = {**base_required, **new_inputs["required"]}
+        optional = base_optional
+
+        return {
+            "required": required,
+            "optional": optional
+        }
+
     FUNCTION = "apply_effect"
 
     @classmethod
@@ -622,4 +635,4 @@ class FlexAudioVisualizerCircular(FlexAudioVisualizerBase):
             if len(points) > 2:
                 cv2.polylines(image, [points], isClosed=True, color=(1.0, 1.0, 1.0), thickness=line_width)
 
-        return image
+        return image

nodes/flex/flex_base.py

@@ -10,12 +10,15 @@ class FlexBase(ABC):
             "required": {
                 "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                 "feature_threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "feature_param": (cls.get_modifiable_params(), {"default": cls.get_modifiable_params()[0] if cls.get_modifiable_params() else "None"}),
+                "feature_param": (
+                    cls.get_modifiable_params(), 
+                    {"default": cls.get_modifiable_params()[0] if cls.get_modifiable_params() else "None"}
+                ),
                 "feature_mode": (["relative", "absolute"], {"default": "relative"}),
             },
             "optional": {
-                "feature": ("FEATURE",),
-                "feature_pipe": ("FEATURE_PIPE",),
+                "opt_feature": ("FEATURE",),
+                "opt_feature_pipe": ("FEATURE_PIPE",),
             }
         }
 

nodes/images/image_base.py

@@ -15,10 +15,7 @@ class FlexImageBase(RyanOnTheInside, FlexBase):
                 **super().INPUT_TYPES()["required"],
                 "images": ("IMAGE",),
             },
-            "optional": {
-                "opt_feature": ("FEATURE",),
-                "opt_feature_pipe": ("FEATURE_PIPE",),
-            }
+            # Optional inputs are inherited from FlexBase
         }
 
     CATEGORY = "RyanOnTheInside/FlexImage"
@@ -50,11 +47,23 @@ class FlexImageBase(RyanOnTheInside, FlexBase):
         else:  # absolute
             return param_value * feature_value * strength
 
-    def apply_effect(self, images, strength, feature_threshold, feature_param, feature_mode, feature=None, feature_pipe=None, **kwargs):
-        if (feature is None) != (feature_pipe is None):
-            raise ValueError("Both feature and feature_pipe must be provided together, or neither should be provided.")
+    def apply_effect(
+        self, 
+        images, 
+        strength, 
+        feature_threshold, 
+        feature_param, 
+        feature_mode, 
+        opt_feature=None, 
+        opt_feature_pipe=None, 
+        **kwargs
+    ):
+        if (opt_feature is None) != (opt_feature_pipe is None):
+            raise ValueError(
+                "Both opt_feature and opt_feature_pipe must be provided together, or neither should be provided."
+            )
 
-        if feature is None and feature_pipe is None:
+        if opt_feature is None and opt_feature_pipe is None:
             # Process all frames without modulation
             num_frames = images.shape[0]
             images_np = images.cpu().numpy()
@@ -75,7 +84,7 @@ class FlexImageBase(RyanOnTheInside, FlexBase):
                 self.update_progress()
         else:
             # Process frames with modulation based on feature values
-            num_frames = feature_pipe.frame_count
+            num_frames = opt_feature_pipe.frame_count
             images_np = images.cpu().numpy()
 
             self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}")
@@ -83,7 +92,7 @@ class FlexImageBase(RyanOnTheInside, FlexBase):
             result = []
             for i in range(num_frames):
                 image = images_np[i]
-                feature_value = feature.get_value_at_frame(i)
+                feature_value = opt_feature.get_value_at_frame(i)
                 kwargs['frame_index'] = i
                 if feature_value >= feature_threshold:
                     processed_image = self.process_image(
@@ -112,8 +121,15 @@ class FlexImageBase(RyanOnTheInside, FlexBase):
 
         return (result_tensor,)
 
-    def process_image(self, image: np.ndarray, feature_value: float, strength: float,
-                      feature_param: str, feature_mode: str, **kwargs) -> np.ndarray:
+    def process_image(
+        self, 
+        image: np.ndarray, 
+        feature_value: float, 
+        strength: float,
+        feature_param: str, 
+        feature_mode: str, 
+        **kwargs
+    ) -> np.ndarray:
         # Modulate the selected parameter
         for param_name in self.get_modifiable_params():
             if param_name in kwargs and param_name == feature_param: