diff --git a/__init__.py b/__init__.py index 894f250..1a1e5a7 100644 --- a/__init__.py +++ b/__init__.py @@ -99,7 +99,9 @@ from .nodes.flex.feature_extractors_whisper import( ContextModifier, WhisperToPromptTravel, WhisperTextRenderer, - ManualWhisperAlignmentData + ManualWhisperAlignmentData, + WhisperAutoAdjust, + WhisperTimeAdjuster ) from .nodes.flex.feature_extractors_audio import( @@ -483,6 +485,8 @@ NODE_CLASS_MAPPINGS = { "TriggerBuilder": TriggerBuilder, "ContextModifier": ContextModifier, "WhisperTextRenderer": WhisperTextRenderer, + "WhisperAutoAdjust": WhisperAutoAdjust, + "WhisperTimeAdjuster": WhisperTimeAdjuster, "ManualWhisperAlignmentData": ManualWhisperAlignmentData, "FeatureToWeightsStrategy": FeatureToWeightsStrategy, diff --git a/nodes/flex/feature_extractors_whisper.py b/nodes/flex/feature_extractors_whisper.py index 37ab126..c327590 100644 --- a/nodes/flex/feature_extractors_whisper.py +++ b/nodes/flex/feature_extractors_whisper.py @@ -24,6 +24,58 @@ from ... import RyanOnTheInside _category = f"{FeatureExtractorBase.CATEGORY}/Whisper" +class TriggerSet: + def __init__(self): + self.triggers = [] + self._expected_image_shape = None # Store the expected shape for validation + + def add_trigger(self, pattern: str, values: tuple, mode: str, fade: str, + duration: int, blend_mode: str, fill_behavior: str, image=None): + # Validate image shape if provided + if image is not None: + current_shape = tuple(image.shape[1:]) # HWC shape (ignore batch dimension) + + if self._expected_image_shape is None: + # First image sets the expected shape + self._expected_image_shape = current_shape + elif current_shape != self._expected_image_shape: + raise ValueError( + f"Image shape mismatch! Expected shape {self._expected_image_shape}, " + f"but got {current_shape}. All trigger images must have the same dimensions." + ) + + trigger = { + "pattern": pattern, + "values": values, + "mode": mode, + "fade": fade, + "duration": duration, + "blend_mode": blend_mode, + "fill_behavior": fill_behavior, + "image": image, + "image_batch_size": image.shape[0] if image is not None else 1 + } + self.triggers.append(trigger) + return self + + def extend(self, other_trigger_set): + if other_trigger_set: + # Validate image shapes from other trigger set + for trigger in other_trigger_set.triggers: + if trigger.get("image") is not None: + current_shape = tuple(trigger["image"].shape[1:]) + if self._expected_image_shape is None: + self._expected_image_shape = current_shape + elif current_shape != self._expected_image_shape: + raise ValueError( + f"Image shape mismatch in trigger set combination! " + f"Expected shape {self._expected_image_shape}, but got {current_shape}. " + f"All trigger images must have the same dimensions." + ) + + self.triggers.extend(other_trigger_set.triggers) + return self + @apply_tooltips class WhisperFeatureNode(FeatureExtractorBase): @classmethod @@ -33,7 +85,6 @@ class WhisperFeatureNode(FeatureExtractorBase): @classmethod def INPUT_TYPES(cls): parent_inputs = super().INPUT_TYPES()["required"] - parent_inputs.pop("frame_count", None) # Remove frame_count as we'll calculate it return { **super().INPUT_TYPES(), "required": { @@ -47,20 +98,16 @@ class WhisperFeatureNode(FeatureExtractorBase): } } - RETURN_TYPES = ("FEATURE",) + RETURN_TYPES = ("FEATURE", "IMAGE") FUNCTION = "create_feature" CATEGORY = _category - def create_feature(self, frame_rate, width, height, extraction_method, + def create_feature(self, frame_rate, frame_count, width, height, extraction_method, alignment_data, trigger_set=None, context_size=3, overlap_mode="blend"): - # Parse alignment data to find the last timestamp + # Parse alignment data try: - # Handle both string and dict/list alignment data data = alignment_data if not isinstance(alignment_data, str) else json.loads(alignment_data) - last_end_time = max(segment["end"] for segment in data) - # Calculate frame_count from audio duration and frame_rate - frame_count = int(last_end_time * frame_rate) + 1 except (json.JSONDecodeError, KeyError, TypeError) as e: raise ValueError(f"Invalid alignment data format: {e}") @@ -72,20 +119,103 @@ class WhisperFeatureNode(FeatureExtractorBase): except json.JSONDecodeError: print("Warning: Could not parse trigger_set JSON") - # Create the WhisperFeature + # Create the WhisperFeature with the requested extraction method whisper_feature = WhisperFeature( name="whisper_feature", frame_rate=frame_rate, frame_count=frame_count, alignment_data=alignment_data, - trigger_pairs=triggers, + trigger_pairs=trigger_set, feature_name=extraction_method, width=width, height=height ) whisper_feature.extract() - return (whisper_feature,) + + # Process images if triggers contain them + output_images = None + if trigger_set and any(t.get("image") is not None for t in trigger_set.triggers): + try: + print(f"Processing {len(trigger_set.triggers)} triggers with images") + image_sequence = [] + current_frame = 0 + + sorted_triggers = whisper_feature.sort_triggers_by_occurrence(trigger_set.triggers) + print(f"Sorted triggers: {len(sorted_triggers)}") + + # Instead of processing trigger by trigger, let's collect all occurrences first + all_occurrences = [] + for trigger in sorted_triggers: + if trigger.get("image") is not None: + frame_ranges = whisper_feature.find_all_trigger_frames(trigger["pattern"]) + for start_frame, end_frame in frame_ranges: + all_occurrences.append({ + "trigger": trigger, + "start_frame": start_frame, + "end_frame": end_frame + }) + + # Sort all occurrences by start frame to maintain chronological order + all_occurrences.sort(key=lambda x: x["start_frame"]) + print(f"Processing {len(all_occurrences)} total trigger occurrences in chronological order") + + # Process occurrences in order + for occurrence in all_occurrences: + trigger = occurrence["trigger"] + start_frame = occurrence["start_frame"] + end_frame = occurrence["end_frame"] + + # Convert back to tensor + images = torch.tensor(trigger["image"]) if isinstance(trigger["image"], list) else trigger["image"] + batch_size = trigger["image_batch_size"] + frame_count = end_frame - start_frame + + # Handle batch images based on fill behavior + if trigger["fill_behavior"] == "none": + frame_count = 1 + images = images[:1] + else: + # Handle fill behaviors + if batch_size > 1: + if trigger["fill_behavior"] == "loop": + repeats = (frame_count + batch_size - 1) // batch_size + images = images.repeat((repeats, 1, 1, 1))[:frame_count] + elif trigger["fill_behavior"] == "hold": + if frame_count > batch_size: + last_frame = images[-1:] + extra_frames = frame_count - batch_size + images = torch.cat([images, last_frame.repeat((extra_frames, 1, 1, 1))]) + else: + if frame_count > 1: + images = images.repeat((frame_count, 1, 1, 1)) + + if trigger["fill_behavior"] != "none": + # Only fill gaps if not "none" + if current_frame < start_frame: + zero_frames = start_frame - current_frame + zero_image = torch.zeros_like(images[0]) + image_sequence.extend([zero_image] * zero_frames) + + image_sequence.extend(images.unbind(0)) + current_frame = end_frame if trigger["fill_behavior"] != "none" else start_frame + 1 + + if image_sequence: + print(f"Final sequence length: {len(image_sequence)}") + output_images = torch.stack(image_sequence) + else: + print("No images in sequence") + + except Exception as e: + print(f"Error processing trigger images: {e}") + import traceback + traceback.print_exc() + output_images = None + + # Set feature back to requested extraction method + whisper_feature.set_active_feature(extraction_method) + + return (whisper_feature, output_images) if output_images is not None else (whisper_feature, None) @apply_tooltips class TriggerBuilder(RyanOnTheInside): @@ -107,36 +237,18 @@ class TriggerBuilder(RyanOnTheInside): def INPUT_TYPES(cls): return { "required": { - # The word or phrase to match in the transcription "pattern": ("STRING", {"multiline": False, "default": "hello"}), - - # Value when the word starts (0-1) "start_value": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}), - - # Value when the word ends (0-1) "end_value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}), - - # How to match the pattern: - # - exact: Must match whole word exactly - # - contains: Pattern appears anywhere in word - # - regex: Use regular expressions - # - phonetic: Match similar sounding words "match_mode": (["exact", "contains", "regex", "phonetic"], {"default": "contains"}), - - # How to transition between values: - # - none: Instant change - # - linear: Smooth linear fade - # - smooth: Smooth eased transition "fade_type": (["none", "linear", "smooth"], {"default": "linear"}), - - # Duration in frames (0 = use word duration) "duration_frames": ("INT", {"default": 0, "min": 0}), - - # How this trigger blends with previous accumulated results "blend_mode": (["blend", "add", "multiply", "max"], {"default": "blend"}), + "fill_behavior": (["none", "hold", "loop"], {"default": "none"}), }, "optional": { "previous_triggers": ("TRIGGER_SET",), + "trigger_image": ("IMAGE",), } } @@ -146,29 +258,27 @@ class TriggerBuilder(RyanOnTheInside): def build(self, pattern: str, start_value: float, end_value: float, match_mode: str, fade_type: str, duration_frames: int, - blend_mode: str, previous_triggers=None) -> tuple: + blend_mode: str, fill_behavior: str, previous_triggers=None, + trigger_image=None) -> tuple: - # Create current trigger with its blend mode - current_trigger = { - "pattern": pattern, - "values": (start_value, end_value), - "mode": match_mode, - "fade": fade_type, - "duration": duration_frames, - "blend_mode": blend_mode # Store blend mode per trigger - } + trigger_set = TriggerSet() - # Initialize trigger set - trigger_set = { - "triggers": [current_trigger] - } + # Add new trigger + trigger_set.add_trigger( + pattern=pattern, + values=(start_value, end_value), + mode=match_mode, + fade=fade_type, + duration=duration_frames, + blend_mode=blend_mode, + fill_behavior=fill_behavior, + image=trigger_image + ) - # If we have previous triggers, combine them - if previous_triggers: - prev_data = json.loads(previous_triggers) - trigger_set["triggers"] = prev_data["triggers"] + trigger_set["triggers"] + # Combine with previous triggers if any + trigger_set.extend(previous_triggers) - return (json.dumps(trigger_set),) + return (trigger_set,) @apply_tooltips class ContextModifier(RyanOnTheInside): @@ -496,6 +606,153 @@ class ManualWhisperAlignmentData: def parse(self, alignment_text: str) -> tuple: return (alignment_text,) +@apply_tooltips +class WhisperTimeAdjuster: + """Manually adjust timing in Whisper alignment data""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "alignment_data": ("whisper_alignment",), + "time_offset": ("FLOAT", { + "default": 0.0, + "min": -1000.0, + "max": 1000.0, + "step": 0.1, + "description": "Seconds to shift all timestamps (positive = delay, negative = earlier)" + }), + } + } + + RETURN_TYPES = ("whisper_alignment",) + FUNCTION = "adjust" + CATEGORY = _category + + def adjust(self, alignment_data, time_offset): + try: + data = alignment_data if not isinstance(alignment_data, str) else json.loads(alignment_data) + + # Adjust all timestamps by offset + adjusted_data = [] + for segment in data: + adjusted_segment = segment.copy() + adjusted_segment["start"] = segment["start"] + time_offset + adjusted_segment["end"] = segment["end"] + time_offset + adjusted_data.append(adjusted_segment) + + return (adjusted_data,) + + except (json.JSONDecodeError, KeyError) as e: + raise ValueError(f"Invalid alignment data format: {e}") + +@apply_tooltips +class WhisperAutoAdjust: + """Automatically adjusts Whisper timing by detecting meaningful audio start""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "alignment_data": ("whisper_alignment",), + "audio": ("AUDIO",), # Video Helper Suite audio format + "detection_window": ("INT", { + "default": 1024, + "min": 256, + "max": 8192, + "step": 256, + "description": "Window size for energy detection" + }), + "energy_threshold": ("FLOAT", { + "default": 0.0001, + "min": 0.00001, + "max": 0.001, + "step": 0.00001, + "description": "Energy threshold (lower = more sensitive)" + }) + } + } + + RETURN_TYPES = ("whisper_alignment",) + FUNCTION = "adjust" + CATEGORY = _category + + def detect_audio_start(self, audio_dict, window_size, threshold): + # Extract waveform and convert to mono numpy array, like BaseAudioProcessor + audio = audio_dict['waveform'].squeeze(0).mean(axis=0).cpu().numpy() + sample_rate = audio_dict['sample_rate'] + + print(f"Audio shape before processing: {audio_dict['waveform'].shape}") + print(f"Audio shape after mono conversion: {audio.shape}") + print(f"Sample rate: {sample_rate}") + + # Calculate energy in windows using numpy for efficiency + num_windows = (len(audio) - window_size) // (window_size // 2) + energies = np.zeros(num_windows) + + for i in range(num_windows): + start = i * (window_size // 2) + end = start + window_size + window = audio[start:end] + energies[i] = np.mean(window ** 2) + + print(f"Energy shape: {energies.shape}") + print(f"Max energy: {energies.max():.6f}") + print(f"Min energy: {energies.min():.6f}") + + # Find first window above threshold + start_window = np.where(energies > threshold)[0] + if len(start_window) == 0: + print(f"No windows above threshold {threshold}") + return 0 + + # Convert window index to sample + start_sample = start_window[0] * (window_size // 2) + start_time = start_sample / sample_rate + + print(f"Found start time: {start_time:.3f}s") + return start_time + + def adjust(self, alignment_data, audio, detection_window=1024, energy_threshold=0.0001): + try: + data = alignment_data if not isinstance(alignment_data, str) else json.loads(alignment_data) + + print("Original whisper timings:") + for segment in data[:3]: # Print first 3 segments + print(f" {segment['start']:.3f} -> {segment['end']:.3f}: {segment['value']}") + + # Detect meaningful audio start + audio_start = self.detect_audio_start(audio, detection_window, energy_threshold) + print(f"\nDetected audio start: {audio_start}") + + # Get first whisper timestamp + whisper_start = min(segment["start"] for segment in data) + print(f"First whisper timestamp: {whisper_start}") + + # Calculate needed offset + time_offset = audio_start - whisper_start + print(f"Calculated offset: {time_offset}") + + # Adjust all timestamps + adjusted_data = [] + for segment in data: + adjusted_segment = segment.copy() + adjusted_segment["start"] = segment["start"] + time_offset + adjusted_segment["end"] = segment["end"] + time_offset + adjusted_data.append(adjusted_segment) + + print("\nAdjusted whisper timings:") + for segment in adjusted_data[:3]: # Print first 3 segments + print(f" {segment['start']:.3f} -> {segment['end']:.3f}: {segment['value']}") + + return (adjusted_data,) + + except Exception as e: + print(f"Error during adjustment: {str(e)}") + import traceback + traceback.print_exc() + raise ValueError(f"Error during alignment: {str(e)}") + diff --git a/nodes/flex/features.py b/nodes/flex/features.py index 1ebde6a..3fc75a4 100644 --- a/nodes/flex/features.py +++ b/nodes/flex/features.py @@ -637,9 +637,8 @@ class WhisperFeature(BaseFeature): except json.JSONDecodeError: raise ValueError("Invalid alignment data format: must be valid JSON string or dict/list") else: - # If it's already a dict/list, use it directly self.alignment_data = alignment_data - self.trigger_pairs = trigger_pairs # Already a dictionary from TriggerBuilder + self.trigger_pairs = trigger_pairs # Now a TriggerSet instance self.feature_name = feature_name self.available_features = self.get_extraction_methods() @@ -699,11 +698,11 @@ class WhisperFeature(BaseFeature): start_frame = int(item["start"] * self.frame_rate) end_frame = int(item["end"] * self.frame_rate) - for trigger in self.trigger_pairs["triggers"]: + for trigger in self.trigger_pairs.triggers: # Access triggers directly from TriggerSet pattern = trigger["pattern"].lower() if pattern in text: start_val, end_val = trigger["values"] - blend_mode = trigger.get("blend_mode", "blend") # Default to blend if not specified + blend_mode = trigger.get("blend_mode", "blend") # Create temporary array for this trigger's values trigger_data = [0.0] * self.frame_count @@ -772,5 +771,52 @@ class WhisperFeature(BaseFeature): else: raise ValueError(f"Invalid feature name. Available features are: {', '.join(self.available_features)}") + def find_trigger_start_time(self, pattern): + """Find the start time of when a pattern appears in the alignment data""" + for segment in self.alignment_data: + if pattern.lower() in segment["value"].lower(): + return segment["start"] + return None + + def find_trigger_end_time(self, pattern): + """Find the end time of when a pattern appears in the alignment data""" + for segment in self.alignment_data: + if pattern.lower() in segment["value"].lower(): + return segment["end"] + return None + + def sort_triggers_by_occurrence(self, triggers): + """Sort triggers based on when their patterns appear in the text""" + def find_first_occurrence(pattern): + text = " ".join(segment["value"] for segment in self.alignment_data) + return text.lower().find(pattern.lower()) + + return sorted(triggers, key=lambda t: find_first_occurrence(t["pattern"])) + + def get_trigger_frames(self, pattern): + """Get the frame range where a trigger pattern occurs""" + start_time = self.find_trigger_start_time(pattern) + if start_time is None: + return None, None + + end_time = self.find_trigger_end_time(pattern) + start_frame = int(start_time * self.frame_rate) + end_frame = int(end_time * self.frame_rate) + + return start_frame, end_frame + + def find_all_trigger_frames(self, pattern): + """Find all frame ranges where a trigger pattern occurs""" + frame_ranges = [] + pattern = pattern.lower() + + for segment in self.alignment_data: + if pattern in segment["value"].lower(): + start_frame = int(segment["start"] * self.frame_rate) + end_frame = int(segment["end"] * self.frame_rate) + frame_ranges.append((start_frame, end_frame)) + + return frame_ranges + #TODO volume feature