From 7edd336cda17e91420b16773ad062343cb9e3ef1 Mon Sep 17 00:00:00 2001
From: DoctorDiffusion <139909432+DoctorDiffusion@users.noreply.github.com>
Date: Mon, 11 Nov 2024 15:55:21 -0500
Subject: [PATCH] Add files via upload

---
 __init__.py                  | 14 ++++++++
 prompt_schedule_converter.py | 63 ++++++++++++++++++++++++++++++++++++
 requirements.txt             |  3 ++
 whisper_node.py              | 41 +++++++++++++++++++++++
 4 files changed, 121 insertions(+)
 create mode 100644 __init__.py
 create mode 100644 prompt_schedule_converter.py
 create mode 100644 requirements.txt
 create mode 100644 whisper_node.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..0d3487b
--- /dev/null
+++ b/__init__.py
@@ -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']
\ No newline at end of file
diff --git a/prompt_schedule_converter.py b/prompt_schedule_converter.py
new file mode 100644
index 0000000..9a17fcd
--- /dev/null
+++ b/prompt_schedule_converter.py
@@ -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)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..f62783b
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,3 @@
+openai-whisper
+pillow
+uuid
\ No newline at end of file
diff --git a/whisper_node.py b/whisper_node.py
new file mode 100644
index 0000000..02a01da
--- /dev/null
+++ b/whisper_node.py
@@ -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