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

Merge branch 'base-class-refactor' into version2

Browse Source
This commit is contained in:
ryanontheinside
2025-01-03 08:52:27 -05:00
parent 4d8675f866 9cf6d2f20d
commit 9c12f409b5
16 changed files with 581 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+1 -1
View File
@@ -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,
nodes/audio/audio_nodes.py
+1 -2
View File
@@ -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):
nodes/audio/audio_nodes_effects.py
-2
View File
@@ -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):
nodes/audio/audio_nodes_utility.py
-1
View File
@@ -229,7 +229,6 @@ class AudioInfo(AudioUtility):
rms_amplitude,
bit_depth
)
class AudioDither(AudioUtility):
@classmethod
def INPUT_TYPES(cls):
nodes/audio/flex_audio.py
+75 -56
View File
@@ -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)
return time_stretch(audio_frame, rate)
nodes/audio/flex_audio_visualizer.py
+98 -160
View File
@@ -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
nodes/depth/depth_base.py
+81 -44
View File
@@ -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
nodes/flex/flex_base.py
+64
View File
@@ -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
nodes/images/image_base.py → nodes/images/flex_image_base.py
+61 -36
View File
@@ -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
nodes/images/flex_images.py
+2 -1
View File
@@ -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):
nodes/latents/flex_latent_base.py
+105
View File
@@ -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
nodes/latents/latent_base.py → nodes/latents/flex_latents.py
+30 -148
View File
@@ -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
return result
nodes/masks/flex_mask_base.py
+59
View File
@@ -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
nodes/masks/flex_masks.py
+1 -3
View File
@@ -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
nodes/masks/flex_masks_normal.py
+2 -1
View File
@@ -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):
nodes/masks/mask_base.py
+1 -55
View File
@@ -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