diff --git a/__init__.py b/__init__.py index d2c7244..0e223c0 100644 --- a/__init__.py +++ b/__init__.py @@ -95,6 +95,7 @@ from .nodes.flex.feature_extractors_audio import( PitchRangePresetNode, PitchRangeByNoteNode, PitchFeatureExtractor, + RhythmFeatureExtractor, ) from .nodes.flex.feature_extractors_midi import( @@ -155,7 +156,7 @@ from .nodes.utility_nodes import ( from .nodes.images.image_utility_nodes import ( DyeImage, - + ImageCASBatch, ) from .nodes.masks.flex_masks import ( @@ -191,7 +192,8 @@ from .nodes.images.flex_images import ( FlexImageTiltShift, FlexImageParallax, FlexImageContrast, - + FlexImageWarp, + FlexImageVignette, ) from .nodes.video.flex_video import ( @@ -223,9 +225,13 @@ from .nodes.flex.feature_modulation import ( FeatureMath, FeatureRebase, FeatureTruncateOrExtend, + FeatureAccumulate, ) - +from .nodes.audio.flex_audio import ( + FlexAudioPitchShift, + FlexAudioTimeStretch, +) import os import folder_paths @@ -288,6 +294,10 @@ NODE_CLASS_MAPPINGS = { #"FlexMaskNormalLighting": FlexMaskNormalLighting, # "FlexMaskDepthChamberRelative": FlexMaskDepthChamberRelative, + + #flex audio + "FlexAudioPitchShift": FlexAudioPitchShift, + "FlexAudioTimeStretch": FlexAudioTimeStretch, #audio "AudioSeparator": AudioSeparator, "DownloadOpenUnmixModel": DownloadOpenUnmixModel, @@ -314,6 +324,7 @@ NODE_CLASS_MAPPINGS = { #features "AudioFeatureExtractor": AudioFeatureExtractor, "PitchFeatureExtractor": PitchFeatureExtractor, + "RhythmFeatureExtractor": RhythmFeatureExtractor, "PitchRange": PitchRangeNode, "PitchRangePreset": PitchRangePresetNode, "PitchRangeByNoteNode": PitchRangeByNoteNode, @@ -346,7 +357,7 @@ NODE_CLASS_MAPPINGS = { "PreviewFeature": PreviewFeature, "FeatureRebase": FeatureRebase, "FeatureTruncateOrExtend": FeatureTruncateOrExtend, - + "FeatureAccumulate": FeatureAccumulate, #images 'FlexImageEdgeDetect': FlexImageEdgeDetect, "FlexImagePosterize": FlexImagePosterize, @@ -359,7 +370,8 @@ NODE_CLASS_MAPPINGS = { "FlexImageTiltShift": FlexImageTiltShift, "FlexImageParallax": FlexImageParallax, "FlexImageContrast": FlexImageContrast, - + "FlexImageWarp": FlexImageWarp, + "FlexImageVignette": FlexImageVignette, #opacity xp # "FlexDepthBasedMaskOpacity": FlexDepthBasedMaskOpacity, # "DepthBasedMaskOpacity": DepthBasedMaskOpacity, @@ -371,6 +383,7 @@ NODE_CLASS_MAPPINGS = { #garb "DyeImage": DyeImage, + "ImageCASBatch": ImageCASBatch, "MovingShape": MovingShape, "_mfc": _mfc, "TextMaskNode": TextMaskNode, @@ -436,6 +449,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "MotionFeatureNode":"Motion Feature", "FeatureMixer": "FeatureMod Mixer", + "FeatureAccumulate": "FeatureMod Accumulate", "FeatureCombine": "FeatureMod Combine", "FeatureOscillator": "FeatureMod Oscillator", "FeatureScaler": "FeatureMod Scaler", @@ -446,6 +460,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "DyeImage" : "Dye Image", + "ImageCASBatch": "Image Contrast Adaptive Sharpen Batch", "ImageIntervalSelectPercentage": "Image Interval Select %" # "FlexImageAdjustment":"Flex Image Adjustment", # "FlexImageFilter":"Flex Image Filter", diff --git a/nodes/audio/audio_utils.py b/nodes/audio/audio_utils.py index 22256d0..207f626 100644 --- a/nodes/audio/audio_utils.py +++ b/nodes/audio/audio_utils.py @@ -81,6 +81,10 @@ def time_stretch(waveform, rate): elif waveform.dim() > 3: raise ValueError("Input waveform has too many dimensions") + # Move tensor to CPU if it's on GPU + device = waveform.device + waveform = waveform.cpu() + # Convert to numpy for librosa processing waveform_np = waveform.numpy() @@ -93,6 +97,9 @@ def time_stretch(waveform, rate): # Stack channels and convert back to torch tensor stretched_waveform = torch.from_numpy(np.stack(stretched_channels)) + # Move the tensor back to the original device + stretched_waveform = stretched_waveform.to(device) + return stretched_waveform def dither_audio(waveform, bit_depth, noise_shaping): diff --git a/nodes/audio/flex_audio.py b/nodes/audio/flex_audio.py new file mode 100644 index 0000000..f2f7bbe --- /dev/null +++ b/nodes/audio/flex_audio.py @@ -0,0 +1,192 @@ +import torch +import numpy as np +from abc import ABC, abstractmethod +from tqdm import tqdm +from comfy.utils import ProgressBar +from ... import RyanOnTheInside +from .audio_utils import pitch_shift, time_stretch +import torch +import numpy as np +from abc import ABC, abstractmethod + +class FlexAudioBase(RyanOnTheInside, ABC): + @classmethod + def INPUT_TYPES(cls): + 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}), + } + } + + CATEGORY = "RyanOnTheInside/FlexAudio" + RETURN_TYPES = ("AUDIO",) + 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_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): + 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) + + self.start_progress(num_frames, desc=f"Applying {self.__class__.__name__}") + + if waveform.dim() == 2: + waveform = waveform.unsqueeze(0) # Add batch dimension + + batch_size, num_channels, total_samples = waveform.shape + samples_per_frame = max(1, total_samples // num_frames) + + # 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)] + + # Define cross-fade length (e.g., 10% of frame length) + crossfade_length = int(samples_per_frame * 0.1) + + for i in range(num_frames): + frame_start = i * samples_per_frame + frame_end = min((i + 1) * samples_per_frame + crossfade_length, total_samples) + audio_frame = waveform[:, :, frame_start:frame_end] # Shape: [Batch, Channels, Samples] + + feature_value = averaged_features[i] + + kwargs['frame_index'] = i + kwargs['sample_rate'] = sample_rate + + if feature_value >= feature_threshold: + try: + processed_frame = self.process_audio_frame( + audio_frame, feature_value, strength, + feature_param=feature_param, + feature_mode=feature_mode, + **kwargs + ) + except Exception as e: + import traceback + print(f"Error processing frame {i}:") + traceback.print_exc() + processed_frame = audio_frame + else: + processed_frame = audio_frame + + if i > 0: + # Apply cross-fade with previous frame + fade_in = torch.linspace(0, 1, crossfade_length).to(device) + fade_out = 1 - fade_in + crossfade_region = (processed_frames[-1][:, :, -crossfade_length:] * fade_out + + processed_frame[:, :, :crossfade_length] * fade_in) + processed_frames[-1][:, :, -crossfade_length:] = crossfade_region + processed_frame = processed_frame[:, :, crossfade_length:] + + processed_frames.append(processed_frame) + self.update_progress() + + self.end_progress() + + # Concatenate processed frames along the time dimension + result_waveform = torch.cat(processed_frames, dim=-1).cpu() + + return ({"waveform": result_waveform, "sample_rate": sample_rate},) + + 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: + modulated_value = self.modulate_param( + kwargs[feature_param], 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)) + + # Call the child class's implementation + return self.apply_effect_internal(audio_frame, **kwargs) + + @abstractmethod + def apply_effect_internal(self, audio_frame: torch.Tensor, **kwargs) -> torch.Tensor: + """Apply the effect to the audio frame. To be implemented by child classes.""" + pass + + + + + +class FlexAudioPitchShift(FlexAudioBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + **super().INPUT_TYPES()["required"], + "n_steps": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 12.0, "step": 0.1}), + } + } + + @classmethod + def get_modifiable_params(cls): + return ["n_steps"] + + 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 + return pitch_shift(audio_frame, sample_rate, n_steps) + +class FlexAudioTimeStretch(FlexAudioBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + **super().INPUT_TYPES()["required"], + "rate": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.0, "step": 0.01}), + } + } + + @classmethod + def get_modifiable_params(cls): + 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 diff --git a/nodes/flex/feature_extractors_audio.py b/nodes/flex/feature_extractors_audio.py index da768f7..3ea8e59 100644 --- a/nodes/flex/feature_extractors_audio.py +++ b/nodes/flex/feature_extractors_audio.py @@ -1,5 +1,6 @@ -from .feature_extractors import FeatureExtractorBase -from .features_audio import AudioFeature, PitchFeature, PitchRange, BaseFeature +from .feature_extractors import FeatureExtractorBase, FirstFeature +from .feature_pipe import FeaturePipe +from .features_audio import AudioFeature, PitchFeature, PitchRange, BaseFeature, RhythmFeature from ... import RyanOnTheInside class AudioFeatureExtractor(FeatureExtractorBase): @@ -27,6 +28,39 @@ class AudioFeatureExtractor(FeatureExtractorBase): feature.extract() return (feature, feature_pipe) +class RhythmFeatureExtractor(FirstFeature): + @classmethod + def feature_type(cls) -> type[BaseFeature]: + return RhythmFeature + + @classmethod + def INPUT_TYPES(cls): + return { + **super().INPUT_TYPES(), + "required": { + **super().INPUT_TYPES()["required"], + "audio": ("AUDIO",), + "time_signature": ("INT", {"default": 4, "min": 1, "max": 12, "step": 1}), + }, + } + + RETURN_TYPES = ("FEATURE", "FEATURE_PIPE") + FUNCTION = "extract_feature" + CATEGORY = "RyanOnTheInside/FlexFeatures/Audio/Rhythm" + + def extract_feature(self, audio, extraction_method, time_signature, frame_rate, video_frames): + feature_pipe = FeaturePipe(frame_rate, video_frames) + feature = RhythmFeature( + feature_name=extraction_method, + audio=audio, + frame_count=feature_pipe.frame_count, + frame_rate=feature_pipe.frame_rate, + feature_type=extraction_method, + time_signature=time_signature + ) + feature.extract() + return (feature, feature_pipe) + class PitchFeatureExtractor(FeatureExtractorBase): @classmethod def feature_type(cls) -> type[BaseFeature]: diff --git a/nodes/flex/feature_modulation.py b/nodes/flex/feature_modulation.py index 038757a..c41f050 100644 --- a/nodes/flex/feature_modulation.py +++ b/nodes/flex/feature_modulation.py @@ -107,7 +107,7 @@ class FeatureMixer(FeatureModulationBase): "smoothing": ("FLOAT", {"default": 0.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}), "rise_detection_threshold": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}), - "rise_smoothing_factor": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.05}), + "rise_smoothing_factor": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 5.0, "step": 0.05}), **super().INPUT_TYPES()["required"], # Include the invert_output option } } @@ -525,4 +525,49 @@ class FeatureTruncateOrExtend(FeatureModulationBase): adjusted_values = source_values processed_feature = self.create_processed_feature(feature, adjusted_values, "TruncatedOrExtended", invert_output) + return (processed_feature, self.visualize(processed_feature)) + + +class FeatureAccumulate(FeatureModulationBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "feature": ("FEATURE",), + "start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "end": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "skip_thresholded": ("BOOLEAN", {"default": False}), + **super().INPUT_TYPES()["required"], + } + } + + RETURN_TYPES = ("FEATURE", "IMAGE") + RETURN_NAMES = ("FEATURE", "FEATURE_VISUALIZATION") + FUNCTION = "accumulate" + + def accumulate(self, feature, start, end, threshold, skip_thresholded, invert_output): + values = [feature.get_value_at_frame(i) for i in range(feature.frame_count)] + + # Apply threshold and accumulate values + accumulated = [] + current_sum = 0 + for v in values: + if v >= threshold: + current_sum += v + accumulated.append(current_sum) + else: + if skip_thresholded: + accumulated.append(v) # Keep original value + else: + accumulated.append(current_sum) # Keep previous accumulated value + + # Normalize accumulated values between start and end + min_val, max_val = min(accumulated), max(accumulated) + if min_val == max_val: + normalized = [start for _ in accumulated] + else: + normalized = [start + (v - min_val) * (end - start) / (max_val - min_val) for v in accumulated] + + processed_feature = self.create_processed_feature(feature, normalized, "Accumulated", invert_output) return (processed_feature, self.visualize(processed_feature)) \ No newline at end of file diff --git a/nodes/flex/features_audio.py b/nodes/flex/features_audio.py index df7c9f8..613bd72 100644 --- a/nodes/flex/features_audio.py +++ b/nodes/flex/features_audio.py @@ -124,6 +124,156 @@ class AudioFeature(BaseAudioFeature): normalized = np.zeros_like(feature_array) self.features[self.feature_name] = normalized.tolist() +import numpy as np +import librosa + +class RhythmFeature(BaseAudioFeature): + def __init__(self, feature_name, audio, frame_count, frame_rate, feature_type='beat_locations', time_signature=4): + super().__init__(feature_name, audio, frame_count, frame_rate) + self.feature_type = feature_type + self.available_features = self.get_extraction_methods() + self.feature_name = feature_type + self.time_signature = time_signature # Default to 4/4 time + self._prepare_audio() + + @classmethod + def get_extraction_methods(cls): + return [ + 'beat_locations', + 'tempo', + 'onset_strength', + 'beat_emphasis', + 'syncopation', + 'rhythm_regularity', + 'down_beats', + 'up_beats' + ] + + def extract(self): + self.features = {} + + # Extract basic rhythm features + tempo, beat_frames = librosa.beat.beat_track(y=self.audio_array, sr=self.sample_rate) + onset_env = librosa.onset.onset_strength(y=self.audio_array, sr=self.sample_rate) + + # Calculate features based on the selected feature_type + if self.feature_type == 'beat_locations': + self._extract_beat_locations(beat_frames) + elif self.feature_type == 'tempo': + self._extract_tempo(tempo) + elif self.feature_type == 'onset_strength': + self._extract_onset_strength(onset_env) + elif self.feature_type == 'beat_emphasis': + self._extract_beat_emphasis(beat_frames, onset_env) + elif self.feature_type == 'syncopation': + self._extract_syncopation(beat_frames, onset_env) + elif self.feature_type == 'rhythm_regularity': + self._extract_rhythm_regularity(beat_frames) + elif self.feature_type in ['down_beats', 'up_beats']: + self._extract_beat_types(beat_frames) + else: + raise ValueError(f"Unsupported feature type: {self.feature_type}") + + self._normalize_features() + return self + + def _extract_beat_locations(self, beat_frames): + beat_times = librosa.frames_to_time(beat_frames, sr=self.sample_rate) + beat_sequence = np.zeros(self.frame_count) + for beat_time in beat_times: + frame_index = int(beat_time * self.frame_rate) + if frame_index < self.frame_count: + beat_sequence[frame_index] = 1 + self.features[self.feature_name] = beat_sequence.tolist() + + def _extract_tempo(self, tempo): + # For simplicity, we'll use a constant tempo for all frames + self.features[self.feature_name] = [tempo] * self.frame_count + + def _extract_onset_strength(self, onset_env): + # Resample onset strength to match frame count + resampled_onset = np.interp( + np.linspace(0, len(onset_env), self.frame_count), + np.arange(len(onset_env)), + onset_env + ) + self.features[self.feature_name] = resampled_onset.tolist() + + def _extract_beat_emphasis(self, beat_frames, onset_env): + beat_emphasis = np.zeros(self.frame_count) + for beat in beat_frames: + if beat < len(onset_env): + emphasis = onset_env[beat] + frame_index = int(librosa.frames_to_time(beat, sr=self.sample_rate) * self.frame_rate) + if frame_index < self.frame_count: + beat_emphasis[frame_index] = emphasis + self.features[self.feature_name] = beat_emphasis.tolist() + + def _extract_syncopation(self, beat_frames, onset_env): + # A simple syncopation measure: difference between actual onsets and expected beat locations + expected_beats = np.zeros_like(onset_env) + expected_beats[beat_frames] = 1 + syncopation = np.abs(onset_env - expected_beats) + + # Resample to match frame count + resampled_syncopation = np.interp( + np.linspace(0, len(syncopation), self.frame_count), + np.arange(len(syncopation)), + syncopation + ) + self.features[self.feature_name] = resampled_syncopation.tolist() + + def _extract_rhythm_regularity(self, beat_frames): + # Measure regularity by calculating the standard deviation of inter-beat intervals + ibi = np.diff(librosa.frames_to_time(beat_frames, sr=self.sample_rate)) + regularity = 1 / (1 + np.std(ibi)) # Invert so that higher values mean more regular + self.features[self.feature_name] = [regularity] * self.frame_count + + def _extract_beat_types(self, beat_frames): + beat_times = librosa.frames_to_time(beat_frames, sr=self.sample_rate) + down_beats = np.zeros(self.frame_count) + up_beats = np.zeros(self.frame_count) + + for i, beat_time in enumerate(beat_times): + frame_index = int(beat_time * self.frame_rate) + if frame_index < self.frame_count: + if i % self.time_signature == 0: + down_beats[frame_index] = 1 # Down beat + else: + up_beats[frame_index] = 1 # Up beat + + self.features['down_beats'] = down_beats.tolist() + self.features['up_beats'] = up_beats.tolist() + + def _normalize_features(self): + if self.feature_type in ['down_beats', 'up_beats']: + # For beat types, we want to keep the original binary values + self.features[self.feature_type] = self.features[self.feature_type] + else: + feature_array = np.array(self.features[self.feature_name], dtype=np.float32) + min_val = np.min(feature_array) + max_val = np.max(feature_array) + if max_val > min_val: + normalized = (feature_array - min_val) / (max_val - min_val) + else: + normalized = np.zeros_like(feature_array) + self.features[self.feature_name] = normalized.tolist() + + def get_rhythm_feature(self, frame_index): + if self.features is None: + self.extract() + + if self.feature_type in ['down_beats', 'up_beats']: + value = self.features[self.feature_type][frame_index] + beat_type = self.feature_type.rstrip('s') # Remove 's' to get "down_beat" or "up_beat" + return { + 'value': value, + 'beat_type': beat_type if value > 0 else "no_beat" + } + else: + return { + 'normalized': self.features[self.feature_name][frame_index] + } class PitchFeature(BaseAudioFeature): def __init__( diff --git a/nodes/images/flex_images.py b/nodes/images/flex_images.py index 6849ec4..80a4aad 100644 --- a/nodes/images/flex_images.py +++ b/nodes/images/flex_images.py @@ -476,3 +476,139 @@ class FlexImageContrast(FlexImageBase): result *= original_luminosity / current_luminosity return np.clip(result, 0, 1) + +import numpy as np +import cv2 + +class FlexImageWarp(FlexImageBase): + @classmethod + def INPUT_TYPES(cls): + return { + **super().INPUT_TYPES(), + "required": { + **super().INPUT_TYPES()["required"], + "warp_type": (["noise", "twist", "bulge"], {"default": "noise"}), + "warp_strength": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.0, "step": 0.01}), + "center_x": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "center_y": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "radius": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + }, + "optional": { + "warp_frequency": ("FLOAT", {"default": 5.0, "min": 0.1, "max": 20.0, "step": 0.1}), + "warp_octaves": ("INT", {"default": 3, "min": 1, "max": 5, "step": 1}), + "warp_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + } + } + + @classmethod + def get_modifiable_params(cls): + return ["warp_strength", "center_x", "center_y", "radius", "warp_frequency", "warp_octaves", "warp_seed", "None"] + + def apply_effect_internal(self, image: np.ndarray, warp_type: str, warp_strength: float, + center_x: float, center_y: float, radius: float, + warp_frequency: float = 5.0, warp_octaves: int = 3, + warp_seed: int = 0, **kwargs) -> np.ndarray: + h, w = image.shape[:2] + center = (int(w * center_x), int(h * center_y)) + + # Create meshgrid + y, x = np.mgrid[0:h, 0:w].astype(np.float32) + + # Calculate distance from center + dx = x - center[0] + dy = y - center[1] + dist = np.sqrt(dx**2 + dy**2) + + # Create a mask based on the radius + max_dist = np.sqrt(w**2 + h**2) + mask = np.clip(1 - dist / (radius * max_dist), 0, 1) + + if warp_type == "noise": + # Generate noise for warping + np.random.seed(warp_seed) + noise = np.zeros((h, w, 2)) + for _ in range(warp_octaves): + freq = warp_frequency * (2 ** _) + amp = warp_strength / (2 ** _) + noise += amp * np.random.rand(h, w, 2) * np.sin(freq * np.stack((y, x)) / np.array([h, w])) + + x_warped = x + noise[:,:,0] * w * mask + y_warped = y + noise[:,:,1] * h * mask + + elif warp_type == "twist": + angle = np.arctan2(dy, dx) + twist = warp_strength * dist * mask + x_warped = x + np.sin(angle + twist) * dist - dx + y_warped = y - np.cos(angle + twist) * dist - dy + + elif warp_type == "bulge": + bulge = 1 + warp_strength * mask + x_warped = center[0] + dx * bulge + y_warped = center[1] + dy * bulge + + else: + raise ValueError(f"Unknown warp type: {warp_type}") + + # Ensure warped coordinates are within image bounds + x_warped = np.clip(x_warped, 0, w-1) + y_warped = np.clip(y_warped, 0, h-1) + + # Remap image + warped = cv2.remap(image, x_warped, y_warped, cv2.INTER_LINEAR) + + return warped + + +class FlexImageVignette(FlexImageBase): + @classmethod + def INPUT_TYPES(cls): + return { + **super().INPUT_TYPES(), + "required": { + **super().INPUT_TYPES()["required"], + "intensity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "radius": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 2.0, "step": 0.01}), + "feather": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + "center_x": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "center_y": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + @classmethod + def get_modifiable_params(cls): + return ["intensity", "radius", "feather", "center_x", "center_y", "None"] + + def apply_effect_internal(self, image: np.ndarray, intensity: float, radius: float, feather: float, + center_x: float, center_y: float, **kwargs) -> np.ndarray: + h, w = image.shape[:2] + + # Create coordinate grids + y, x = np.ogrid[:h, :w] + center = (int(w * center_x), int(h * center_y)) + + # Calculate distance from center + dist = np.sqrt((x - center[0])**2 + (y - center[1])**2) + + # Normalize distance + max_dist = np.sqrt((max(center[0], w - center[0]))**2 + (max(center[1], h - center[1]))**2) + normalized_dist = dist / max_dist + + # Apply radius + normalized_dist = normalized_dist / radius + + # Create vignette mask + mask = 1 - np.clip(normalized_dist, 0, 1) + + # Apply feathering + mask = np.clip((mask - (1 - feather)) / feather, 0, 1) + + # Apply intensity + mask = mask * intensity + (1 - intensity) + + # Reshape mask to match image dimensions + mask = mask.reshape(h, w, 1) + + # Apply vignette + result = image * mask + + return np.clip(result, 0, 1) \ No newline at end of file diff --git a/nodes/images/image_utility_nodes.py b/nodes/images/image_utility_nodes.py index 5198691..ce1ba33 100644 --- a/nodes/images/image_utility_nodes.py +++ b/nodes/images/image_utility_nodes.py @@ -1,8 +1,22 @@ import torch +import torch.nn.functional as F from ..node_utilities import string_to_rgb from ... import RyanOnTheInside +from comfy.utils import ProgressBar class ImageUtilityNode(RyanOnTheInside): + 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 CATEGORY = "RyanOnTheInside/image/utility" class DyeImage(ImageUtilityNode): @@ -34,3 +48,76 @@ class DyeImage(ImageUtilityNode): return (result,) + +# From https://github.com/Jamy-L/Pytorch-Contrast-Adaptive-Sharpening/ +# THEN from comfyui essentials shoutout MATEO +class ImageCASBatch(ImageUtilityNode): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "amount": ("FLOAT", {"default": 0.8, "min": 0, "max": 1, "step": 0.05}), + "batch_size": ("INT", {"default": 4, "min": 1, "max": 64, "step": 1}), + }, + } + + RETURN_TYPES = ("IMAGE",) + CATEGORY = "essentials/image processing" + FUNCTION = "execute" + + def execute(self, image, amount, batch_size): + # Process images in batches + output_batches = [] + self.start_progress(image.shape[0] // batch_size) + for i in range(0, image.shape[0], batch_size): + batch = image[i:i+batch_size] + output_batch = self.process_batch(batch, amount) + output_batches.append(output_batch) + self.update_progress() + self.end_progress() + # Concatenate all processed batches + output = torch.cat(output_batches, dim=0) + return (output,) + + def process_batch(self, batch, amount): + epsilon = 1e-5 + img = F.pad(batch.permute([0,3,1,2]), pad=(1, 1, 1, 1)) + + a = img[..., :-2, :-2] + b = img[..., :-2, 1:-1] + c = img[..., :-2, 2:] + d = img[..., 1:-1, :-2] + e = img[..., 1:-1, 1:-1] + f = img[..., 1:-1, 2:] + g = img[..., 2:, :-2] + h = img[..., 2:, 1:-1] + i = img[..., 2:, 2:] + + # Computing contrast + cross = (b, d, e, f, h) + mn = torch.min(torch.stack(cross), dim=0).values + mx = torch.max(torch.stack(cross), dim=0).values + + diag = (a, c, g, i) + mn2 = torch.min(torch.stack(diag), dim=0).values + mx2 = torch.max(torch.stack(diag), dim=0).values + mx = mx + mx2 + mn = mn + mn2 + + # Computing local weight + inv_mx = torch.reciprocal(mx + epsilon) + amp = inv_mx * torch.minimum(mn, (2 - mx)) + + # scaling + amp = torch.sqrt(amp) + w = - amp * (amount * (1/5 - 1/8) + 1/8) + div = torch.reciprocal(1 + 4*w) + + output = ((b + d + f + h)*w + e) * div + output = output.clamp(0, 1) + #output = torch.nan_to_num(output) + + output = output.permute([0,2,3,1]) + + return output diff --git a/nodes/video/config.yaml b/nodes/video/config.yaml new file mode 100644 index 0000000..b99d4a7 --- /dev/null +++ b/nodes/video/config.yaml @@ -0,0 +1,3 @@ +#Plz don't delete this file, just edit it when neccessary. +ckpts_path: "./ckpts" +ops_backend: "cupy" #Either "taichi" or "cupy" \ No newline at end of file diff --git a/nodes/video/flex_video.py b/nodes/video/flex_video.py index 24f0b72..64f18fa 100644 --- a/nodes/video/flex_video.py +++ b/nodes/video/flex_video.py @@ -69,14 +69,16 @@ class FlexVideoSeek(FlexVideoBase): @classmethod def get_modifiable_params(cls): - return ["seek_speed"] + return ["seek"] @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "seek_speed": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}), - }) + + # Override the feature_mode input to remove the "absolute" option + inputs["required"]["feature_mode"] = (["relative"], {"default": "relative"}) + # Add the new reverse parameter + inputs["required"]["reverse"] = ("BOOLEAN", {"default": False}) return inputs FUNCTION = "apply_effect" @@ -85,33 +87,40 @@ class FlexVideoSeek(FlexVideoBase): self, video: np.ndarray, feature_values: np.ndarray, - seek_speed: float, + reverse: bool, **kwargs, ) -> np.ndarray: num_frames = video.shape[0] processed_video = np.empty_like(video) - # Get strength from kwargs strength = kwargs.get('strength', 1.0) - # Compute cumulative feature values to determine frame positions - feature_values_clipped = np.clip(feature_values, 0.0, 1.0) # Ensure values are between 0 and 1 - adjusted_speeds = feature_values_clipped * seek_speed * strength# Adjust speeds according to feature values + seek_speed = 1.0 - # Normalize adjusted speeds to ensure total frames match input + # Reverse the video if the reverse parameter is True + if reverse: + video = video[::-1] + + # Clip feature values between 0 and 1 + feature_values_clipped = np.clip(feature_values, 0.0, 1.0) + + # Apply seek_speed to feature values + adjusted_speeds = feature_values_clipped * seek_speed * strength + + # Ensure the total adjusted speed matches the number of frames total_speed = np.sum(adjusted_speeds) if total_speed == 0: - adjusted_speeds = np.ones(num_frames) * (num_frames / num_frames) + adjusted_speeds = np.ones(num_frames) else: - adjusted_speeds = adjusted_speeds / total_speed * (num_frames - 1) + adjusted_speeds = adjusted_speeds / total_speed * num_frames - cumulative_speeds = np.cumsum(adjusted_speeds) + # Calculate cumulative frame positions + cumulative_positions = np.cumsum(adjusted_speeds) - # Map frame indices based on cumulative speeds - frame_indices = np.clip(cumulative_speeds.astype(int), 0, num_frames - 1) + # Map frame indices based on cumulative positions + frame_indices = np.clip(cumulative_positions.astype(int), 0, num_frames - 1) - # Ensure the output video has the same number of frames + # Create the processed video by selecting frames based on the mapped indices for idx in range(num_frames): - frame_idx = frame_indices[idx] - processed_video[idx] = video[frame_idx] + processed_video[idx] = video[frame_indices[idx]] return processed_video diff --git a/nodes/video/flex_video_speed.py b/nodes/video/flex_video_speed.py index c2d05b3..8ba3b12 100644 --- a/nodes/video/flex_video_speed.py +++ b/nodes/video/flex_video_speed.py @@ -5,7 +5,7 @@ from torch.hub import download_url_to_file, get_dir import folder_paths import traceback import torch -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames +from vfi_utils import preprocess_frames, postprocess_frames from .video_base import FlexVideoBase import numpy as np from ..flex.feature_pipe import FeaturePipe