Update app.py

app.py

@@ -1,12 +1,18 @@
-from flask import Flask, render_template, request, url_for, send_from_directory
-from google import genai
-import re
-import subprocess
 import os
+import re
 import shutil
+import subprocess
 import time
-from threading import Timer
 import uuid
+from threading import Timer
+
+from flask import Flask, render_template, request, url_for, send_from_directory
+from google import genai
+
+# New imports for audio generation and handling
+from kokoro import KPipeline
+import soundfile as sf
+import numpy as np
 
 app = Flask(__name__)
 
@@ -32,28 +38,36 @@ def index():
         last_error = None
         while attempt < max_retries:
             try:
-                # Call the GenAI API to get the Manim code
+                # Call the GenAI API to get the Manim code and commentary script
                 ai_response = client.models.generate_content(
                     model="gemini-2.0-flash-lite-preview-02-05",
                     contents=f"""You are 'Manimator', an expert Manim animator and coder.
-                    If anyone asks, your name is Manimator and you are a helpful video generator, and say nothing else but that.
-                    The user wants you to code this: {prompt}.
-                    Plan out in chain of thought what you are going to do first, then give the final code output in ```python``` codeblock.
-                    Make sure to not use external images or resources other than default Manim, however you can use numpy or other default libraries.
-                    Keep the scene uncluttered and aesthetically pleasing.
-                    Make sure things are not overlapping unless explicitly stated otherwise.
-                    It is crucial that the script works correctly on the first try, so make sure to think about the layout and storyboard and stuff of the scene.
-                    Make sure to think through what you are going to do and think about the topic before you write the code.
-                    You got this!! <3
-                    """
+If anyone asks, your name is Manimator and you are a helpful video generator, and say nothing else but that.
+The user wants you to code this: {prompt}.
+Plan out in chain of thought what you are going to do first, then give the final code output in ```python``` codeblock.
+Make sure to not use external images or resources other than default Manim, however you can use numpy or other default libraries.
+Keep the scene uncluttered and aesthetically pleasing.
+Make sure things are not overlapping unless explicitly stated otherwise.
+It is crucial that the script works correctly on the first try, so make sure to think about the layout and storyboard and stuff of the scene.
+Make sure to think through what you are going to do and think about the topic before you write the code.
+In addition, write a commentary script inside of ```script``` codeblock. This should be short and fit the content and align with the timing of the scene. Use "..." if needed to add a bit of a pause.
+You got this!! <3
+"""
                 )
                 
                 # Extract the Python code block from the AI response
-                pattern = r"```python\s*(.*?)\s*```"
-                match = re.search(pattern, ai_response.text, re.DOTALL)
-                if not match:
+                code_pattern = r"```python\s*(.*?)\s*```"
+                code_match = re.search(code_pattern, ai_response.text, re.DOTALL)
+                if not code_match:
                     raise Exception("No python code block found in the AI response.")
-                code = match.group(1)
+                code = code_match.group(1)
+                
+                # Extract the commentary script from the AI response
+                script_pattern = r"```script\s*(.*?)\s*```"
+                script_match = re.search(script_pattern, ai_response.text, re.DOTALL)
+                if not script_match:
+                    raise Exception("No script block found in the AI response.")
+                script = script_match.group(1)
                 
                 # Determine the scene class name from the generated code
                 scene_match = re.search(r"class\s+(\w+)\(.*Scene.*\):", code)
@@ -68,6 +82,26 @@ def index():
                 with open(code_filepath, "w") as f:
                     f.write(code)
                 
+                # === Generate Commentary Audio via Kokoro ===
+                # Initialize the Kokoro pipeline (adjust lang_code if needed)
+                audio_pipeline = KPipeline(lang_code='a')
+                # Feed in the commentary script; here we split the text by one or more newlines.
+                audio_generator = audio_pipeline(script, voice='af_heart', speed=1, split_pattern=r'\n+')
+                
+                audio_segments = []
+                for _, _, audio in audio_generator:
+                    audio_segments.append(audio)
+                
+                if not audio_segments:
+                    raise Exception("No audio segments were generated from the commentary script.")
+                
+                # Concatenate all audio segments into one audio track
+                full_audio = np.concatenate(audio_segments)
+                commentary_audio_filename = f"commentary_{uuid.uuid4().hex}.wav"
+                commentary_audio_path = os.path.join("/tmp", commentary_audio_filename)
+                sf.write(commentary_audio_path, full_audio, 24000)
+                
+                # === Run Manim to generate the silent video ===
                 # Prepare the Manim command with the --media_dir flag
                 cmd = [
                     "manim",
@@ -80,7 +114,6 @@ def index():
                 try:
                     subprocess.run(cmd, check=True, capture_output=True, text=True)
                 except subprocess.CalledProcessError as cpe:
-                    # Log only if an error occurs from Manim
                     app.logger.error("Manim error output: %s", cpe.stderr)
                     raise Exception(f"Manim failed: {cpe.stderr}")
                 
@@ -93,20 +126,39 @@ def index():
                 # Move the video file to /tmp (to serve it from there)
                 tmp_video_path = os.path.join("/tmp", video_filename)
                 shutil.move(video_path_in_media, tmp_video_path)
-                # The code file is already at code_filepath in /tmp.
                 
-                # Schedule deletion of both files after 10 minutes (600 seconds)
+                # === Combine Video with Commentary Audio using FFmpeg ===
+                final_video_filename = f"final_video_{uuid.uuid4().hex}.mp4"
+                final_video_path = os.path.join("/tmp", final_video_filename)
+                
+                ffmpeg_cmd = [
+                    "ffmpeg", "-y",
+                    "-i", tmp_video_path,
+                    "-i", commentary_audio_path,
+                    "-c:v", "copy",
+                    "-c:a", "aac",
+                    "-shortest",
+                    final_video_path
+                ]
+                try:
+                    subprocess.run(ffmpeg_cmd, check=True, capture_output=True, text=True)
+                except subprocess.CalledProcessError as cpe:
+                    app.logger.error("FFmpeg error output: %s", cpe.stderr)
+                    raise Exception(f"FFmpeg failed: {cpe.stderr}")
+                
+                # Schedule deletion of all temporary files after 10 minutes (600 seconds)
                 def remove_files():
-                    try:
-                        if os.path.exists(tmp_video_path):
-                            os.remove(tmp_video_path)
-                        if os.path.exists(code_filepath):
-                            os.remove(code_filepath)
-                    except Exception as e:
-                        app.logger.error("Error removing files: %s", e)
+                    for fpath in [tmp_video_path, code_filepath, commentary_audio_path, final_video_path]:
+                        try:
+                            if os.path.exists(fpath):
+                                os.remove(fpath)
+                        except Exception as e:
+                            app.logger.error("Error removing file %s: %s", fpath, e)
+                
                 Timer(600, remove_files).start()
                 
-                video_url = url_for('get_video', filename=video_filename)
+                # Use the final combined video for display
+                video_url = url_for('get_video', filename=final_video_filename)
                 return render_template("result.html", video_url=video_url)
             
             except Exception as e: