Add files via upload
This commit is contained in:
+14
@@ -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']
|
||||
@@ -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)
|
||||
@@ -0,0 +1,3 @@
|
||||
openai-whisper
|
||||
pillow
|
||||
uuid
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user