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

whisper with alignment

Browse Source
This commit is contained in:
RyanOnTheInside
2025-01-12 22:37:03 +00:00
parent b223e7b875
commit 1db4fef68f
3 changed files with 362 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+5 -1
View File
@@ -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,
nodes/flex/feature_extractors_whisper.py
+307 -50
View File
@@ -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)}")
nodes/flex/features.py
+50 -4
View File
@@ -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