Update app.py

app.py

@@ -21,7 +21,6 @@ from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageEnhance
 from datetime import datetime
 from threading import Lock
 import base64
-import json
 import io
 
 # --- Dependency Check ---
@@ -790,39 +789,82 @@ def remove_video_background(video_path):
     shutil.rmtree(job_temp_dir)
     return output_video_path
 
+def generate_ass_from_whisper(result):
+    """Generates an ASS subtitle file content from a Whisper result object with word timestamps."""
+    ass_content = [
+        "[Script Info]",
+        "Title: Generated by Skriptz",
+        "ScriptType: v4.00+",
+        "WrapStyle: 0",
+        "PlayResX: 1920",
+        "PlayResY: 1080",
+        "\n[V4+ Styles]",
+        "Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding",
+        "Style: Default,Arial,55,&H00FFFFFF,&H0000FFFF,&H00000000,&H00000000,0,0,0,0,100,100,0,0,1,2,1,2,10,10,25,1",
+        "\n[Events]",
+        "Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text"
+    ]
+
+    def format_time(s):
+        h, r = divmod(s, 3600)
+        m, s = divmod(r, 60)
+        cs = int((s - int(s)) * 100)
+        return f"{int(h)}:{int(m):02}:{int(s):02}.{cs:02}"
+
+    for segment in result['segments']:
+        start_time = format_time(segment['start'])
+        end_time = format_time(segment['end'])
+        
+        karaoke_line = ""
+        for word_info in segment['words']:
+            word = word_info['word'].strip()
+            duration_cs = int((word_info['end'] - word_info['start']) * 100)
+            karaoke_line += f"{{\\k{duration_cs}}}{word} "
+            
+        dialogue_line = f"Dialogue: 0,{start_time},{end_time},Default,,0,0,0,,{karaoke_line.strip()}"
+        ass_content.append(dialogue_line)
+
+    return "\n".join(ass_content)
+
 def transcribe_media(media_path, model_name):
     if media_path is None: raise gr.Error("Please upload a video or audio file first.")
     model = load_whisper_model(model_name)
     if model is None: raise gr.Error("Whisper model is not available.")
+    
     audio_path = media_path.name
     base_name = os.path.splitext(os.path.basename(media_path.name))[0]
     
     if audio_path.lower().endswith(('.mp4', '.mov', '.mkv', '.avi', '.webm')):
         audio_path_temp = os.path.join(TEMP_DIR, f"{base_name}.mp3")
         try:
-            run_ffmpeg_command(["ffmpeg", "-i", audio_path, "-q:a", "0", "-map", "a", "-y", audio_path_temp])
+            run_ffmpeg_command(["ffmpeg", "-y", "-i", audio_path, "-q:a", "0", "-map", "a", audio_path_temp])
             audio_path = audio_path_temp
         except gr.Error as e:
             if "does not contain any stream" in str(e): raise gr.Error("The uploaded video has no audio track.")
             else: raise e
     
-    result = model.transcribe(audio_path, verbose=False) 
+    result = model.transcribe(audio_path, word_timestamps=True, verbose=False)
     
-    def format_ts(s):
+    def format_ts(s, separator=','):
         h, r = divmod(s, 3600); m, s = divmod(r, 60)
-        return f"{int(h):02}:{int(m):02}:{int(s):02},{int((s-int(s))*1000):03}"
+        return f"{int(h):02}:{int(m):02}:{int(s):02}{separator}{int((s-int(s))*1000):03}"
         
     srt_path = os.path.join(TEMP_DIR, f"{base_name}.srt")
     vtt_path = os.path.join(TEMP_DIR, f"{base_name}.vtt")
+    ass_path = os.path.join(TEMP_DIR, f"{base_name}.ass")
     
     with open(srt_path, "w", encoding="utf-8") as srt_f, open(vtt_path, "w", encoding="utf-8") as vtt_f:
         vtt_f.write("WEBVTT\n\n")
         for i, seg in enumerate(result["segments"]):
             start, end, text = seg['start'], seg['end'], seg['text'].strip()
             srt_f.write(f"{i + 1}\n{format_ts(start)} --> {format_ts(end)}\n{text}\n\n")
-            vtt_f.write(f"{format_ts(start).replace(',', '.')} --> {format_ts(end).replace(',', '.')}\n{text}\n\n")
+            vtt_f.write(f"{format_ts(start, '.')} --> {format_ts(end, '.')}\n{text}\n\n")
             
-    return result["text"], [srt_path, vtt_path]
+    ass_content = generate_ass_from_whisper(result)
+    with open(ass_path, "w", encoding="utf-8") as ass_f:
+        ass_f.write(ass_content)
+
+    return result["text"], [srt_path, vtt_path, ass_path]
 
 def transcribe_and_prep_burn(media_file, model_name):
     if not media_file: raise gr.Error("Please upload a file first.")
@@ -832,10 +874,6 @@ def transcribe_and_prep_burn(media_file, model_name):
     else: return text, files, None, gr.update(visible=False)
 
 def reformat_srt_for_word_wrap(original_srt_path, words_per_line):
-    """
-    Reads an SRT file and adds line breaks to the text blocks without changing timestamps,
-    then writes the result to a new temporary SRT file.
-    """
     if not original_srt_path or not os.path.exists(original_srt_path):
         return None
 
@@ -844,20 +882,14 @@ def reformat_srt_for_word_wrap(original_srt_path, words_per_line):
     
     with open(original_srt_path, 'r', encoding='utf-8') as f_in, \
          open(reformatted_path, 'w', encoding='utf-8') as f_out:
-        
-        content = f_in.read()
-        blocks = content.strip().split('\n\n')
-        
-        for block in blocks:
+        content = f_in.read().strip().split('\n\n')
+        for block in content:
             lines = block.split('\n')
-            if len(lines) < 3: # Not a valid block (e.g., just a number and timestamp)
+            if len(lines) < 3:
                 f_out.write(block + '\n\n')
                 continue
-
-            text_lines = lines[2:]
-            full_text = ' '.join(text_lines)
-            words = full_text.split()
-            
+            text_lines = ' '.join(lines[2:])
+            words = text_lines.split()
             new_text_lines = []
             current_line = []
             for word in words:
@@ -865,29 +897,16 @@ def reformat_srt_for_word_wrap(original_srt_path, words_per_line):
                 if len(current_line) >= words_per_line:
                     new_text_lines.append(' '.join(current_line))
                     current_line = []
-            
-            if current_line: # Add any remaining words
-                new_text_lines.append(' '.join(current_line))
-                
-            # Reconstruct the block with the new text formatting
-            f_out.write(lines[0] + '\n') # Index
-            f_out.write(lines[1] + '\n') # Timestamp
-            f_out.write('\n'.join(new_text_lines) + '\n\n')
-
+            if current_line: new_text_lines.append(' '.join(current_line))
+            f_out.write(f"{lines[0]}\n{lines[1]}\n{'\n'.join(new_text_lines)}\n\n")
     return reformatted_path
 
-def burn_subtitles(video_path, srt_file_obj, font_size_scale, font_color, words_per_line):
-    if not video_path: raise gr.Error("Original video path not found. Please re-transcribe.")
-    if not srt_file_obj or not srt_file_obj[0].name: raise gr.Error("SRT file not found. Please re-transcribe.")
-    
+def burn_block_subtitles(video_path, srt_file_obj, font_size_scale, font_color, words_per_line):
     original_srt_path = srt_file_obj[0].name
     reformatted_srt_path = None
-    
     try:
-        # Reformat the SRT file based on the slider value
         reformatted_srt_path = reformat_srt_for_word_wrap(original_srt_path, words_per_line)
-        if not reformatted_srt_path:
-            raise gr.Error("Failed to reformat subtitle file.")
+        if not reformatted_srt_path: raise gr.Error("Failed to reformat subtitle file.")
         
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         output_video_path = os.path.join(TEMP_DIR, f"subtitled_video_{timestamp}.mp4")
@@ -897,21 +916,66 @@ def burn_subtitles(video_path, srt_file_obj, font_size_scale, font_color, words_
         calculated_font_size = int(video_h / divisor)
         color_bgr = font_color[5:7] + font_color[3:5] + font_color[1:3]
         ffmpeg_color = f"&H00{color_bgr.upper()}"
+        escaped_srt_path = reformatted_srt_path.replace('\\', '/').replace(':', r'\\:')
         
-        # Use the new, reformatted path and escape it for FFMPEG
-        escaped_srt_path = reformatted_srt_path.replace('\\', '/').replace(':', r'\:')
-        
-        vf_filter = f"subtitles='{escaped_srt_path}':force_style='Fontsize={calculated_font_size},PrimaryColour={ffmpeg_color},BorderStyle=1,Outline=1,Shadow=0.5,MarginV=15'"
-        cmd = ["ffmpeg", "-i", video_path, "-vf", vf_filter, "-c:a", "copy", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-y", output_video_path]
-        run_ffmpeg_command(cmd, "Burning subtitles into video...")
-        
+        vf_filter = f"subtitles='{escaped_srt_path}':force_style='Fontsize={calculated_font_size},PrimaryColour={ffmpeg_color},BorderStyle=1,Outline=1,Shadow=0.5,MarginV=25'"
+        cmd = ["ffmpeg", "-y", "-i", video_path, "-vf", vf_filter, "-c:a", "copy", "-c:v", "libx264", "-pix_fmt", "yuv420p", output_video_path]
+        run_ffmpeg_command(cmd, "Burning block subtitles into video...")
         return output_video_path
-        
     finally:
-        # Clean up the temporary reformatted file
         if reformatted_srt_path and os.path.exists(reformatted_srt_path):
             os.remove(reformatted_srt_path)
 
+def burn_karaoke_subtitles(video_path, subtitle_files, font_size_scale, base_color, highlight_color):
+    ass_file_path = subtitle_files[2].name
+    temp_ass_path = None
+    try:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        temp_ass_path = os.path.join(TEMP_DIR, f"style_applied_{timestamp}.ass")
+        output_video_path = os.path.join(TEMP_DIR, f"karaoke_video_{timestamp}.mp4")
+
+        _, video_h = get_video_dimensions(video_path)
+        if video_h == 0: video_h = 720
+        calculated_font_size = int((video_h / 20) * (font_size_scale / 5))
+
+        def format_ass_color(hex_color):
+            if hex_color.startswith('#'): hex_color = hex_color[1:]
+            r, g, b = hex_color[0:2], hex_color[2:4], hex_color[4:6]
+            return f"&H00{b.upper()}{g.upper()}{r.upper()}"
+
+        primary_color_ass = format_ass_color(highlight_color)
+        secondary_color_ass = format_ass_color(base_color)
+
+        with open(ass_file_path, 'r', encoding='utf-8') as f_in, open(temp_ass_path, 'w', encoding='utf-8') as f_out:
+            for line in f_in:
+                if line.startswith("Style:"):
+                    parts = line.split(',')
+                    parts[2] = str(calculated_font_size) # Fontsize
+                    parts[3] = secondary_color_ass # PrimaryColour (Base text)
+                    parts[4] = primary_color_ass   # SecondaryColour (Karaoke fill)
+                    f_out.write(','.join(parts))
+                else:
+                    f_out.write(line)
+        
+        escaped_ass_path = temp_ass_path.replace('\\', '/').replace(':', r'\\:')
+        vf_filter = f"subtitles='{escaped_ass_path}'"
+        cmd = ["ffmpeg", "-y", "-i", video_path, "-vf", vf_filter, "-c:a", "copy", "-c:v", "libx264", "-pix_fmt", "yuv420p", output_video_path]
+        run_ffmpeg_command(cmd, "Burning karaoke subtitles into video...")
+        return output_video_path
+    finally:
+        if temp_ass_path and os.path.exists(temp_ass_path):
+            os.remove(temp_ass_path)
+
+def burn_subtitles_wrapper(video_path, subtitle_files, style, font_size_scale, block_font_color, block_words_per_line, kara_base_color, kara_highlight_color):
+    if not video_path or not subtitle_files: raise gr.Error("Missing video or subtitle files. Please transcribe first.")
+    
+    if style == "Block":
+        return burn_block_subtitles(video_path, subtitle_files, font_size_scale, block_font_color, block_words_per_line)
+    elif style == "Karaoke":
+        return burn_karaoke_subtitles(video_path, subtitle_files, font_size_scale, kara_base_color, kara_highlight_color)
+    else:
+        raise gr.Error("Invalid subtitle style selected.")
+
 def remove_background_single(input_path, output_path, **kwargs):
     with Image.open(input_path) as img:
         remove(img).save(output_path)
@@ -1546,17 +1610,23 @@ def update_clip_properties(timeline_state, selected_index, new_duration):
 def handle_timeline_action(timeline_state, selected_index, action):
     if selected_index == -1 or not (0 <= selected_index < len(timeline_state)):
         gr.Warning("Please select a clip from the timeline first.")
-        return timeline_state
+        return timeline_state, gr.update()
     
     new_list = list(timeline_state)
+    new_index = selected_index
+
     if action == "up" and selected_index > 0:
         new_list.insert(selected_index - 1, new_list.pop(selected_index))
+        new_index = selected_index - 1
     elif action == "down" and selected_index < len(new_list) - 1:
         new_list.insert(selected_index + 1, new_list.pop(selected_index))
+        new_index = selected_index + 1
     elif action == "remove":
         new_list.pop(selected_index)
+        new_index = -1 # Deselect after removing
         
-    return new_list
+    # Return the new list and tell the UI to select the new index
+    return new_list, gr.update(selected_index=new_index if new_index != -1 else None)
 
 def set_resolution_from_first_asset(timeline_state):
     if not timeline_state:
@@ -1606,14 +1676,15 @@ def create_animatic(timeline_data, audio_path, out_w, out_h, keep_original_audio
         output_clip_path = os.path.join(job_temp_dir, f"clip_{i:04d}.mp4")
         vf_filter = f"scale={output_res_str}:force_original_aspect_ratio=decrease,pad={output_res_str}:(ow-iw)/2:(oh-ih)/2,setsar=1"
 
-        cmd = ["ffmpeg"]
+        cmd = ["ffmpeg", "-y"]
         
         # --- Input Handling ---
-        # For images, we might need a silent audio source for concat to work
-        if item_type == "image" and keep_original_audio:
-            cmd.extend(["-f", "lavfi", "-i", "anullsrc=channel_layout=stereo:sample_rate=44100"])
+        if item_type == "image":
+            cmd.extend(["-loop", "1"])
+            # For images, we need a silent audio source if we're keeping audio in other clips
+            if keep_original_audio:
+                cmd.extend(["-f", "lavfi", "-i", "anullsrc=channel_layout=stereo:sample_rate=44100"])
         
-        # For videos, handle the start time trim
         if item_type == "video":
             start_time = item.get('start_time', 0)
             if start_time > 0:
@@ -1629,21 +1700,21 @@ def create_animatic(timeline_data, audio_path, out_w, out_h, keep_original_audio
             "-vf", vf_filter
         ]
         
-        # --- CORRECTED AUDIO LOGIC ---
+        # --- Corrected Audio Logic ---
         if keep_original_audio:
-            # If it's an image, map the silent audio source we added.
             if item_type == "image":
-                 cmd.extend(["-c:v", "libx264", "-t", str(item_duration), "-pix_fmt", "yuv420p", "-vf", vf_filter])
-                 output_options = ["-shortest", "-c:a", "aac"]
-            # For videos, just re-encode the existing audio to a standard format.
+                 # Map the video from the image and the silent audio
+                 cmd.extend(["-map", "1:v:0", "-map", "0:a:0"])
+                 output_options.extend(["-shortest", "-c:a", "aac"])
             else:
+                 # For videos, just re-encode the existing audio to a standard format.
                  output_options.extend(["-c:a", "aac"])
         else:
             # If we don't keep audio, make the clip silent.
             output_options.append("-an")
 
         cmd.extend(output_options)
-        cmd.extend(["-y", output_clip_path])
+        cmd.append(output_clip_path)
         
         run_ffmpeg_command(cmd, f"Processing clip {i+1}/{len(timeline_data)}")
         clip_paths.append(output_clip_path)
@@ -1653,22 +1724,19 @@ def create_animatic(timeline_data, audio_path, out_w, out_h, keep_original_audio
         for path in clip_paths:
             f.write(f"file '{os.path.abspath(path)}'\n")
     
-    # This path now contains the final video, either silent or with original audio
     combined_video_path = os.path.join(job_temp_dir, "combined_video.mp4")
     run_ffmpeg_command(["ffmpeg", "-f", "concat", "-safe", "0", "-i", file_list_path, "-c", "copy", "-y", combined_video_path], "Joining clips...")
 
     final_output_path = os.path.join(TEMP_DIR, f"animatic_final_{timestamp}.mp4")
     
-    # Decide what to do with the audio
     if not keep_original_audio and audio_path:
-        # Mux the project audio track onto the silent video
-        run_ffmpeg_command(["ffmpeg", "-i", combined_video_path, "-i", audio_path, "-c:v", "copy", "-c:a", "aac", "-shortest", "-y", final_output_path], "Muxing audio")
+        run_ffmpeg_command(["ffmpeg", "-y", "-i", combined_video_path, "-i", audio_path, "-c:v", "copy", "-c:a", "aac", "-shortest", final_output_path], "Muxing audio")
     else:
-        # Either we kept original audio, or there's no project audio, so just use the combined video as-is
         shutil.move(combined_video_path, final_output_path)
         
     shutil.rmtree(job_temp_dir)
     return final_output_path
+
 def detect_bpm(audio_path):
     if not audio_path:
         return "Please upload an audio track first."
@@ -2681,9 +2749,9 @@ with gr.Blocks(
             set_clip_start_btn.click(fn=None, js=storyboard_get_time_js, outputs=clip_start_time_input)
             set_clip_end_btn.click(fn=None, js=storyboard_get_time_js, outputs=clip_end_time_input)
 
-            timeline_up_btn.click(handle_timeline_action, [timeline_state, selected_timeline_index_state, gr.State("up")], timeline_state)
-            timeline_down_btn.click(handle_timeline_action, [timeline_state, selected_timeline_index_state, gr.State("down")], timeline_state)
-            timeline_remove_btn.click(handle_timeline_action, [timeline_state, selected_timeline_index_state, gr.State("remove")], timeline_state)
+            timeline_up_btn.click(handle_timeline_action, [timeline_state, selected_timeline_index_state, gr.State("up")], [timeline_state, timeline_df])
+            timeline_down_btn.click(handle_timeline_action, [timeline_state, selected_timeline_index_state, gr.State("down")], [timeline_state, timeline_df])
+            timeline_remove_btn.click(handle_timeline_action, [timeline_state, selected_timeline_index_state, gr.State("remove")], [timeline_state, timeline_df])
             
             clear_assets_btn.click(lambda: ([], gr.update(value=None), None), outputs=[assets_state, asset_gallery, asset_preview_gallery])
             clear_timeline_btn.click(lambda: [], outputs=[timeline_state])
@@ -3382,7 +3450,7 @@ with gr.Blocks(
                         
                         with gr.TabItem("Transcription", visible=(whisper is not None)):
                             gr.Markdown("## Transcribe Speech and Burn Subtitles")
-                            gr.Info("Uses OpenAI's Whisper model. First run will download model files. After transcribing a video, options to burn subtitles will appear.")
+                            gr.Info("Uses OpenAI's Whisper model with word-level timestamps. First run will download model files.")
                             transcribed_video_path_state = gr.State(None)
                             with gr.Row():
                                 with gr.Column():
@@ -3393,17 +3461,27 @@ with gr.Blocks(
                                     with gr.Row():
                                         transcribe_text = gr.Textbox(label="Transcription Result", lines=10, interactive=True, elem_id="transcription_textbox")
                                         copy_transcription_btn = gr.Button("📋 Copy")
-                                    transcribe_files = gr.File(label="Download Subtitle Files (.srt, .vtt)", file_count="multiple", interactive=False)
+                                    transcribe_files = gr.File(label="Download Subtitle Files (.srt, .vtt, .ass)", file_count="multiple", interactive=False)
                             
                             with gr.Accordion("🔥 Burn Subtitles onto Video", open=True, visible=False) as burn_accordion:
                                 gr.Markdown("Set styling and burn the generated subtitles into the video.")
                                 with gr.Row():
+                                    burn_style = gr.Radio(["Block", "Karaoke"], value="Block", label="Subtitle Style")
                                     burn_font_size = gr.Slider(1, 10, 5, step=1, label="Relative Font Size")
+                                with gr.Row(visible=True) as block_style_row:
                                     burn_words_per_line = gr.Slider(1, 20, 7, step=1, label="Max Words Per Line", info="Splits long subtitle lines for better readability.")
-                                    burn_font_color = gr.ColorPicker(value="#FFFFFF", label="Font Color")
+                                    burn_block_font_color = gr.ColorPicker(value="#FFFFFF", label="Font Color")
+                                with gr.Row(visible=False) as karaoke_style_row:
+                                    burn_karaoke_base_color = gr.ColorPicker(value="#FFFFFF", label="Base Color")
+                                    burn_karaoke_highlight_color = gr.ColorPicker(value="#FFFF00", label="Highlight Color")
+
                                 burn_btn = gr.Button("🔥 Burn Subtitles", variant="primary")
                                 burn_output_video = gr.Video(label="Video with Burned-in Subtitles", interactive=True, show_download_button=True, height=480)
+
+                            def toggle_subtitle_styles(style_choice):
+                                return gr.update(visible=(style_choice == "Block")), gr.update(visible=(style_choice == "Karaoke"))
                             
+                            burn_style.change(toggle_subtitle_styles, burn_style, [block_style_row, karaoke_style_row])
                             copy_transcription_btn.click(fn=None, js=copy_transcription_js)
                             transcribe_btn.click(
                                 fn=transcribe_and_prep_burn, 
@@ -3411,8 +3489,8 @@ with gr.Blocks(
                                 outputs=[transcribe_text, transcribe_files, transcribed_video_path_state, burn_accordion]
                             ).then(fn=None, js=fire_confetti_and_sound_js)
                             burn_btn.click(
-                                fn=burn_subtitles,
-                                inputs=[transcribed_video_path_state, transcribe_files, burn_font_size, burn_font_color, burn_words_per_line],
+                                fn=burn_subtitles_wrapper,
+                                inputs=[transcribed_video_path_state, transcribe_files, burn_style, burn_font_size, burn_block_font_color, burn_words_per_line, burn_karaoke_base_color, burn_karaoke_highlight_color],
                                 outputs=burn_output_video
                             ).then(fn=None, js=fire_confetti_and_sound_js)
 
@@ -3509,4 +3587,6 @@ if __name__ == "__main__":
         try: shutil.rmtree(TEMP_DIR)
         except OSError as e: print(f"Error removing temp directory {TEMP_DIR}: {e}")
     os.makedirs(TEMP_DIR, exist_ok=True)
-    demo.launch(inbrowser=True)
+    if whisper:
+        load_whisper_model("base") # Pre-load the default model on startup
+    demo.launch(inbrowser=True)