nodes/ryanontheinside-ComfyUI_RyanOnTheInside
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

flex base

Browse Source
This commit is contained in:
ryanontheinstide
2024-10-23 14:42:06 -04:00
parent 27c1ffe3c0
commit e15ac1fe30
3 changed files with 106 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/flex/flex_base.py
+61
View File
@@ -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
nodes/images/image_base.py
+44 -35
View File
@@ -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
nodes/masks/mask_base.py
+1 -1
View File
@@ -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