diff --git a/nodes/flex/flex_base.py b/nodes/flex/flex_base.py new file mode 100644 index 0000000..1e9a6c5 --- /dev/null +++ b/nodes/flex/flex_base.py @@ -0,0 +1,61 @@ +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": { + "feature": ("FEATURE",), + "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/image_base.py index 7f405a7..cdf211a 100644 --- a/nodes/images/image_base.py +++ b/nodes/images/image_base.py @@ -4,17 +4,16 @@ 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",), @@ -51,28 +50,32 @@ 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): - 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.") + 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.") - if opt_feature is None and opt_feature_pipe is None: - # If neither feature nor feature_pipe is provided, process all frames without modulation + if feature is None and feature_pipe is None: + # 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 - num_frames = opt_feature_pipe.frame_count + # Process frames with modulation based on feature values + num_frames = feature_pipe.frame_count images_np = images.cpu().numpy() self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") @@ -80,13 +83,17 @@ class FlexImageBase(RyanOnTheInside, ABC): result = [] for i in range(num_frames): image = images_np[i] - feature_value = opt_feature.get_value_at_frame(i) + feature_value = 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 +104,31 @@ 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, + 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/masks/mask_base.py b/nodes/masks/mask_base.py index ed36032..7363361 100644 --- a/nodes/masks/mask_base.py +++ b/nodes/masks/mask_base.py @@ -275,7 +275,7 @@ 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 +#TODO SOON TO UPDATE TO SUB FLEXBASE. UNTIL THEN REFER TO IMAGES AUDIO OR VIDEO FOR EXAMPLES ON HOW TO EXTEND class FlexMaskBase(MaskBase): feature_threshold_default=0.0 @classmethod