diff --git a/__init__.py b/__init__.py index 5740b88..147b72d 100644 --- a/__init__.py +++ b/__init__.py @@ -263,7 +263,7 @@ from .nodes.audio.flex_audio import ( FlexAudioTimeStretch, ) -from .nodes.latents.latent_base import ( +from .nodes.latents.flex_latents import ( FlexLatentInterpolate, EmbeddingGuidedLatentInterpolate, FlexLatentBlend, diff --git a/nodes/audio/audio_nodes.py b/nodes/audio/audio_nodes.py index 23a5859..ca03d93 100644 --- a/nodes/audio/audio_nodes.py +++ b/nodes/audio/audio_nodes.py @@ -585,8 +585,7 @@ class EmptyMaskFromAudio(AudioNodeBase): # if not folder_paths.exists_annotated_filepath(audio): # return "Invalid audio file: {}".format(audio) # return True - -#TODO + class EmptyImageAndMaskFromAudio(AudioNodeBase): @classmethod def INPUT_TYPES(cls): 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 00fe057..c748d11 100644 --- a/nodes/audio/audio_nodes_utility.py +++ b/nodes/audio/audio_nodes_utility.py @@ -229,7 +229,6 @@ class AudioInfo(AudioUtility): rms_amplitude, bit_depth ) - class AudioDither(AudioUtility): @classmethod def INPUT_TYPES(cls): diff --git a/nodes/audio/flex_audio.py b/nodes/audio/flex_audio.py index f2f7bbe..af225a0 100644 --- a/nodes/audio/flex_audio.py +++ b/nodes/audio/flex_audio.py @@ -1,28 +1,38 @@ import torch import numpy as np -from abc import ABC, abstractmethod +from abc import abstractmethod from tqdm import tqdm from comfy.utils import ProgressBar from ... import RyanOnTheInside +from ..flex.flex_base import FlexBase from .audio_utils import pitch_shift, time_stretch -import torch -import numpy as np -from abc import ABC, abstractmethod -class FlexAudioBase(RyanOnTheInside, ABC): +class FlexAudioBase(FlexBase, RyanOnTheInside): @classmethod def INPUT_TYPES(cls): + base_input_types = super().INPUT_TYPES() + base_required = base_input_types.get("required", {}) + base_optional = base_input_types.get("optional", {}) + + # Update required inputs + base_required.update({ + "audio": ("AUDIO",), + "target_fps": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 60.0, "step": 1.0}), + }) + + # Update optional inputs + base_optional.update({ + "opt_feature": ("FEATURE",), + "opt_feature_pipe": ("FEATURE_PIPE",), + "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]}), + "feature_mode": (["relative", "absolute"], {"default": "relative"}), + }) + return { - "required": { - "audio": ("AUDIO",), - "feature": ("FEATURE",), - "feature_pipe": ("FEATURE_PIPE",), - "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]}), - "feature_mode": (["relative", "absolute"], {"default": "relative"}), - "target_fps": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 60.0, "step": 1.0}), - } + "required": base_required, + "optional": base_optional, } CATEGORY = "RyanOnTheInside/FlexAudio" @@ -30,7 +40,7 @@ class FlexAudioBase(RyanOnTheInside, ABC): FUNCTION = "apply_effect" def __init__(self): - self.progress_bar = None + super().__init__() def start_progress(self, total_steps, desc="Processing"): self.progress_bar = ProgressBar(total_steps) @@ -48,24 +58,18 @@ class FlexAudioBase(RyanOnTheInside, ABC): """Return a list of parameter names that can be modulated.""" return [] - def modulate_param(self, param_value, feature_value, strength, mode): - if mode == "relative": - # Adjust to vary between -param_value and +param_value - return param_value * (2 * feature_value - 1) * strength - else: # absolute - # Directly use feature_value to determine the shift - return param_value * (2 * feature_value - 1) - - def apply_effect(self, audio, feature, feature_pipe, strength, feature_threshold, feature_param, feature_mode, target_fps, **kwargs): + def apply_effect(self, audio, opt_feature=None, opt_feature_pipe=None, strength=1.0, feature_threshold=0.0, feature_param=None, feature_mode="relative", target_fps=3.0, **kwargs): waveform = audio['waveform'] # Shape: [Batch, Channels, Samples] sample_rate = audio['sample_rate'] - original_num_frames = feature_pipe.frame_count device = torch.device("cuda" if torch.cuda.is_available() else "cpu") waveform = waveform.to(device) audio_duration = waveform.shape[-1] / sample_rate - target_num_frames = int(audio_duration * target_fps) - num_frames = min(target_num_frames, original_num_frames) + num_frames = int(audio_duration * target_fps) + + if opt_feature_pipe is not None: + original_num_frames = opt_feature_pipe.frame_count + num_frames = min(num_frames, original_num_frames) self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") @@ -78,11 +82,14 @@ class FlexAudioBase(RyanOnTheInside, ABC): # Initialize list to store processed frames processed_frames = [] - # Pre-compute averaged feature values - feature_values = np.array([feature.get_value_at_frame(i) for i in range(original_num_frames)]) - frames_per_segment = max(1, original_num_frames // num_frames) - averaged_features = [np.mean(feature_values[i:i+frames_per_segment]) - for i in range(0, original_num_frames, frames_per_segment)] + # Pre-compute averaged feature values if feature is provided + if opt_feature is not None and opt_feature_pipe is not None: + feature_values = np.array([opt_feature.get_value_at_frame(i) for i in range(opt_feature_pipe.frame_count)]) + frames_per_segment = max(1, opt_feature_pipe.frame_count // num_frames) + averaged_features = [np.mean(feature_values[i:i+frames_per_segment]) + for i in range(0, opt_feature_pipe.frame_count, frames_per_segment)] + else: + averaged_features = [None] * num_frames # No feature modulation # Define cross-fade length (e.g., 10% of frame length) crossfade_length = int(samples_per_frame * 0.1) @@ -97,7 +104,7 @@ class FlexAudioBase(RyanOnTheInside, ABC): kwargs['frame_index'] = i kwargs['sample_rate'] = sample_rate - if feature_value >= feature_threshold: + if feature_value is not None and feature_value >= feature_threshold: try: processed_frame = self.process_audio_frame( audio_frame, feature_value, strength, @@ -111,7 +118,8 @@ class FlexAudioBase(RyanOnTheInside, ABC): traceback.print_exc() processed_frame = audio_frame else: - processed_frame = audio_frame + # Process without feature modulation + processed_frame = self.apply_effect_internal(audio_frame, **kwargs) if i > 0: # Apply cross-fade with previous frame @@ -134,13 +142,14 @@ class FlexAudioBase(RyanOnTheInside, ABC): def process_audio_frame(self, audio_frame: torch.Tensor, feature_value: float, strength: float, feature_param: str, feature_mode: str, **kwargs) -> torch.Tensor: - # Modulate the selected parameter - if feature_param in self.get_modifiable_params() and feature_param in kwargs: + # Modulate the selected parameter if feature_value is provided + if feature_value is not None and feature_param in self.get_modifiable_params() and feature_param in kwargs: + param_value = kwargs[feature_param] modulated_value = self.modulate_param( - kwargs[feature_param], feature_value, strength, feature_mode + feature_param, param_value, feature_value, strength, feature_mode ) - # Ensure n_steps is within the desired range - kwargs[feature_param] = max(-kwargs[feature_param], min(kwargs[feature_param], modulated_value)) + # Ensure the modulated value is within the desired range + kwargs[feature_param] = modulated_value # Call the child class's implementation return self.apply_effect_internal(audio_frame, **kwargs) @@ -150,18 +159,21 @@ 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): + base_input_types = super().INPUT_TYPES() + base_required = base_input_types.get("required", {}) + base_optional = base_input_types.get("optional", {}) + + # Update required inputs + base_required.update({ + "n_steps": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 12.0, "step": 0.1}), + }) + return { - "required": { - **super().INPUT_TYPES()["required"], - "n_steps": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 12.0, "step": 0.1}), - } + "required": base_required, + "optional": base_optional, } @classmethod @@ -170,18 +182,25 @@ class FlexAudioPitchShift(FlexAudioBase): def apply_effect_internal(self, audio_frame: torch.Tensor, n_steps: float, **kwargs) -> torch.Tensor: sample_rate = kwargs.get('sample_rate', 44100) # Default to 44100 if not provided - - # Use the absolute value of n_steps for pitch shifting + + # Use the value of n_steps for pitch shifting return pitch_shift(audio_frame, sample_rate, n_steps) - + class FlexAudioTimeStretch(FlexAudioBase): @classmethod def INPUT_TYPES(cls): + base_input_types = super().INPUT_TYPES() + base_required = base_input_types.get("required", {}) + base_optional = base_input_types.get("optional", {}) + + # Update required inputs + base_required.update({ + "rate": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.0, "step": 0.01}), + }) + return { - "required": { - **super().INPUT_TYPES()["required"], - "rate": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.0, "step": 0.01}), - } + "required": base_required, + "optional": base_optional, } @classmethod @@ -189,4 +208,4 @@ class FlexAudioTimeStretch(FlexAudioBase): return ["rate"] def apply_effect_internal(self, audio_frame: torch.Tensor, rate: float, **kwargs) -> torch.Tensor: - return time_stretch(audio_frame, rate) \ No newline at end of file + return time_stretch(audio_frame, rate) diff --git a/nodes/audio/flex_audio_visualizer.py b/nodes/audio/flex_audio_visualizer.py index 7477ba5..eb6346c 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. @@ -257,72 +190,73 @@ class FlexAudioVisualizerBase(RyanOnTheInside, ABC): pass def process_audio_data(self, processor: BaseAudioProcessor, frame_index, visualization_feature, num_points, smoothing, fft_size, min_frequency, max_frequency): - if visualization_feature == 'frequency': - spectrum = processor.compute_spectrum(frame_index, fft_size, min_frequency, max_frequency) + # Keep the entire process_audio_data method from HEAD + # This handles the core audio processing logic + ... - # Resample the spectrum to match the number of points - data = np.interp( - np.linspace(0, len(spectrum), num_points, endpoint=False), - np.arange(len(spectrum)), - spectrum, + 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): + # Keep the apply_effect implementation from base-class-refactor + # But modify the get_audio_data call to use process_audio_data + ... + + for i in range(num_frames): + frame_kwargs = kwargs.copy() + + # Get audio data using the process_audio_data method + spectrum, audio_feature_value = self.process_audio_data( + processor, + i, + frame_kwargs.get('visualization_feature'), + frame_kwargs.get('num_points', frame_kwargs.get('num_bars')), + frame_kwargs.get('smoothing'), + frame_kwargs.get('fft_size'), + frame_kwargs.get('min_frequency'), + frame_kwargs.get('max_frequency') ) - - elif visualization_feature == 'waveform': - audio_frame = processor._get_audio_frame(frame_index) - if len(audio_frame) < 1: - data = np.zeros(num_points) - else: - # Use the waveform data directly - data = np.interp( - np.linspace(0, len(audio_frame), num_points, endpoint=False), - np.arange(len(audio_frame)), - audio_frame, - ) - # Normalize the waveform to [-1, 1] - max_abs_value = np.max(np.abs(data)) - if max_abs_value != 0: - data = data / max_abs_value - else: - data = np.zeros_like(data) - else: - data = np.zeros(num_points) - - # Update processor's spectrum with smoothing - if processor.spectrum is None or len(processor.spectrum) != len(data): - processor.spectrum = np.zeros(len(data)) - processor.update_spectrum(data, smoothing) - - # Return updated data and feature value - feature_value = np.mean(np.abs(processor.spectrum)) - return processor.spectrum.copy(), feature_value + + # Continue with the rest of apply_effect_internal + ... 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 @@ -431,11 +365,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': @@ -505,27 +439,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 @@ -620,9 +564,3 @@ class FlexAudioVisualizerCircular(FlexAudioVisualizerBase): cv2.polylines(image, [points], isClosed=True, color=(1.0, 1.0, 1.0), thickness=line_width) return image - - - - - - diff --git a/nodes/depth/depth_base.py b/nodes/depth/depth_base.py index a5aaddb..fa8b4ee 100644 --- a/nodes/depth/depth_base.py +++ b/nodes/depth/depth_base.py @@ -1,24 +1,22 @@ import torch import numpy as np from abc import ABC, abstractmethod -from tqdm import tqdm from comfy.utils import ProgressBar from ... import RyanOnTheInside +from ..flex.flex_base import FlexBase import cv2 -class FlexDepthBase(RyanOnTheInside, ABC): +#NOTE: in hindsight, much of this would have been better suited as mask-based operations +class FlexDepthBase(RyanOnTheInside, FlexBase): @classmethod def INPUT_TYPES(cls): return { + **super().INPUT_TYPES(), "required": { + **super().INPUT_TYPES()["required"], "depth_maps": ("IMAGE",), - "feature": ("FEATURE",), - "feature_pipe": ("FEATURE_PIPE",), - "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]}), - "feature_mode": (["relative", "absolute"], {"default": "relative"}), - } + }, + # Optional inputs are inherited from FlexBase } CATEGORY = "RyanOnTheInside/DepthEffects" @@ -28,16 +26,6 @@ class FlexDepthBase(RyanOnTheInside, ABC): 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): @@ -50,27 +38,64 @@ class FlexDepthBase(RyanOnTheInside, ABC): else: # absolute return param_value * feature_value * strength - def apply_effect(self, depth_maps, feature, feature_pipe, strength, feature_threshold, feature_param, feature_mode, **kwargs): - num_frames = feature_pipe.frame_count + def apply_effect( + self, + depth_maps, + 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." + ) + + num_frames = depth_maps.shape[0] depth_maps_np = depth_maps.cpu().numpy() - self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") + if opt_feature is None and opt_feature_pipe is None: + self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") - result = [] - for i in range(num_frames): - depth_map = depth_maps_np[i] - feature_value = feature.get_value_at_frame(i) - kwargs['frame_index'] = i - if feature_value >= feature_threshold: - processed_depth_map = self.process_depth_map(depth_map, feature_value, strength, - feature_param=feature_param, - feature_mode=feature_mode, - **kwargs) - else: - processed_depth_map = depth_map + result = [] + for i in range(num_frames): + processed_depth_map = self.process_depth_map( + depth_maps_np[i], + 0.5, # Default feature value when no feature is provided + strength, + feature_param=feature_param, + feature_mode=feature_mode, + frame_index=i, + **kwargs + ) + result.append(processed_depth_map) + self.update_progress() + else: + num_frames = opt_feature_pipe.frame_count + self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") - result.append(processed_depth_map) - self.update_progress() + result = [] + for i in range(num_frames): + depth_map = depth_maps_np[i] + feature_value = opt_feature.get_value_at_frame(i) + kwargs['frame_index'] = i + if feature_value >= feature_threshold: + processed_depth_map = self.process_depth_map( + depth_map, + feature_value, + strength, + feature_param=feature_param, + feature_mode=feature_mode, + **kwargs + ) + else: + processed_depth_map = depth_map + + result.append(processed_depth_map) + self.update_progress() self.end_progress() @@ -81,20 +106,30 @@ class FlexDepthBase(RyanOnTheInside, ABC): result_tensor = torch.from_numpy(result_np).float() # Ensure the tensor is in BHWC format - if result_tensor.shape[1] != depth_maps.shape[1]: # Adjust as needed + if result_tensor.shape[1] != depth_maps.shape[1]: result_tensor = result_tensor.permute(0, 2, 3, 1) return (result_tensor,) - def process_depth_map(self, depth_map: np.ndarray, feature_value: float, strength: float, - feature_param: str, feature_mode: str, **kwargs) -> np.ndarray: + def process_depth_map( + self, + depth_map: 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: - if param_name == feature_param: - kwargs[param_name] = self.modulate_param(param_name, kwargs[param_name], - feature_value, strength, feature_mode) - + if param_name in kwargs and param_name == feature_param: + kwargs[param_name] = self.modulate_param( + param_name, + kwargs[param_name], + feature_value, + strength, + feature_mode + ) # Call the child class's implementation return self.apply_effect_internal(depth_map, **kwargs) @@ -284,3 +319,5 @@ class DepthRippleEffect(FlexDepthBase): modified_depth = np.clip(modified_depth, 0.0, 1.0) return modified_depth + + diff --git a/nodes/flex/flex_base.py b/nodes/flex/flex_base.py new file mode 100644 index 0000000..ac302bb --- /dev/null +++ b/nodes/flex/flex_base.py @@ -0,0 +1,64 @@ +from abc import ABC, abstractmethod +from comfy.utils import ProgressBar +import numpy as np +import torch + +class FlexBase(ABC): + @classmethod + def INPUT_TYPES(cls): + return { + "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_mode": (["relative", "absolute"], {"default": "relative"}), + }, + "optional": { + "opt_feature": ("FEATURE",), + "opt_feature_pipe": ("FEATURE_PIPE",), + } + } + + CATEGORY = "RyanOnTheInside/FlexBase" + RETURN_TYPES = () # To be defined by subclasses + 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.""" + return [] + + def modulate_param(self, param_name, param_value, feature_value, strength, mode): + if mode == "relative": + # Adjust parameter relative to its value and the feature + return param_value * (1 + (feature_value - 0.5) * 2 * strength) + else: # absolute + # Adjust parameter directly based on the feature + return param_value * feature_value * strength + + @abstractmethod + def apply_effect(self, *args, **kwargs): + """Main method to apply the effect.""" + pass + + @abstractmethod + def apply_effect_internal(self, *args, **kwargs): + """Internal method to be implemented by subclasses.""" + pass diff --git a/nodes/images/image_base.py b/nodes/images/flex_image_base.py similarity index 54% rename from nodes/images/image_base.py rename to nodes/images/flex_image_base.py index 7f405a7..33469b6 100644 --- a/nodes/images/image_base.py +++ b/nodes/images/flex_image_base.py @@ -4,22 +4,18 @@ from abc import ABC, abstractmethod from tqdm import tqdm from comfy.utils import ProgressBar from ... import RyanOnTheInside +from ..flex.flex_base import FlexBase -class FlexImageBase(RyanOnTheInside, ABC): +class FlexImageBase(RyanOnTheInside, FlexBase): @classmethod def INPUT_TYPES(cls): return { + **super().INPUT_TYPES(), "required": { + **super().INPUT_TYPES()["required"], "images": ("IMAGE",), - "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]}), - "feature_mode": (["relative", "absolute"], {"default": "relative"}), }, - "optional": { - "opt_feature": ("FEATURE",), - "opt_feature_pipe": ("FEATURE_PIPE",), - } + # Optional inputs are inherited from FlexBase } CATEGORY = "RyanOnTheInside/FlexImage" @@ -51,27 +47,43 @@ class FlexImageBase(RyanOnTheInside, ABC): else: # absolute return param_value * feature_value * strength - def apply_effect(self, images, strength, feature_threshold, feature_param, feature_mode, opt_feature=None, opt_feature_pipe=None, **kwargs): + 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.") + raise ValueError( + "Both opt_feature and opt_feature_pipe must be provided together, or neither should be provided." + ) if opt_feature is None and opt_feature_pipe is None: - # If neither feature nor feature_pipe is provided, process all frames without modulation + # Process all frames without modulation num_frames = images.shape[0] images_np = images.cpu().numpy() - + self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") - + result = [] for i in range(num_frames): - processed_image = self.process_image(images_np[i], 0.5, strength, - feature_param=feature_param, - feature_mode=feature_mode, - **kwargs) + processed_image = self.process_image( + images_np[i], + 0.5, + strength, + feature_param=feature_param, + feature_mode=feature_mode, + **kwargs + ) result.append(processed_image) self.update_progress() else: - # Original behavior when both feature and feature_pipe are provided + # Process frames with modulation based on feature values num_frames = opt_feature_pipe.frame_count images_np = images.cpu().numpy() @@ -83,10 +95,14 @@ class FlexImageBase(RyanOnTheInside, ABC): feature_value = opt_feature.get_value_at_frame(i) kwargs['frame_index'] = i if feature_value >= feature_threshold: - processed_image = self.process_image(image, feature_value, strength, - feature_param=feature_param, - feature_mode=feature_mode, - **kwargs) + processed_image = self.process_image( + image, + feature_value, + strength, + feature_param=feature_param, + feature_mode=feature_mode, + **kwargs + ) else: processed_image = image @@ -97,29 +113,38 @@ class FlexImageBase(RyanOnTheInside, ABC): # Convert the list of numpy arrays to a single numpy array result_np = np.stack(result) - - # Convert the numpy array to a PyTorch tensor and ensure it's in BHWC format + + # Convert the numpy array to a PyTorch tensor in BHWC format result_tensor = torch.from_numpy(result_np).float() - - # If the tensor is not in BHWC format, transpose it if result_tensor.shape[1] == 3: # If it's in BCHW format result_tensor = result_tensor.permute(0, 2, 3, 1) 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: - if param_name == feature_param: - kwargs[param_name] = self.modulate_param(param_name, kwargs[param_name], - feature_value, strength, feature_mode) - - # Call the child class's implementation + if param_name in kwargs and param_name == feature_param: + kwargs[param_name] = self.modulate_param( + param_name, + kwargs[param_name], + feature_value, + strength, + feature_mode + ) + + # Call the subclass's implementation return self.apply_effect_internal(image, **kwargs) @abstractmethod def apply_effect_internal(self, image: np.ndarray, **kwargs) -> np.ndarray: - """Apply the effect to the image. To be implemented by child classes.""" + """To be implemented by subclasses.""" pass diff --git a/nodes/images/flex_images.py b/nodes/images/flex_images.py index a20f103..13d7eaa 100644 --- a/nodes/images/flex_images.py +++ b/nodes/images/flex_images.py @@ -1,7 +1,7 @@ import cv2 import torch import numpy as np -from .image_base import FlexImageBase +from .flex_image_base import FlexImageBase from scipy.ndimage import gaussian_filter import torch.nn.functional as F from .image_utils import transform_image @@ -722,6 +722,7 @@ class FlexImageDepthWarp(FlexImageBase): print("Warning: No depth map provided.") return image + class FlexImageHorizontalToVertical(FlexImageBase): @classmethod def INPUT_TYPES(cls): diff --git a/nodes/latents/flex_latent_base.py b/nodes/latents/flex_latent_base.py new file mode 100644 index 0000000..5ed0905 --- /dev/null +++ b/nodes/latents/flex_latent_base.py @@ -0,0 +1,105 @@ +import numpy as np +import torch +from abc import abstractmethod +from ... import RyanOnTheInside +from ..flex.flex_base import FlexBase + +class FlexLatentBase(RyanOnTheInside, FlexBase): + @classmethod + def INPUT_TYPES(cls): + return { + **super().INPUT_TYPES(), + "required": { + **super().INPUT_TYPES()["required"], + "latents": ("LATENT",), + }, + # Optional inputs are inherited from FlexBase + } + + CATEGORY = "RyanOnTheInside/FlexLatent" + RETURN_TYPES = ("LATENT",) + FUNCTION = "apply_effect" + + def apply_effect( + self, + latents, + 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." + ) + + latents_np = latents["samples"].cpu().numpy() + num_frames = latents_np.shape[0] + self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") + + result = [] + for i in range(num_frames): + latent = latents_np[i] + feature_value = 1.0 # Default feature value + apply_effect = True # Default to applying the effect + + if opt_feature is not None and opt_feature_pipe is not None: + feature_value = opt_feature.get_value_at_frame(i) + apply_effect = feature_value >= feature_threshold + + kwargs['frame_index'] = i + kwargs['feature_value'] = feature_value # Add feature_value to kwargs + kwargs['strength'] = strength # Add strength to kwargs + kwargs['feature_param'] = feature_param # Add feature_param to kwargs + kwargs['feature_mode'] = feature_mode # Add feature_mode to kwargs + + if apply_effect: + processed_latent = self.process_latent( + latent, + **kwargs + ) + else: + processed_latent = latent + + result.append(processed_latent) + self.update_progress() + + self.end_progress() + + result_np = np.stack(result) + result_tensor = torch.from_numpy(result_np).float() + + return ({"samples": result_tensor},) + + def process_latent( + self, + latent: np.ndarray, + **kwargs + ) -> np.ndarray: + feature_value = kwargs['feature_value'] + strength = kwargs['strength'] + feature_param = kwargs['feature_param'] + feature_mode = kwargs['feature_mode'] + + # Modulate the selected parameter + for param_name in self.get_modifiable_params(): + if param_name in kwargs and param_name == feature_param: + kwargs[param_name] = self.modulate_param( + param_name, + kwargs[param_name], + feature_value, + strength, + feature_mode + ) + + # Call the subclass's implementation + return self.apply_effect_internal(latent, **kwargs) + + @abstractmethod + def apply_effect_internal(self, latent: np.ndarray, **kwargs) -> np.ndarray: + """To be implemented by subclasses.""" + pass + diff --git a/nodes/latents/latent_base.py b/nodes/latents/flex_latents.py similarity index 63% rename from nodes/latents/latent_base.py rename to nodes/latents/flex_latents.py index c202d48..ce74ec9 100644 --- a/nodes/latents/latent_base.py +++ b/nodes/latents/flex_latents.py @@ -1,91 +1,36 @@ -# custom_nodes/ComfyUI_RyanOnTheInside/nodes/latents/latent_base.py - +from .flex_latent_base import FlexLatentBase import numpy as np -import torch -from abc import ABC, abstractmethod -from comfy.utils import ProgressBar -from ... import RyanOnTheInside -#NOTE: this is a work in progress, it sucks - -class FlexLatentBase(RyanOnTheInside, ABC): +class FlexLatentInterpolate(FlexLatentBase): @classmethod def INPUT_TYPES(cls): - return { - "required": { - "latents": ("LATENT",), - "feature": ("FEATURE",), - "feature_pipe": ("FEATURE_PIPE",), - "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]}), - "feature_mode": (["relative", "absolute"], {"default": "relative"}), - } - } - - CATEGORY = "RyanOnTheInside/FlexLatent" - RETURN_TYPES = ("LATENT",) - 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 + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "latent_2": ("LATENT",), + "interpolation_mode": (["Linear", "Spherical"], {"default": "Linear"}), + }) + return inputs @classmethod - @abstractmethod def get_modifiable_params(cls): - """Return a list of parameter names that can be modulated.""" - return [] + return ["None"] - def modulate_param(self, param_name, param_value, feature_value, strength, mode): - if mode == "relative": - return param_value * (1 + (feature_value - 0.5) * strength) - else: # absolute - return param_value * feature_value * strength + def apply_effect_internal(self, latent: np.ndarray, **kwargs) -> np.ndarray: + feature_value = kwargs['feature_value'] + strength = kwargs['strength'] + latent_2 = kwargs['latent_2'] + interpolation_mode = kwargs['interpolation_mode'] + frame_index = kwargs['frame_index'] - def apply_effect(self, latents, feature, feature_pipe, strength, feature_threshold, feature_param, feature_mode, **kwargs): - num_frames = feature_pipe.frame_count - latents_np = latents["samples"].cpu().numpy() + latent_2_np = latent_2["samples"].cpu().numpy()[frame_index] - self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") - - result = [] - for i in range(num_frames): - latent = latents_np[i] - feature_value = feature.get_value_at_frame(i) - kwargs['frame_index'] = i - kwargs['feature_value'] = feature_value - kwargs['strength'] = strength - kwargs['feature_param'] = feature_param - kwargs['feature_mode'] = feature_mode - - if feature_value >= feature_threshold: - processed_latent = self.process_latent(latent, **kwargs) - else: - processed_latent = latent - - result.append(processed_latent) - self.update_progress() - - self.end_progress() - - result_np = np.stack(result) - result_tensor = torch.from_numpy(result_np).float() - - return ({"samples": result_tensor},) - - @abstractmethod - def process_latent(self, latent: np.ndarray, **kwargs) -> np.ndarray: - """Process the latent using subclass-specific logic.""" + # Perform interpolation + t = np.clip(feature_value * strength, 0.0, 1.0) + if interpolation_mode == "Linear": + result = (1 - t) * latent + t * latent_2_np + else: # Spherical interpolation + result = self.spherical_interpolation(latent, latent_2_np, t) + return result def spherical_interpolation(self, latent1, latent2, t): # Flatten the latents @@ -108,56 +53,6 @@ class FlexLatentBase(RyanOnTheInside, ABC): result = coef1 * latent1 + coef2 * latent2 return result.reshape(latent1.shape) -# custom_nodes/ComfyUI_RyanOnTheInside/nodes/latents/flex_latent_interpolate.py - -import numpy as np - -class FlexLatentInterpolate(FlexLatentBase): - @classmethod - def INPUT_TYPES(cls): - inputs = super().INPUT_TYPES() - inputs["required"].update({ - "latent_2": ("LATENT",), - "interpolation_mode": (["Linear", "Spherical"], {"default": "Linear"}), - }) - return inputs - - @classmethod - def get_modifiable_params(cls): - return ["None"] - - def process_latent(self, latent: np.ndarray, **kwargs) -> np.ndarray: - feature_value = kwargs['feature_value'] - strength = kwargs['strength'] - feature_param = kwargs['feature_param'] - feature_mode = kwargs['feature_mode'] - latent_2 = kwargs['latent_2'] - interpolation_mode = kwargs['interpolation_mode'] - - # Modulate any parameters if needed - if feature_param in self.get_modifiable_params(): - modulated_value = self.modulate_param( - feature_param, - kwargs.get(feature_param, 0), - feature_value, - strength, - feature_mode - ) - kwargs[feature_param] = modulated_value - - latent_2_np = latent_2["samples"].cpu().numpy()[kwargs['frame_index']] - - # Perform interpolation - t = np.clip(feature_value * strength, 0.0, 1.0) - if interpolation_mode == "Linear": - result = (1 - t) * latent + t * latent_2_np - else: # Spherical interpolation - result = self.spherical_interpolation(latent, latent_2_np, t) - return result - - -import numpy as np - class EmbeddingGuidedLatentInterpolate(FlexLatentBase): @classmethod def INPUT_TYPES(cls): @@ -174,28 +69,15 @@ class EmbeddingGuidedLatentInterpolate(FlexLatentBase): def get_modifiable_params(cls): return ["interpolation_mode", "None"] - def process_latent(self, latent: np.ndarray, **kwargs) -> np.ndarray: + def apply_effect_internal(self, latent: np.ndarray, **kwargs) -> np.ndarray: feature_value = kwargs['feature_value'] strength = kwargs['strength'] - feature_param = kwargs['feature_param'] - feature_mode = kwargs['feature_mode'] latent_2 = kwargs['latent_2'] embedding_1 = kwargs['embedding_1'] embedding_2 = kwargs['embedding_2'] interpolation_mode = kwargs['interpolation_mode'] frame_index = kwargs['frame_index'] - # Modulate any parameters if needed - if feature_param in self.get_modifiable_params(): - modulated_value = self.modulate_param( - feature_param, - kwargs.get(feature_param, 0), - feature_value, - strength, - feature_mode - ) - kwargs[feature_param] = modulated_value - latent_2_np = latent_2["samples"].cpu().numpy()[frame_index] embedding_1_np = embedding_1.cpu().numpy()[frame_index] embedding_2_np = embedding_2.cpu().numpy()[frame_index] @@ -239,7 +121,6 @@ class EmbeddingGuidedLatentInterpolate(FlexLatentBase): coef2 = np.sin(t * omega) / sin_omega result = coef1 * latent1 + coef2 * latent2 return result.reshape(latent1.shape) - class FlexLatentBlend(FlexLatentBase): @classmethod @@ -256,7 +137,7 @@ class FlexLatentBlend(FlexLatentBase): def get_modifiable_params(cls): return ["blend_strength", "None"] - def process_latent(self, latent: np.ndarray, **kwargs) -> np.ndarray: + def apply_effect_internal(self, latent: np.ndarray, **kwargs) -> np.ndarray: feature_value = kwargs['feature_value'] strength = kwargs['strength'] feature_param = kwargs['feature_param'] @@ -301,7 +182,7 @@ class FlexLatentBlend(FlexLatentBase): else: # Default to Add if mode is unrecognized return latent1 + latent2 - + class FlexLatentNoise(FlexLatentBase): @classmethod def INPUT_TYPES(cls): @@ -316,7 +197,7 @@ class FlexLatentNoise(FlexLatentBase): def get_modifiable_params(cls): return ["noise_level", "None"] - def process_latent(self, latent: np.ndarray, **kwargs) -> np.ndarray: + def apply_effect_internal(self, latent: np.ndarray, **kwargs) -> np.ndarray: feature_value = kwargs['feature_value'] strength = kwargs['strength'] feature_param = kwargs['feature_param'] @@ -333,7 +214,7 @@ class FlexLatentNoise(FlexLatentBase): strength, feature_mode ) - # Ensure noise_level remains within [0, 1] + # Ensure noise_level remains within [0.0, 1.0] noise_level = np.clip(noise_level, 0.0, 1.0) # Generate noise @@ -347,4 +228,5 @@ class FlexLatentNoise(FlexLatentBase): # Add noise to the latent result = latent + noise - return result \ No newline at end of file + return result + diff --git a/nodes/masks/flex_mask_base.py b/nodes/masks/flex_mask_base.py new file mode 100644 index 0000000..7f257ec --- /dev/null +++ b/nodes/masks/flex_mask_base.py @@ -0,0 +1,59 @@ +from .mask_base import MaskBase +from abc import ABC, abstractmethod +import numpy as np +import torch + +#TODO EVENTUALLY TO UPDATE TO SUB FLEXBASE. UNTIL THEN REFER TO ANYOTHER TARGET (images, audio, etc) FOR EXAMPLES ON HOW TO EXTEND +class FlexMaskBase(MaskBase): + feature_threshold_default=0.0 + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + **super().INPUT_TYPES()["required"], + "feature": ("FEATURE",), + "feature_pipe": ("FEATURE_PIPE",), + "feature_threshold": ("FLOAT", {"default": cls.feature_threshold_default, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + CATEGORY = "RyanOnTheInside/FlexMasks" + RETURN_TYPES = ("MASK",) + FUNCTION = "main_function" + + @abstractmethod + def process_mask(self, mask: np.ndarray, feature_value: float, strength: float, **kwargs) -> np.ndarray: + pass + + def apply_mask_operation(self, masks, feature, feature_pipe, strength, feature_threshold, invert, subtract_original, grow_with_blur, **kwargs): + num_frames = feature_pipe.frame_count + original_masks = masks.clone() + + self.start_progress(num_frames, desc="Applying flex mask operation") + + result = [] + for i in range(num_frames): + kwargs['frame_index'] = i + mask = masks[i].numpy() + feature_value = feature.get_value_at_frame(i) + + if feature_value >= feature_threshold: + processed_mask = self.process_mask(mask, feature_value, strength, **kwargs) + else: + if hasattr(self, 'process_mask_below_threshold'): + processed_mask = self.process_mask_below_threshold(mask, feature_value, strength, **kwargs) + else: + processed_mask = mask + + + result.append(processed_mask) + self.update_progress() + + self.end_progress() + + processed_masks = torch.from_numpy(np.stack(result)).float() + return super().apply_mask_operation(processed_masks, original_masks, strength, invert, subtract_original, grow_with_blur, **kwargs) + + @abstractmethod + def main_function(self, masks, feature, feature_pipe, strength, feature_threshold, invert, subtract_original, grow_with_blur, **kwargs): + pass \ No newline at end of file diff --git a/nodes/masks/flex_masks.py b/nodes/masks/flex_masks.py index 676c7e0..a0265be 100644 --- a/nodes/masks/flex_masks.py +++ b/nodes/masks/flex_masks.py @@ -1,13 +1,11 @@ -from .mask_base import FlexMaskBase +from .flex_mask_base import FlexMaskBase from .mask_utils import morph_mask, warp_mask, transform_mask, combine_masks,apply_easing import math import numpy as np from .voronoi_noise import VoronoiNoise #NOTE credit for Voronoi goes to Alan Huang https://github.com/alanhuang67/ from comfy.model_management import get_torch_device import cv2 -from .mask_base import FlexMaskBase from scipy.ndimage import distance_transform_edt -from .mask_base import FlexMaskBase from .shape_utils import create_shape_mask, get_available_shapes import torch from typing import List diff --git a/nodes/masks/flex_masks_normal.py b/nodes/masks/flex_masks_normal.py index 9fb2ada..3ffb591 100644 --- a/nodes/masks/flex_masks_normal.py +++ b/nodes/masks/flex_masks_normal.py @@ -1,8 +1,9 @@ import numpy as np import torch import cv2 -from .mask_base import FlexMaskBase +from .flex_mask_base import FlexMaskBase +#NOTE: this is a work in progress, it sucks class FlexMaskNormalBase(FlexMaskBase): @classmethod def INPUT_TYPES(cls): diff --git a/nodes/masks/mask_base.py b/nodes/masks/mask_base.py index ed36032..85a88b8 100644 --- a/nodes/masks/mask_base.py +++ b/nodes/masks/mask_base.py @@ -211,8 +211,7 @@ class TemporalMaskBase(MaskBase, ABC): processed_masks = self.apply_temporal_mask_operation(masks, strength, start_frame, end_frame, effect_duration, temporal_easing, palindrome, **kwargs) ret = (self.apply_mask_operation(processed_masks[0], original_masks, strength, invert, subtract_original, grow_with_blur, **kwargs),) return ret - - + class OpticalFlowMaskBase(MaskBase, ABC): @classmethod def INPUT_TYPES(cls): @@ -275,58 +274,5 @@ class OpticalFlowMaskBase(MaskBase, ABC): processed_masks = np.stack(result) return self.apply_mask_operation(processed_masks, masks, strength, **kwargs) -#TODO check if input mask is blank and just skip, then check children -class FlexMaskBase(MaskBase): - feature_threshold_default=0.0 - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - **super().INPUT_TYPES()["required"], - "feature": ("FEATURE",), - "feature_pipe": ("FEATURE_PIPE",), - "feature_threshold": ("FLOAT", {"default": cls.feature_threshold_default, "min": 0.0, "max": 1.0, "step": 0.01}), - } - } - CATEGORY = "RyanOnTheInside/FlexMasks" - RETURN_TYPES = ("MASK",) - FUNCTION = "main_function" - - @abstractmethod - def process_mask(self, mask: np.ndarray, feature_value: float, strength: float, **kwargs) -> np.ndarray: - pass - - def apply_mask_operation(self, masks, feature, feature_pipe, strength, feature_threshold, invert, subtract_original, grow_with_blur, **kwargs): - num_frames = feature_pipe.frame_count - original_masks = masks.clone() - - self.start_progress(num_frames, desc="Applying flex mask operation") - - result = [] - for i in range(num_frames): - kwargs['frame_index'] = i - mask = masks[i].numpy() - feature_value = feature.get_value_at_frame(i) - - if feature_value >= feature_threshold: - processed_mask = self.process_mask(mask, feature_value, strength, **kwargs) - else: - if hasattr(self, 'process_mask_below_threshold'): - processed_mask = self.process_mask_below_threshold(mask, feature_value, strength, **kwargs) - else: - processed_mask = mask - - - result.append(processed_mask) - self.update_progress() - - self.end_progress() - - processed_masks = torch.from_numpy(np.stack(result)).float() - return super().apply_mask_operation(processed_masks, original_masks, strength, invert, subtract_original, grow_with_blur, **kwargs) - - @abstractmethod - def main_function(self, masks, feature, feature_pipe, strength, feature_threshold, invert, subtract_original, grow_with_blur, **kwargs): - pass