diff --git a/nodes/audio/audio_nodes_effects.py b/nodes/audio/audio_nodes_effects.py index 92b936b..3db0465 100644 --- a/nodes/audio/audio_nodes_effects.py +++ b/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): diff --git a/nodes/audio/audio_nodes_utility.py b/nodes/audio/audio_nodes_utility.py index b93a71a..bfdc6b8 100644 --- a/nodes/audio/audio_nodes_utility.py +++ b/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): diff --git a/nodes/audio/flex_audio.py b/nodes/audio/flex_audio.py index 4d9e91c..67bd174 100644 --- a/nodes/audio/flex_audio.py +++ b/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): diff --git a/nodes/audio/flex_audio_visualizer.py b/nodes/audio/flex_audio_visualizer.py index 5ad3e75..a28e0ea 100644 --- a/nodes/audio/flex_audio_visualizer.py +++ b/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 \ No newline at end of file + return image diff --git a/nodes/flex/flex_base.py b/nodes/flex/flex_base.py index 1e9a6c5..ac302bb 100644 --- a/nodes/flex/flex_base.py +++ b/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",), } } diff --git a/nodes/images/image_base.py b/nodes/images/image_base.py index cdf211a..33469b6 100644 --- a/nodes/images/image_base.py +++ b/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: