Add files via upload

This commit is contained in:
DoctorDiffusion
2024-11-11 15:55:21 -05:00
committed by GitHub
parent e07d186306
commit 7edd336cda
4 changed files with 121 additions and 0 deletions
+14
View File
@@ -0,0 +1,14 @@
from .whisper_node import WhisperNode
from .prompt_schedule_converter import PromptScheduleConverter
NODE_CLASS_MAPPINGS = {
"whisper_node" : WhisperNode,
"prompt_schedule_converter" : PromptScheduleConverter
}
NODE_DISPLAY_NAME_MAPPINGS = {
"whisper_node" : "Whisper Node",
"prompt_schedule_converter" : "Prompt Schedule Converter"
}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+63
View File
@@ -0,0 +1,63 @@
import re
class PromptScheduleConverter:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"text_with_timestamps": ("STRING",), # Subtitle text in custom timestamp format
"fps": ("FLOAT",), # Frames per second
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("formatted_text",)
FUNCTION = "convert_time_to_frame"
CATEGORY = "Schedulizer"
def convert_time_to_frame(self, text_with_timestamps, fps):
# Define the pattern for extracting time ranges and text
text_pattern = r'\[(\d+\.\d+)s - (\d+\.\d+)s\] (.+)' # Capture non-empty text after time range
# Find all matches for start time, end time, and associated text
matches = re.findall(text_pattern, text_with_timestamps)
# Check if any matches were found
if not matches:
print("No valid matches found in input.")
return {"ui": {"text": "No valid input detected."}, "result": ("No valid input detected.",)}
formatted_text = []
last_text = None # Track the last text to avoid duplicates
first_entry = True # Track if we're adding the first entry
# Loop over each matched time range and text
for start_time, end_time, line in matches:
# Convert start time to frame number
start_frame = self.time_to_frames(start_time, fps)
# Ensure the text is stripped of excess whitespace
line = line.strip()
# Check if this is the first entry
if first_entry:
# Always set the first entry to frame "0"
formatted_text.append(f'"0" : "{line}"')
last_text = line # Update last_text to the current line
first_entry = False # Disable the first entry flag
continue # Skip to the next loop iteration
# Only add the frame-text pair if the text has changed
if line != last_text:
formatted_text.append(f'"{start_frame}" : "{line}"')
last_text = line # Update last_text to the current line
# Join all frame-text pairs with line breaks, adding a comma after each except the last
formatted_output = ",\n".join(formatted_text)
# Return formatted output
return {"ui": {"text": formatted_output}, "result": (formatted_output,)}
def time_to_frames(self, time_str, fps):
# Convert time string (ss.s) to frame count
return int(float(time_str) * fps)
+3
View File
@@ -0,0 +1,3 @@
openai-whisper
pillow
uuid
+41
View File
@@ -0,0 +1,41 @@
import whisper
import os
import folder_paths
import uuid
import torchaudio
class WhisperNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"audio": ("AUDIO",),
"model": (["base", "tiny", "small", "medium", "large", "large-v2"],),
}
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("text", "text_with_timestamps")
FUNCTION = "whisper_node"
CATEGORY = "Schedulizer"
def whisper_node(self, audio, model):
# Save audio bytes to file
temp_dir = folder_paths.get_temp_directory()
os.makedirs(temp_dir, exist_ok=True)
audio_save_path = os.path.join(temp_dir, f"{uuid.uuid1()}.wav")
torchaudio.save(audio_save_path, audio['waveform'].squeeze(0), audio["sample_rate"])
# Transcribe using Whisper
model = whisper.load_model(model)
result = model.transcribe(audio_save_path, word_timestamps=True)
# Extract plain text and timestamped text
plain_text = result["text"].strip()
text_with_timestamps = "\n".join(
f"[{segment['start']:.2f}s - {segment['end']:.2f}s] {segment['text'].strip()}"
for segment in result['segments']
)
return plain_text, text_with_timestamps