diff --git a/__init__.py b/__init__.py index 565bfe3..3917608 100644 --- a/__init__.py +++ b/__init__.py @@ -409,7 +409,15 @@ NODE_CLASS_MAPPINGS = { #features "AudioFeatureExtractor": AudioFeatureExtractor, + +#TODO add framecount from audio +#TODO time feature from rhythm or framerate or something....tempo? need tempo extractor! frames per measure, frames per quarter note, etc +#TODO increase shift amount chromatic aberation +#TODO make feature info JS display info +#TODO removethis +#TODO duplicate progress bar "AudioFeatureExtractorFirst": AudioFeatureExtractorFirst, + "PitchFeatureExtractor": PitchFeatureExtractor, "RhythmFeatureExtractor": RhythmFeatureExtractor, diff --git a/nodes/flex/flex_base.py b/nodes/flex/flex_base.py index 66459fc..c74be2c 100644 --- a/nodes/flex/flex_base.py +++ b/nodes/flex/flex_base.py @@ -3,12 +3,13 @@ from comfy.utils import ProgressBar import numpy as np import torch from ...tooltips import apply_tooltips +from .parameter_scheduling import ParameterScheduler @apply_tooltips class FlexBase(ABC): @classmethod def INPUT_TYPES(cls): - return { + base_inputs = { "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}), @@ -22,13 +23,37 @@ class FlexBase(ABC): "opt_feature": ("FEATURE",), } } + + # Add list_ok to all float and int inputs + for section in ["required", "optional"]: + if section in base_inputs: + for key, value in base_inputs[section].items(): + if isinstance(value, tuple) and len(value) == 2: + input_type, config = value + if input_type in ["FLOAT", "INT"] and isinstance(config, dict): + config["list_ok"] = True + base_inputs[section][key] = (input_type, config) + + return base_inputs - CATEGORY = "RyanOnTheInside/FlexBase" - RETURN_TYPES = () # To be defined by subclasses - FUNCTION = "apply_effect" + @classmethod + def update_input_types(cls): + """Helper method to add list_ok to all float and int inputs in a class's INPUT_TYPES""" + input_types = cls.INPUT_TYPES() + for section in ["required", "optional"]: + if section in input_types: + for key, value in input_types[section].items(): + if isinstance(value, tuple) and len(value) == 2: + input_type, config = value + if input_type in ["FLOAT", "INT"] and isinstance(config, dict): + config["list_ok"] = True + input_types[section][key] = (input_type, config) + return input_types def __init__(self): self.progress_bar = None + self.parameter_scheduler = None + self.frame_count = None def start_progress(self, total_steps, desc="Processing"): self.progress_bar = ProgressBar(total_steps) @@ -40,6 +65,41 @@ class FlexBase(ABC): def end_progress(self): self.progress_bar = None + def initialize_scheduler(self, frame_count: int, **kwargs): + """Initialize parameter scheduler with all numeric parameters""" + self.frame_count = frame_count + self.parameter_scheduler = ParameterScheduler(frame_count) + for key, value in kwargs.items(): + if isinstance(value, (int, float, list, tuple)): + self.parameter_scheduler.register_parameter(key, value) + + def get_parameter_value(self, param_name: str, frame_index: int, + feature_value: float = None, strength: float = 1.0, + mode: str = "relative") -> float: + """Get parameter value considering both scheduling and feature modulation""" + # Get base value (either scheduled or static) + if self.parameter_scheduler and self.parameter_scheduler.is_scheduled(param_name): + base_value = self.parameter_scheduler.get_value(param_name, frame_index) + else: + base_value = getattr(self, param_name, None) + + # If no feature value provided, return base value + if feature_value is None: + return base_value + + # Apply feature modulation to the base value + return self.modulate_param(param_name, base_value, feature_value, strength, mode) + + def get_feature_value(self, frame_index: int, feature=None, param_name=None): + """Get feature value from either a provided feature or a scheduled parameter""" + if feature is not None: + return feature.get_value_at_frame(frame_index) + elif self.parameter_scheduler and param_name: + feature_seq = self.parameter_scheduler.get_as_feature(param_name) + if feature_seq is not None: + return feature_seq[frame_index] + return None + @classmethod @abstractmethod def get_modifiable_params(cls): @@ -56,10 +116,61 @@ class FlexBase(ABC): @abstractmethod def apply_effect(self, *args, **kwargs): - """Main method to apply the effect.""" + """Apply the effect with potential parameter scheduling""" pass @abstractmethod def apply_effect_internal(self, *args, **kwargs): """Internal method to be implemented by subclasses.""" pass + + def process_parameters(self, frame_index: int = 0, feature_value: float = None, + strength: float = 1.0, feature_param: str = None, + feature_mode: str = "relative", **kwargs) -> dict: + """Process parameters considering both scheduling and feature modulation""" + # Initialize parameter scheduler if not already done + if self.parameter_scheduler is None: + # Determine frame count from the first array/list parameter we find + frame_count = kwargs.get('frame_count', 1) + for value in kwargs.values(): + if isinstance(value, (list, tuple, np.ndarray)): + frame_count = len(value) + break + self.initialize_scheduler(frame_count, **kwargs) + + # Get all parameters that could be scheduled + processed_kwargs = {} + for param_name, value in kwargs.items(): + # Skip non-numeric parameters and special parameters + if param_name in ['frame_count', 'frame_index']: + processed_kwargs[param_name] = value + continue + + try: + # Handle different types of inputs + if isinstance(value, (list, tuple, np.ndarray)): + # Convert numpy arrays to lists if needed + if isinstance(value, np.ndarray): + if value.ndim > 1: # If multi-dimensional array + value = value.flatten().tolist() # Flatten and convert to list + else: + value = value.tolist() + # Use frame_index to get the current value + try: + base_value = float(value[frame_index]) + except (IndexError, TypeError): + base_value = float(value[0]) # Fallback to first value + else: + # Single value + base_value = float(value) + + # Apply feature modulation if this is the target parameter + if param_name == feature_param and feature_value is not None: + processed_kwargs[param_name] = self.modulate_param(param_name, base_value, feature_value, strength, feature_mode) + else: + processed_kwargs[param_name] = base_value + except (ValueError, TypeError): + # If we can't convert to float, pass through unchanged + processed_kwargs[param_name] = value + + return processed_kwargs diff --git a/nodes/flex/parameter_scheduling.py b/nodes/flex/parameter_scheduling.py new file mode 100644 index 0000000..48e319a --- /dev/null +++ b/nodes/flex/parameter_scheduling.py @@ -0,0 +1,77 @@ +from typing import Union, List, Any +import numpy as np +from abc import ABC, abstractmethod + +class ScheduledParameter: + """Wrapper class for parameters that can be either single values or sequences""" + def __init__(self, value: Union[float, int, List[Union[float, int]]], frame_count: int): + self.original_value = value + self.frame_count = frame_count + self._sequence = None + self._initialize_sequence() + + def _initialize_sequence(self): + if isinstance(self.original_value, (list, tuple)): + # If it's a sequence, interpolate to match frame count + x = np.linspace(0, 1, len(self.original_value)) + y = np.array(self.original_value) + f = np.interp(np.linspace(0, 1, self.frame_count), x, y) + self._sequence = f + else: + # If it's a single value, repeat it + self._sequence = np.full(self.frame_count, self.original_value) + + def get_value(self, frame_index: int) -> Union[float, int]: + """Get the parameter value for a specific frame""" + if frame_index < 0 or frame_index >= self.frame_count: + raise ValueError(f"Frame index {frame_index} out of bounds [0, {self.frame_count})") + return float(self._sequence[frame_index]) + + def get_normalized_sequence(self) -> np.ndarray: + """Get the sequence normalized to [0,1] range for feature-like behavior""" + if self._sequence is None: + return None + seq = np.array(self._sequence) + min_val = np.min(seq) + max_val = np.max(seq) + if max_val > min_val: + return (seq - min_val) / (max_val - min_val) + return np.full_like(seq, 0.5) # If all values are the same, return 0.5 + + @property + def is_scheduled(self) -> bool: + """Returns True if the parameter is a sequence, False if it's a single value""" + return isinstance(self.original_value, (list, tuple)) + +class ParameterScheduler: + """Helper class to manage scheduled parameters for a node""" + def __init__(self, frame_count: int): + self.frame_count = frame_count + self.parameters = {} + + def register_parameter(self, name: str, value: Any) -> None: + """Register a parameter that might be scheduled""" + if isinstance(value, (int, float)) or (isinstance(value, (list, tuple)) and all(isinstance(x, (int, float)) for x in value)): + self.parameters[name] = ScheduledParameter(value, self.frame_count) + + def get_value(self, name: str, frame_index: int) -> Any: + """Get the value of a parameter for a specific frame""" + if name in self.parameters: + return self.parameters[name].get_value(frame_index) + raise KeyError(f"Parameter {name} not registered") + + def get_as_feature(self, name: str) -> np.ndarray: + """Get a parameter's sequence normalized as a feature (0-1 range)""" + if name in self.parameters: + return self.parameters[name].get_normalized_sequence() + return None + + def is_scheduled(self, name: str) -> bool: + """Check if a parameter is scheduled""" + if name in self.parameters: + return self.parameters[name].is_scheduled + return False + + def has_scheduled_parameters(self) -> bool: + """Check if any parameters are scheduled""" + return any(param.is_scheduled for param in self.parameters.values()) \ No newline at end of file diff --git a/nodes/images/flex_image_base.py b/nodes/images/flex_image_base.py index fd33dee..a7523ac 100644 --- a/nodes/images/flex_image_base.py +++ b/nodes/images/flex_image_base.py @@ -11,20 +11,21 @@ from ...tooltips import apply_tooltips class FlexImageBase(RyanOnTheInside, FlexBase): @classmethod def INPUT_TYPES(cls): - return { - **super().INPUT_TYPES(), - "required": { - **super().INPUT_TYPES()["required"], - "images": ("IMAGE",), - }, - # Optional inputs are inherited from FlexBase - } + base_inputs = super().INPUT_TYPES() + base_inputs["required"].update({ + "images": ("IMAGE",), + }) + base_inputs["optional"].update({ + "opt_feature_pipe": ("FEATURE_PIPE",), + }) + return base_inputs - CATEGORY = "RyanOnTheInside/FlexImage" + CATEGORY = "RyanOnTheInside/FlexImages" RETURN_TYPES = ("IMAGE",) FUNCTION = "apply_effect" def __init__(self): + super().__init__() # Initialize FlexBase self.progress_bar = None def start_progress(self, total_steps, desc="Processing"): @@ -117,30 +118,22 @@ 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: - # 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(image, **kwargs) - @abstractmethod def apply_effect_internal(self, image: np.ndarray, **kwargs) -> np.ndarray: - """To be implemented by subclasses.""" + """Apply the effect with processed parameters. To be implemented by child classes.""" pass + + def process_image(self, image: np.ndarray, feature_value: float, strength: float, + feature_param: str, feature_mode: str, frame_index: int = 0, **kwargs) -> np.ndarray: + # Process parameters using base class functionality + processed_kwargs = self.process_parameters( + frame_index=frame_index, + feature_value=feature_value, + strength=strength, + feature_param=feature_param, + feature_mode=feature_mode, + **kwargs + ) + + # Call the child class's implementation with processed parameters + return self.apply_effect_internal(image, **processed_kwargs) diff --git a/nodes/masks/flex_mask_base.py b/nodes/masks/flex_mask_base.py index 3f0ebfd..6eb0abf 100644 --- a/nodes/masks/flex_mask_base.py +++ b/nodes/masks/flex_mask_base.py @@ -57,26 +57,24 @@ class FlexMaskBase(FlexBase, MaskBase): return ["None"] @abstractmethod - def process_mask(self, mask: np.ndarray, feature_value: float, strength: float, **kwargs) -> np.ndarray: - """Process a single mask with feature modulation. + def process_mask(self, mask: np.ndarray, feature_value: float, strength: float, + frame_index: int = 0, **kwargs) -> np.ndarray: + # Process parameters using base class functionality + processed_kwargs = self.process_parameters( + frame_index=frame_index, + feature_value=feature_value, + strength=strength, + feature_param=kwargs.get('feature_param'), + feature_mode=kwargs.get('feature_mode'), + **kwargs + ) - Args: - mask: Input mask to process - feature_value: Current feature value for modulation - strength: Strength of the feature modulation - **kwargs: Additional parameters - - Returns: - Processed mask as numpy array - """ - pass + # Call the child class's implementation with processed parameters + return self.apply_effect_internal(mask, **processed_kwargs) - def apply_effect_internal(self, mask: np.ndarray, feature_value: float, strength: float, **kwargs) -> np.ndarray: - """Internal implementation for the Flex system. - - This bridges between FlexBase and our mask-specific processing. - """ - return self.process_mask(mask, feature_value, strength, **kwargs) + def apply_effect_internal(self, mask: np.ndarray, **kwargs) -> np.ndarray: + """Apply the effect with processed parameters.""" + return self.process_mask(mask, **kwargs) def apply_effect(self, masks, opt_feature=None, strength=1.0, feature_threshold=0.0, mask_strength=1.0, invert=False, subtract_original=0.0, grow_with_blur=0.0, **kwargs): """Main entry point for the Flex system. @@ -96,7 +94,9 @@ class FlexMaskBase(FlexBase, MaskBase): # When feature_param is "None", always apply the effect with default feature value if kwargs.get('feature_param') == "None": feature_value = 0.5 # Default feature value - processed_mask = self.apply_effect_internal(mask, feature_value, strength, **kwargs) + kwargs['feature_value'] = feature_value + kwargs['strength'] = strength + processed_mask = self.apply_effect_internal(mask, **kwargs) else: # Normal feature-based behavior if opt_feature is not None: @@ -107,10 +107,14 @@ class FlexMaskBase(FlexBase, MaskBase): apply_effect = True if apply_effect: - processed_mask = self.apply_effect_internal(mask, feature_value, strength, **kwargs) + kwargs['feature_value'] = feature_value + kwargs['strength'] = strength + processed_mask = self.apply_effect_internal(mask, **kwargs) else: if hasattr(self, 'process_mask_below_threshold'): - processed_mask = self.process_mask_below_threshold(mask, feature_value, strength, **kwargs) + kwargs['feature_value'] = feature_value + kwargs['strength'] = strength + processed_mask = self.process_mask_below_threshold(mask, **kwargs) else: processed_mask = mask @@ -120,8 +124,14 @@ class FlexMaskBase(FlexBase, MaskBase): self.end_progress() processed_masks = torch.from_numpy(np.stack(result)).float() + + # Remove strength from kwargs to avoid duplicate argument + kwargs_for_mask_op = kwargs.copy() + kwargs_for_mask_op.pop('strength', None) + kwargs_for_mask_op.pop('feature_value', None) + # Use mask_strength instead of strength for the final mask operation - return (super().apply_mask_operation(processed_masks, original_masks, mask_strength, invert, subtract_original, grow_with_blur, **kwargs),) + return (super().apply_mask_operation(processed_masks, original_masks, mask_strength, invert, subtract_original, grow_with_blur, **kwargs_for_mask_op),) def main_function(self, masks, opt_feature=None, strength=1.0, feature_threshold=0.0, mask_strength=1.0, invert=False, subtract_original=0.0, grow_with_blur=0.0, **kwargs): """Implementation of MaskBase's abstract main_function. diff --git a/web/js/drawable_feature.js b/web/js/drawable_feature.js index 01dc75f..a423074 100644 --- a/web/js/drawable_feature.js +++ b/web/js/drawable_feature.js @@ -177,7 +177,7 @@ app.registerExtension({ ctx.fillStyle = "#888"; ctx.font = "12px Arial"; ctx.textAlign = "center"; - ctx.fillText("Click empty space to add • Click point to delete • Drag points to move", + ctx.fillText("Click empty space to add • Double-click point to delete • Drag points to move", margin + graphWidth / 2, graphY + 20); // Draw grid