From eefa76397bedcc038ba84041f1a6ffd76b0ae98d Mon Sep 17 00:00:00 2001 From: ryanontheinstide Date: Sat, 26 Oct 2024 14:00:01 -0400 Subject: [PATCH] refactor --- __init__.py | 2 +- nodes/audio/audio_nodes.py | 3 +- nodes/audio/flex_audio.py | 4 +- nodes/depth/depth_base.py | 125 +++++--- nodes/flex/feature_extractors_midi.py | 2 +- .../{midi_feature.py => features_midi.py} | 0 nodes/flex/midi_feature_extractor.py | 163 ---------- nodes/flex/proximity_feature.py | 92 ------ nodes/flex/proximity_feature_extractor.py | 279 ------------------ .../{image_base.py => flex_image_base.py} | 0 nodes/images/flex_images.py | 2 +- nodes/latents/flex_latent_base.py | 105 +++++++ .../{latent_base.py => flex_latents.py} | 162 +--------- 13 files changed, 203 insertions(+), 736 deletions(-) rename nodes/flex/{midi_feature.py => features_midi.py} (100%) delete mode 100644 nodes/flex/midi_feature_extractor.py delete mode 100644 nodes/flex/proximity_feature.py delete mode 100644 nodes/flex/proximity_feature_extractor.py rename nodes/images/{image_base.py => flex_image_base.py} (100%) create mode 100644 nodes/latents/flex_latent_base.py rename nodes/latents/{latent_base.py => flex_latents.py} (59%) diff --git a/__init__.py b/__init__.py index f562581..ae7d0ba 100644 --- a/__init__.py +++ b/__init__.py @@ -253,7 +253,7 @@ from .nodes.audio.flex_audio import ( FlexAudioTimeStretch, ) -from .nodes.latents.latent_base import ( +from .nodes.latents.flex_latents import ( FlexLatentInterpolate, EmbeddingGuidedLatentInterpolate, FlexLatentBlend, diff --git a/nodes/audio/audio_nodes.py b/nodes/audio/audio_nodes.py index b51674b..b49e048 100644 --- a/nodes/audio/audio_nodes.py +++ b/nodes/audio/audio_nodes.py @@ -496,8 +496,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): diff --git a/nodes/audio/flex_audio.py b/nodes/audio/flex_audio.py index cf40917..af225a0 100644 --- a/nodes/audio/flex_audio.py +++ b/nodes/audio/flex_audio.py @@ -1,13 +1,13 @@ 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 -class FlexAudioBase(FlexBase): +class FlexAudioBase(FlexBase, RyanOnTheInside): @classmethod def INPUT_TYPES(cls): base_input_types = super().INPUT_TYPES() diff --git a/nodes/depth/depth_base.py b/nodes/depth/depth_base.py index a5aaddb..fa8b4ee 100644 --- a/nodes/depth/depth_base.py +++ b/nodes/depth/depth_base.py @@ -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 + + diff --git a/nodes/flex/feature_extractors_midi.py b/nodes/flex/feature_extractors_midi.py index 6fd40ac..9a390d0 100644 --- a/nodes/flex/feature_extractors_midi.py +++ b/nodes/flex/feature_extractors_midi.py @@ -5,7 +5,7 @@ import folder_paths from server import PromptServer from aiohttp import web import shutil -from .midi_feature import MIDIFeature +from .features_midi import MIDIFeature from .feature_extractors import FeatureExtractorBase class MIDILoadAndExtract(FeatureExtractorBase): diff --git a/nodes/flex/midi_feature.py b/nodes/flex/features_midi.py similarity index 100% rename from nodes/flex/midi_feature.py rename to nodes/flex/features_midi.py diff --git a/nodes/flex/midi_feature_extractor.py b/nodes/flex/midi_feature_extractor.py deleted file mode 100644 index b065939..0000000 --- a/nodes/flex/midi_feature_extractor.py +++ /dev/null @@ -1,163 +0,0 @@ -import mido -import os -from .feature_pipe import FeaturePipe -import folder_paths -from server import PromptServer -from aiohttp import web -import shutil -from .midi_feature import MIDIFeature -from .feature_extractors import FeatureExtractorBase - -class MIDILoadAndExtract(FeatureExtractorBase): - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "midi_file": (folder_paths.get_filename_list("midi_files"),), - "track_selection": (["all"],), - "attribute": (MIDIFeature.get_attribute_names(), {"default": "Note On/Off"}), - "frame_rate": ("FLOAT", {"default": 30, "min": 0.1, "max": 120, "step": 0.1}), - "video_frames": ("IMAGE",), - "chord_only": ("BOOLEAN", {"default": False}), - "notes": ("STRING", {"default": ""}), - }, - } - - RETURN_TYPES = ("MIDI", "FEATURE", "FEATURE_PIPE") - FUNCTION = "process_midi" - CATEGORY = "RyanOnTheInside/Audio" - - def process_midi(self, midi_file, track_selection, notes, attribute, frame_rate, video_frames, chord_only=False): - try: - midi_path = folder_paths.get_full_path("midi_files", midi_file) - if not midi_path or not os.path.exists(midi_path): - raise FileNotFoundError(f"MIDI file not found: {midi_file}") - - midi_data = mido.MidiFile(midi_path) - - selected_notes = [int(n.strip()) for n in notes.split(',') if n.strip().isdigit()] - feature_pipe = FeaturePipe(frame_rate, video_frames) - - # Convert friendly attribute name to internal attribute name - internal_attribute = MIDIFeature.get_attribute_value(attribute) - - feature = MIDIFeature(f"midi_{internal_attribute}", midi_data, internal_attribute, - feature_pipe.frame_rate, feature_pipe.frame_count, - notes=selected_notes, chord_only=chord_only) - - feature.extract() - - return (midi_data, feature, feature_pipe) - - except Exception as e: - # error_msg = f"Error in MIDILoadAndExtract.process_midi: {type(e).__name__}: {str(e)}\n" - # error_msg += traceback.format_exc() - # print(error_msg) - raise RuntimeError(f"Error processing MIDI file: {type(e).__name__}: {str(e)}") - - @classmethod - def analyze_midi(cls, midi_path): - midi_data = mido.MidiFile(midi_path) - tracks = ["all"] - all_notes = set() - track_notes = {} - for i, track in enumerate(midi_data.tracks): - track_notes[str(i)] = set() - for msg in track: - if msg.type == 'note_on': - track_notes[str(i)].add(msg.note) - all_notes.add(msg.note) - if len(track_notes[str(i)]) == 0: - tracks.append(f"{i}: (Empty)") - else: - tracks.append(f"{i}: {track.name or f'Track {i}'}") - - return { - "tracks": tracks, - "all_notes": ",".join(map(str, sorted(set(all_notes)))), - "track_notes": {k: ",".join(map(str, sorted(v))) for k, v in track_notes.items()} - } - - @classmethod - def VALIDATE_INPUTS(cls, midi_file, track_selection, notes, attribute, frame_rate, video_frames): - if not folder_paths.exists_and_is_file(midi_file): - return "MIDI file not found: {}".format(midi_file) - - if notes != "all": - try: - note_list = [int(n.strip()) for n in notes.split(',') if n.strip()] - if not all(0 <= n <= 127 for n in note_list): - return "Invalid note value. All notes must be between 0 and 127." - except ValueError: - return "Invalid notes format. Please provide comma-separated integers or 'all'." - - return True - - -routes = PromptServer.instance.routes -@PromptServer.instance.routes.post('/get_track_notes') -async def get_track_notes(request): - data = await request.json() - midi_file = data.get('midi_file') - - if not midi_file: - return web.json_response({"error": "Missing required parameters"}, status=400) - - midi_path = folder_paths.get_full_path("midi_files", midi_file) - if not midi_path or not os.path.exists(midi_path): - return web.json_response({"error": "MIDI file not found"}, status=404) - - analysis = MIDILoadAndExtract.analyze_midi(midi_path) - return web.json_response(analysis) - -@routes.post('/upload_midi') -async def upload_midi(request): - data = await request.post() - midi_file = data['file'] - - if midi_file and midi_file.filename: - safe_filename = os.path.basename(midi_file.filename) - - midi_dir = folder_paths.get_folder_paths("midi_files")[0] - - if not os.path.exists(midi_dir): - os.makedirs(midi_dir, exist_ok=True) - - midi_path = os.path.join(midi_dir, safe_filename) - - with open(midi_path, 'wb') as f: - shutil.copyfileobj(midi_file.file, f) - - midi_files = folder_paths.get_filename_list("midi_files") - analysis = MIDILoadAndExtract.analyze_midi(midi_path) - - return web.json_response({ - "status": "success", - "uploaded_file": safe_filename, - "midi_files": midi_files, - "analysis": analysis - }) - else: - return web.json_response({"status": "error", "message": "No file uploaded"}, status=400) - -@PromptServer.instance.routes.post('/refresh_midi_data') -async def refresh_midi_data(request): - data = await request.json() - midi_file = data.get('midi_file') - track_selection = data.get('track_selection') - - if not midi_file: - return web.json_response({"error": "Missing required parameters"}, status=400) - - midi_path = folder_paths.get_full_path("midi_files", midi_file) - if not midi_path or not os.path.exists(midi_path): - return web.json_response({"error": "MIDI file not found"}, status=404) - - analysis = MIDILoadAndExtract.analyze_midi(midi_path) - - # Filter notes based on track selection - if track_selection != "all": - track_index = track_selection.split(':')[0] - analysis['all_notes'] = analysis['track_notes'].get(track_index, "") - - return web.json_response(analysis) \ No newline at end of file diff --git a/nodes/flex/proximity_feature.py b/nodes/flex/proximity_feature.py deleted file mode 100644 index 281646c..0000000 --- a/nodes/flex/proximity_feature.py +++ /dev/null @@ -1,92 +0,0 @@ -from .features import BaseFeature -import torch -import numpy as np - -import numpy as np -from .features import BaseFeature - -class ProximityFeature(BaseFeature): - def __init__(self, name, anchor_locations, query_locations, frame_rate, frame_count, frame_dimensions, normalization_method='frame'): - super().__init__(name, "proximity", frame_rate, frame_count) - self.anchor_locations = anchor_locations - self.query_locations = query_locations - self.frame_diagonal = np.sqrt(frame_dimensions[0]**2 + frame_dimensions[1]**2) - self.proximity_values = None - self.normalization_method = normalization_method - - def extract(self): - from scipy.spatial.distance import cdist - - proximities = [] - for anchor, query in zip(self.anchor_locations, self.query_locations): - if len(anchor) == 0 or len(query) == 0: - proximities.append(self.frame_diagonal if self.normalization_method == 'frame' else float('inf')) - else: - distances = cdist(anchor.points, query.points) - min_distance = np.min(distances) - proximities.append(min_distance) - - proximities = np.array(proximities) - - if self.normalization_method == 'frame': - self.proximity_values = 1 - np.clip(proximities / self.frame_diagonal, 0, 1) - elif self.normalization_method == 'minmax': - min_proximity = np.min(proximities) - max_proximity = np.max(proximities) - if max_proximity > min_proximity: - self.proximity_values = 1 - (proximities - min_proximity) / (max_proximity - min_proximity) - else: - self.proximity_values = np.ones_like(proximities) - else: - raise ValueError(f"Unknown normalization method: {self.normalization_method}") - - self.smooth_proximities() - return self - - def smooth_proximities(self, window_size=5): - kernel = np.ones(window_size) / window_size - self.proximity_values = np.convolve(self.proximity_values, kernel, mode='same') - - def get_value_at_frame(self, frame_index): - return self.proximity_values[frame_index] - -class Location: - def __init__(self, x, y, z=None): - x = np.asarray(x).reshape(-1) - y = np.asarray(y).reshape(-1) - - if z is not None: - z = np.asarray(z).reshape(-1) - # Ensure z has the same length as x and y, with a small tolerance for floating-point imprecision - if abs(len(z) - len(x)) <= 1: # Allow for off-by-one errors - z = z[:len(x)] if len(z) > len(x) else np.pad(z, (0, len(x) - len(z)), 'constant', constant_values=np.nan) - elif len(z) != len(x): - raise ValueError(f"Mismatch in dimensions: x,y have length {len(x)}, but z has length {len(z)}") - self.points = np.column_stack((x, y, z)) - else: - self.points = np.column_stack((x, y)) - - def __len__(self): - return len(self.points) - - def __getitem__(self, idx): - return self.points[idx] - - @property - def x(self): - return self.points[:, 0] - - @property - def y(self): - return self.points[:, 1] - - @property - def z(self): - return self.points[:, 2] if self.points.shape[1] > 2 else None - - @classmethod - def from_tensor(cls, tensor): - return cls(tensor[:, 0], tensor[:, 1], tensor[:, 2] if tensor.shape[1] > 2 else None) - - def to_tensor(self): - return torch.from_numpy(self.points) \ No newline at end of file diff --git a/nodes/flex/proximity_feature_extractor.py b/nodes/flex/proximity_feature_extractor.py deleted file mode 100644 index e8d75b6..0000000 --- a/nodes/flex/proximity_feature_extractor.py +++ /dev/null @@ -1,279 +0,0 @@ -from .proximity_feature import Location, ProximityFeature -from .feature_extractors import FirstFeature -from ... import RyanOnTheInside -from .feature_pipe import FeaturePipe -import numpy as np -import cv2 -import torch -from scipy.spatial.distance import cdist # Add this import - -class ProximityFeatureNode(FirstFeature): - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "video_frames": ("IMAGE",), - "frame_rate": ("FLOAT", {"default": 30.0, "min": 1.0, "max": 120.0, "step": 0.1}), - "anchor_locations": ("LOCATION",), - "query_locations": ("LOCATION",), - "normalization_method": (["frame", "minmax"],), - } - } - - RETURN_TYPES = ("FEATURE", "FEATURE_PIPE") - RETURN_NAMES = ("proximity_feature", "feature_pipe") - FUNCTION = "create_feature" - - def create_feature(self, video_frames, frame_rate, anchor_locations, query_locations, normalization_method): - feature_pipe = FeaturePipe(frame_rate, video_frames) - frame_dimensions = (video_frames.shape[2], video_frames.shape[1]) # width, height - - proximity_feature = ProximityFeature( - name="proximity_feature", - anchor_locations=anchor_locations, - query_locations=query_locations, - frame_rate=frame_rate, - frame_count=len(video_frames), - frame_dimensions=frame_dimensions, - normalization_method=normalization_method - ) - - # Ensure anchor and query points have the same number of columns - anchor_dim = anchor_locations[0].points.shape[1] - query_dim = query_locations[0].points.shape[1] - - if anchor_dim != query_dim: - if anchor_dim == 3: - for location in anchor_locations: - location.points = location.points[:, :2] - elif query_dim == 3: - for location in query_locations: - location.points = location.points[:, :2] - print("Depth ignored. Depth must be included for all points or no points.") - - # self.start_progress(len(video_frames), desc="Calculating proximity") - proximity_feature.extract() - # self.end_progress() - - return (proximity_feature, feature_pipe) - -class ProximityFeatureInput(RyanOnTheInside): - CATEGORY="RyanOnTheInside/Proximity" - -class LocationFromMask(ProximityFeatureInput): - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "masks": ("MASK",), - "method": (["mask_boundary", "mask_center"],), - }, - "optional": { - "depth_maps": ("IMAGE",), - } - } - - RETURN_TYPES = ("LOCATION",) - FUNCTION = "generate_locations" - - def generate_locations(self, masks, method, depth_maps=None): - locations = [] - for i in range(masks.shape[0]): - mask = masks[i].cpu().numpy() - if np.sum(mask) == 0: # Check if the mask is empty - print(f"No mask found in frame {i}.") - locations.append(Location(np.array([-1]), np.array([-1]), np.array([-1]) if depth_maps is not None else None)) - continue - - if method == "mask_boundary": - contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - if contours: - boundary = np.vstack(contours).squeeze() - else: - boundary = np.array([[0, 0]]) # Fallback to a single point - elif method == "mask_center": - y, x = np.where(mask > 0.5) - if len(x) > 0 and len(y) > 0: - center_x, center_y = np.mean(x), np.mean(y) - boundary = np.array([[center_x, center_y]]) - else: - boundary = np.array([[mask.shape[1] / 2, mask.shape[0] / 2]]) - - if depth_maps is not None: - depth_map = depth_maps[i].cpu().numpy() - # Ensure boundary points are within the depth map dimensions - valid_points = (boundary[:, 0] < depth_map.shape[1]) & (boundary[:, 1] < depth_map.shape[0]) - boundary = boundary[valid_points] - - if boundary.size == 0: # Check if boundary is empty after filtering - print(f"No valid boundary points in frame {i}.") - locations.append(Location(np.array([-1]), np.array([-1]), np.array([-1]))) - continue - - # Extract depth values for valid boundary points - z = np.mean(depth_map[boundary[:, 1].astype(int), boundary[:, 0].astype(int)], axis=-1) - - location = Location(boundary[:, 0], boundary[:, 1], z) - else: - location = Location(boundary[:, 0], boundary[:, 1]) - - locations.append(location) - - return (locations,) - -class LocationFromPoint(ProximityFeatureInput): - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "x": ("FLOAT", {"default": 0.0,"min":0.0,"step":0.01}), - "y": ("FLOAT", {"default": 0.0,"min":0.0,"step":0.01}), - "z": ("FLOAT", {"default": 0.0, "max":1.0, "min":0.0,"step":0.01}), - "batch_count": ("INT", {"default": 1, "min": 1}), - } - } - - RETURN_TYPES = ("LOCATION",) - FUNCTION = "generate_locations" - - def generate_locations(self, x, y, batch_count, z): - locations = [] - for _ in range(batch_count): - location = Location(np.array([x]), np.array([y]), np.array([z])) - locations.append(location) - - return (locations,) - -class LocationTransform(ProximityFeatureInput): - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "locations": ("LOCATION",), - "feature": ("FEATURE",), - "transformation_type": (["translate", "scale"],), - "transformation_value": ("FLOAT", {"default": 1.0}), - } - } - - RETURN_TYPES = ("LOCATION",) - FUNCTION = "transform_locations" - - def transform_locations(self, locations, feature, transformation_type, transformation_value): - transformed_locations = [] - for frame_index, location in enumerate(locations): - feature_value = feature.get_value_at_frame(frame_index) - if transformation_type == "translate": - new_x = location.x + feature_value * transformation_value - new_y = location.y + feature_value * transformation_value - if location.z is not None: - new_z = location.z + feature_value * transformation_value - transformed_location = Location(new_x, new_y, new_z) - else: - transformed_location = Location(new_x, new_y) - elif transformation_type == "scale": - new_x = location.x * feature_value * transformation_value - new_y = location.y * feature_value * transformation_value - if location.z is not None: - new_z = location.z * feature_value * transformation_value - transformed_location = Location(new_x, new_y, new_z) - else: - transformed_location = Location(new_x, new_y) - transformed_locations.append(transformed_location) - - return (transformed_locations,) - - - -class ProximityVisualizer(RyanOnTheInside): - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "video_frames": ("IMAGE",), - "anchor_locations": ("LOCATION",), - "query_locations": ("LOCATION",), - "proximity_feature": ("FEATURE",), - "anchor_color": ("STRING", {"default": "(255,0,0)"}), - "query_color": ("STRING", {"default": "(0,255,0)"}), - "line_color": ("STRING", {"default": "(0,0,255)"}), - "text_color": ("STRING", {"default": "(255,255,255)"}), - "font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}), - } - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "visualize_proximity" - - def visualize_proximity(self, video_frames, anchor_locations, query_locations, proximity_feature, - anchor_color, query_color, line_color, text_color, font_scale): - anchor_color = self.parse_color(anchor_color) - query_color = self.parse_color(query_color) - line_color = self.parse_color(line_color) - text_color = self.parse_color(text_color) - - output_frames = [] - height, width = video_frames.shape[1:3] # Extract height and width from video_frames - - # Calculate the frame diagonal from the video frames - frame_diagonal = np.sqrt(width**2 + height**2) - scale_factor = frame_diagonal / proximity_feature.frame_diagonal - for frame_index in range(len(video_frames)): - frame = video_frames[frame_index].cpu().numpy() - frame = (frame * 255).astype(np.uint8) - frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) - - anchor = anchor_locations[frame_index] - query = query_locations[frame_index] - - # Calculate scaling factors based on the frame diagonal - - - # Draw all anchor and query points (using only x and y coordinates) - for point in anchor: - scaled_point = (int(point[0] * scale_factor), int(point[1] * scale_factor)) - cv2.circle(frame, scaled_point, 2, anchor_color, -1) - for point in query: - scaled_point = (int(point[0] * scale_factor), int(point[1] * scale_factor)) - cv2.circle(frame, scaled_point, 2, query_color, -1) - - # Find the closest pair of points - if len(anchor) > 0 and len(query) > 0: - distances = cdist(anchor.points[:, :2], query.points[:, :2]) # Use only x and y for distance calculation - min_idx = np.unravel_index(distances.argmin(), distances.shape) - closest_anchor = anchor[min_idx[0]] - closest_query = query[min_idx[1]] - - # Draw line between closest points (using only x and y coordinates) - cv2.line(frame, - (int(closest_anchor[0] * scale_factor), int(closest_anchor[1] * scale_factor)), - (int(closest_query[0] * scale_factor), int(closest_query[1] * scale_factor)), - line_color, 2) - - # Display coordinates of closest points - anchor_text = f"Anchor: ({int(closest_anchor[0] * scale_factor)}, {int(closest_anchor[1] * scale_factor)}" - query_text = f"Query: ({int(closest_query[0] * scale_factor)}, {int(closest_query[1] * scale_factor)}" - - if anchor.z is not None and query.z is not None: - anchor_text += f", {closest_anchor[2]:.2f})" - query_text += f", {closest_query[2]:.2f})" - else: - anchor_text += ") [2D]" - query_text += ") [2D]" - - cv2.putText(frame, anchor_text, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, font_scale, text_color, 1) - cv2.putText(frame, query_text, (10, 90), cv2.FONT_HERSHEY_SIMPLEX, font_scale, text_color, 1) - - # Add proximity value to the frame - proximity_value = proximity_feature.get_value_at_frame(frame_index) - cv2.putText(frame, f"Proximity: {proximity_value:.2f}", (10, 30), - cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2) - - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - output_frames.append(frame) - - output_tensor = torch.from_numpy(np.stack(output_frames)).float() / 255.0 - return (output_tensor,) - - def parse_color(self, color_string): - return tuple(map(int, color_string.strip("()").split(","))) \ No newline at end of file diff --git a/nodes/images/image_base.py b/nodes/images/flex_image_base.py similarity index 100% rename from nodes/images/image_base.py rename to nodes/images/flex_image_base.py diff --git a/nodes/images/flex_images.py b/nodes/images/flex_images.py index 012ab7a..a3dfca4 100644 --- a/nodes/images/flex_images.py +++ b/nodes/images/flex_images.py @@ -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 diff --git a/nodes/latents/flex_latent_base.py b/nodes/latents/flex_latent_base.py new file mode 100644 index 0000000..5ed0905 --- /dev/null +++ b/nodes/latents/flex_latent_base.py @@ -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 + diff --git a/nodes/latents/latent_base.py b/nodes/latents/flex_latents.py similarity index 59% rename from nodes/latents/latent_base.py rename to nodes/latents/flex_latents.py index a7da4d9..ce74ec9 100644 --- a/nodes/latents/latent_base.py +++ b/nodes/latents/flex_latents.py @@ -1,94 +1,4 @@ -# custom_nodes/ComfyUI_RyanOnTheInside/nodes/latents/latent_base.py - -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): - @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 - - @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": - return param_value * (1 + (feature_value - 0.5) * strength) - else: # absolute - return param_value * feature_value * strength - - 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() - - 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.""" - -# custom_nodes/ComfyUI_RyanOnTheInside/nodes/latents/flex_latent_interpolate.py - +from .flex_latent_base import FlexLatentBase import numpy as np class FlexLatentInterpolate(FlexLatentBase): @@ -105,26 +15,14 @@ class FlexLatentInterpolate(FlexLatentBase): def get_modifiable_params(cls): return ["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'] 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()[kwargs['frame_index']] + latent_2_np = latent_2["samples"].cpu().numpy()[frame_index] # Perform interpolation t = np.clip(feature_value * strength, 0.0, 1.0) @@ -154,31 +52,6 @@ class FlexLatentInterpolate(FlexLatentBase): coef2 = np.sin(t * omega) / sin_omega result = coef1 * latent1 + coef2 * latent2 return result.reshape(latent1.shape) - -# custom_nodes/ComfyUI_RyanOnTheInside/nodes/latents/embedding_guided_latent_interpolate.py - - def spherical_interpolation(self, latent1, latent2, t): - # Flatten the latents - latent1_flat = latent1.flatten() - latent2_flat = latent2.flatten() - - # Compute dot product and norms - dot = np.dot(latent1_flat, latent2_flat) - norm1 = np.linalg.norm(latent1_flat) - norm2 = np.linalg.norm(latent2_flat) - denominator = norm1 * norm2 + 1e-8 # Avoid division by zero - omega = np.arccos(np.clip(dot / denominator, -1.0, 1.0)) - - if np.isclose(omega, 0): - return latent1 - else: - sin_omega = np.sin(omega) - coef1 = np.sin((1 - t) * omega) / sin_omega - coef2 = np.sin(t * omega) / sin_omega - result = coef1 * latent1 + coef2 * latent2 - return result.reshape(latent1.shape) - -import numpy as np class EmbeddingGuidedLatentInterpolate(FlexLatentBase): @classmethod @@ -196,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] @@ -261,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 @@ -278,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'] @@ -323,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): @@ -338,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'] @@ -355,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 @@ -369,4 +228,5 @@ class FlexLatentNoise(FlexLatentBase): # Add noise to the latent result = latent + noise - return result \ No newline at end of file + return result +