Update app.py

app.py

@@ -1,374 +1,181 @@
-import gradio as gr
+import streamlit as st
 import subprocess
-import os
+import tempfile
 import re
-import shutil
 from pathlib import Path
-import tempfile
-
-# Try to import google.generativeai
-try:
-    import google.generativeai as genai
-    GENAI_AVAILABLE = True
-except ImportError:
-    GENAI_AVAILABLE = False
-    print("⚠️ google-generativeai not installed")
+import google.generativeai as genai
+import os
 
-# Get API key from environment
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
 
-def fix_all_known_errors(code):
-    """Fix ALL known Manim errors"""
-    
-    # 1. Fix colors to base 10 only
+st.set_page_config(page_title="Manim Video Generator", page_icon="🎬")
+
+def fix_code(code):
+    """Fix common Manim errors"""
     color_map = {
-        'ORANGE_A': 'ORANGE', 'ORANGE_B': 'ORANGE', 'ORANGE_C': 'ORANGE', 'ORANGE_D': 'ORANGE', 'ORANGE_E': 'ORANGE',
-        'GRAY_A': 'WHITE', 'GRAY_B': 'WHITE', 'GRAY_C': 'WHITE', 'GRAY_D': 'WHITE', 'GRAY_E': 'WHITE',
-        'GREY': 'WHITE', 'GRAY': 'WHITE', 'GREY_A': 'WHITE', 'GREY_B': 'WHITE',
-        'PURPLE_A': 'PURPLE', 'PURPLE_B': 'PURPLE', 'PURPLE_C': 'PURPLE', 'PURPLE_D': 'PURPLE', 'PURPLE_E': 'PURPLE',
-        'RED_A': 'RED', 'RED_B': 'RED', 'RED_C': 'RED', 'RED_D': 'RED', 'RED_E': 'RED',
-        'BLUE_A': 'BLUE', 'BLUE_B': 'BLUE', 'BLUE_C': 'BLUE', 'BLUE_D': 'BLUE', 'BLUE_E': 'BLUE',
-        'GREEN_A': 'GREEN', 'GREEN_B': 'GREEN', 'GREEN_C': 'GREEN', 'GREEN_D': 'GREEN', 'GREEN_E': 'GREEN',
-        'YELLOW_A': 'YELLOW', 'YELLOW_B': 'YELLOW', 'YELLOW_C': 'YELLOW', 'YELLOW_D': 'YELLOW', 'YELLOW_E': 'YELLOW',
-        'BROWN': 'ORANGE', 'GOLD': 'YELLOW', 'SILVER': 'WHITE', 'BRONZE': 'ORANGE',
-        'MAROON': 'RED', 'NAVY': 'BLUE', 'CYAN': 'TEAL', 'MAGENTA': 'PINK', 'LIME': 'GREEN',
+        'ORANGE_A': 'ORANGE', 'GRAY_A': 'WHITE', 'GREY': 'WHITE',
+        'PURPLE_A': 'PURPLE', 'RED_A': 'RED', 'BLUE_A': 'BLUE',
+        'BROWN': 'ORANGE', 'GOLD': 'YELLOW'
     }
     for old, new in color_map.items():
         code = re.sub(rf'\b{old}\b', new, code)
     
-    # 2. Fix invalid objects
-    code = re.sub(r'Checkmark\([^)]*\)', 'Circle(radius=0.3, color=GREEN, fill_opacity=1)', code)
-    code = re.sub(r'CheckMark\([^)]*\)', 'Circle(radius=0.3, color=GREEN, fill_opacity=1)', code)
-    code = re.sub(r'XMark\([^)]*\)', 'Cross(color=RED).scale(0.5)', code)
-    code = re.sub(r'SVGMobject\([^)]+\)', 'Circle(radius=0.8, color=PINK, fill_opacity=0.8)', code)
-    code = re.sub(r'ImageMobject\([^)]+\)', 'Square(side_length=2, color=BLUE, fill_opacity=0.5)', code)
-    
-    # 3. Fix DecimalNumber methods
-    code = re.sub(r'\.add_prefix\([^)]+\)', '', code)
-    code = re.sub(r'\.add_suffix\([^)]+\)', '', code)
-    
-    # 4. Fix axis labels - replace with Text
-    code = re.sub(r'(\w+)\s*=\s*axes\.get_x_axis_label\([^)]+\)', 
-                  r'\1 = Text("X", font_size=24).next_to(axes.x_axis, DOWN, buff=0.3)', code)
-    code = re.sub(r'(\w+)\s*=\s*axes\.get_y_axis_label\([^)]+\)', 
-                  r'\1 = Text("Y", font_size=24).next_to(axes.y_axis, LEFT, buff=0.3).rotate(90*DEGREES)', code)
     code = re.sub(r'\.get_x_axis_label\([^)]+\)', '', code)
     code = re.sub(r'\.get_y_axis_label\([^)]+\)', '', code)
-    
-    # 5. Limit font sizes
+    code = re.sub(r'Checkmark\([^)]*\)', 'Circle(radius=0.3, color=GREEN)', code)
     code = re.sub(r'font_size=(\d+)', lambda m: f'font_size={min(int(m.group(1)), 36)}', code)
     
-    # 6. Fix shifts to safe values
-    code = re.sub(r'\.shift\(UP\s*\*\s*\d+\.?\d*\)', '.shift(UP*1.5)', code)
-    code = re.sub(r'\.shift\(DOWN\s*\*\s*\d+\.?\d*\)', '.shift(DOWN*1.5)', code)
-    code = re.sub(r'\.shift\(LEFT\s*\*\s*\d+\.?\d*\)', '.shift(LEFT*3)', code)
-    code = re.sub(r'\.shift\(RIGHT\s*\*\s*\d+\.?\d*\)', '.shift(RIGHT*3)', code)
-    
-    # 7. Force buff on all to_edge
-    code = re.sub(r'\.to_edge\((UP|DOWN|LEFT|RIGHT)\)(?!\s*,)', r'.to_edge(\1, buff=1)', code)
-    
-    # 8. Replace Tex with MathTex
-    code = re.sub(r'\bTex\(', 'MathTex(', code)
-    
-    # 9. CRITICAL FIX: Add color=BLACK to Text objects that don't have color specified
-    # This prevents invisible white text on white background
-    # Match Text(...) that doesn't have color= parameter
-    def add_black_color(match):
-        text_content = match.group(0)
-        # Check if color parameter already exists
-        if 'color=' not in text_content:
-            # Find the closing parenthesis of Text()
-            # Add color=BLACK before the last closing paren
-            if ', font_size=' in text_content:
-                # Insert after font_size
-                text_content = text_content.replace(')', ', color=BLACK)', 1)
-            else:
-                # Insert before closing paren
-                text_content = text_content[:-1] + ', color=BLACK)'
-        return text_content
-    
-    # Apply to all Text() objects
-    code = re.sub(r'Text\([^)]+\)', add_black_color, code)
+    # Add color=BLACK to Text without color
+    def add_color(match):
+        text = match.group(0)
+        if 'color=' not in text:
+            text = text[:-1] + ', color=BLACK)'
+        return text
+    code = re.sub(r'Text\([^)]+\)', add_color, code)
     
     return code
 
-def generate_code_with_gemini(prompt):
-    """Generate Manim code using Gemini API"""
-    
-    if not GENAI_AVAILABLE:
-        return None, "❌ google-generativeai package not installed"
-    
+def generate_code(prompt):
+    """Generate Manim code using Gemini"""
     if not GEMINI_API_KEY:
-        return None, "❌ GEMINI_API_KEY not set in environment variables"
+        return None, "❌ GEMINI_API_KEY not set in secrets"
     
     try:
         genai.configure(api_key=GEMINI_API_KEY)
-        models = [m.name for m in genai.list_models() if 'generateContent' in m.supported_generation_methods]
-        model_name = next((m for m in models if 'flash' in m.lower()), models[0])
-        
-        model = genai.GenerativeModel(model_name)
+        model = genai.GenerativeModel('gemini-1.5-flash')
         
-        full_prompt = f"""You are a Manim code generator. Generate WORKING Manim animation code.
+        full_prompt = f"""Generate simple Manim code.
 
-USER REQUEST: {prompt}
-
-CRITICAL RULES YOU MUST FOLLOW:
-
-1. ALWAYS start with: from manim import *
-
-2. Class structure MUST be:
-```python
-from manim import *
-
-class MyScene(Scene):
-    def construct(self):
-        # Your animation code here
-        # MUST have at least one self.play() or self.add()
-```
-
-3. COLORS - ONLY use these 10: BLUE, RED, GREEN, YELLOW, ORANGE, PURPLE, PINK, TEAL, WHITE, BLACK
-   - NO color variants like BLUE_A, RED_B, etc.
-
-4. TEXT RULES:
-   - Use: Text("message", font_size=36, color=BLACK)
-   - ALWAYS FadeOut text before creating new text in same position
-   - Example:
-   ```python
-   text1 = Text("First", font_size=36, color=BLACK)
-   self.play(Write(text1))
-   self.wait(1)
-   self.play(FadeOut(text1))  # Must FadeOut!
-   
-   text2 = Text("Second", font_size=36, color=BLACK)
-   self.play(Write(text2))
+RULES:
+1. Use ONLY: BLUE, RED, GREEN, YELLOW, ORANGE, PURPLE, PINK, TEAL, WHITE, BLACK
+2. Background is WHITE - ALL Text MUST have color=BLACK
+3. Always FadeOut before new text:
+   text = Text("Hello", color=BLACK)
+   self.play(Write(text))
    self.wait(1)
-   self.play(FadeOut(text2))
-   ```
-
-5. SHAPES:
-   - Circle: Circle(radius=1, color=BLUE)
-   - Square: Square(side_length=2, color=RED)
-   - Rectangle: Rectangle(width=3, height=2, color=GREEN)
-   - Line: Line(start=LEFT, end=RIGHT, color=YELLOW)
-
-6. POSITIONING:
-   - Use .shift(UP), .shift(DOWN), .shift(LEFT), .shift(RIGHT)
-   - Use .to_edge(UP), .to_edge(DOWN), .to_edge(LEFT), .to_edge(RIGHT)
-   - Use .move_to(ORIGIN)
+   self.play(FadeOut(text))
+4. Keep simple, 10-15 seconds max
+5. MUST include: from manim import *
 
-7. ANIMATIONS:
-   - Show objects: Create(), Write(), FadeIn()
-   - Hide objects: FadeOut(), Uncreate()
-   - Transform: Transform(obj1, obj2)
-   - Move: obj.animate.shift(UP)
+User request: {prompt}
 
-8. NEVER USE:
-   - get_x_axis_label(), get_y_axis_label()
-   - SVGMobject, ImageMobject
-   - Checkmark, XMark
-   - .add_prefix(), .add_suffix()
-
-9. SIMPLE AND WORKING:
-   - Keep it simple
-   - Test each animation
-   - Make sure there's visible content
-   - Animation should be 5-10 seconds long
-
-Generate COMPLETE, WORKING code that creates a visible animation for: {prompt}
-
-The code must actually show something on screen!"""
+Generate complete working code:"""
         
         response = model.generate_content(full_prompt)
-        code = response.text
-        
-        return code, None
-        
+        return response.text, None
     except Exception as e:
-        return None, f"❌ Gemini API error: {str(e)}"
+        return None, str(e)
 
-def render_video(prompt, quality="low"):
-    """Main function to generate and render Manim video"""
-    
-    if not prompt or not prompt.strip():
-        yield "❌ Please enter a prompt!", None, None
-        return
+def render_video(prompt):
+    """Render Manim video"""
+    temp_dir = Path(tempfile.mkdtemp())
     
-    try:
-        # Create temp directory
-        temp_dir = Path(tempfile.mkdtemp(prefix="manim_"))
-        
-        yield "🤖 Generating code with Gemini AI...", None, None
-        
-        # Generate code
-        code, error = generate_code_with_gemini(prompt)
+    # Generate code
+    with st.spinner("🤖 Generating code..."):
+        code, error = generate_code(prompt)
         
         if error:
-            yield error, None, None
-            return
-        
-        if not code:
-            yield "❌ Failed to generate code", None, None
-            return
+            st.error(error)
+            return None, None
         
-        # Extract Python code if wrapped in markdown
+        # Extract Python code
         if "```python" in code:
             code = code.split("```python")[1].split("```")[0]
         elif "```" in code:
             code = code.split("```")[1].split("```")[0]
         code = code.strip()
         
-        # Apply fixes
-        code = fix_all_known_errors(code)
-        
-        # Find Scene class
-        match = re.search(r'class\s+(\w+)\s*\(Scene\)', code)
-        if not match:
-            yield "❌ No Scene class found in generated code!", code, None
-            return
-        
-        class_name = match.group(1)
-        
-        # Save code to file
-        code_file = temp_dir / "animation.py"
-        with open(code_file, 'w', encoding='utf-8') as f:
-            f.write(code)
-        
-        yield f"✓ Code generated (Scene: {class_name})\n🎬 Rendering video (this may take 1-2 minutes)...", code, None
-        
-        # Render video with absolute paths
-        quality_map = {'low': '-ql', 'medium': '-qm', 'high': '-qh'}
-        quality_flag = quality_map.get(quality, '-ql')
-        
-        abs_code_file = str(code_file.absolute())
-        media_dir = str((temp_dir / "media").absolute())
-        
-        command = f"manim {abs_code_file} {class_name} {quality_flag} --disable_caching --media_dir {media_dir}"
-        
-        process = subprocess.Popen(
-            command, shell=True,
-            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
-            text=True
-        )
-        
-        output_lines = []
-        has_error = False
-        for line in process.stdout:
-            output_lines.append(line)
-            # Check for common errors in output
-            if any(err in line.lower() for err in ['error', 'exception', 'traceback', 'failed']):
-                has_error = True
-        
-        process.wait()
-        
-        # Get full output for debugging
-        full_output = ''.join(output_lines)
+        # Fix code
+        code = fix_code(code)
+    
+    # Find Scene class
+    match = re.search(r'class\s+(\w+)\s*\(Scene\)', code)
+    if not match:
+        st.error("No Scene class found!")
+        return None, code
+    
+    class_name = match.group(1)
+    
+    # Save code
+    code_file = temp_dir / "animation.py"
+    code_file.write_text(code)
+    
+    # Render video
+    with st.spinner(f"🎬 Rendering {class_name}... (this takes 1-2 minutes)"):
+        cmd = f"manim {code_file} {class_name} -ql --disable_caching --media_dir {temp_dir}/media"
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=180)
         
-        if process.returncode == 0:
-            # Find the generated video
-            media_path = temp_dir / "media"
-            video_files = list(media_path.rglob("*.mp4"))
-            
-            if video_files:
-                video_path = str(video_files[0])
-                # Check if video file is too small (blank video indicator)
-                video_size = os.path.getsize(video_path)
-                
-                if video_size < 5000:  # Less than 5KB means blank/error
-                    yield f"⚠️ Video rendered but seems blank (only {video_size} bytes)!\n\nMaybe there's an error in the code. Check the generated code below.\n\nRender output (last 40 lines):\n{chr(10).join(output_lines[-40:])}", code, None
-                elif has_error:
-                    yield f"⚠️ Video rendered but errors detected!\n\nRender output (last 40 lines):\n{chr(10).join(output_lines[-40:])}", code, video_path
-                else:
-                    yield f"✅ Video rendered successfully! 🎉", code, video_path
+        if result.returncode == 0:
+            # Find video
+            videos = list((temp_dir / "media").rglob("*.mp4"))
+            if videos:
+                return videos[0], code
             else:
-                yield "❌ Video file not found after rendering", code, None
+                st.error("Video file not found after rendering")
+                st.code(result.stdout, language="text")
+                return None, code
         else:
-            error_msg = ''.join(output_lines[-40:])  # Last 40 lines
-            yield f"❌ Rendering failed:\n\n{error_msg}", code, None
-        
-        # Cleanup
-        try:
-            shutil.rmtree(temp_dir)
-        except:
-            pass
+            st.error("Rendering failed!")
+            st.code(result.stderr, language="text")
+            return None, code
+
+# UI
+st.title("🎬 Manim Video Generator")
+st.markdown("Create mathematical animations using AI!")
+
+if not GEMINI_API_KEY:
+    st.error("⚠️ GEMINI_API_KEY not set! Add it in Settings → Secrets")
+    st.markdown("Get your free API key: https://aistudio.google.com/app/apikey")
+    st.stop()
+
+# Input
+prompt = st.text_area(
+    "What animation do you want?",
+    placeholder="e.g., Show a blue circle morphing into a red square",
+    height=100
+)
+
+# Examples
+st.markdown("**Examples:**")
+col1, col2, col3 = st.columns(3)
+with col1:
+    if st.button("🔵 Circle → Square"):
+        prompt = "Show a blue circle morphing into a red square"
+with col2:
+    if st.button("📐 Pythagorean"):
+        prompt = "Animate the Pythagorean theorem: a² + b² = c²"
+with col3:
+    if st.button("〰️ Sine Wave"):
+        prompt = "Show a sine wave moving"
+
+# Generate button
+if st.button("🎬 Generate Video", type="primary", use_container_width=True):
+    if not prompt:
+        st.warning("Please enter a prompt!")
+    else:
+        video_path, code = render_video(prompt)
+        
+        if video_path:
+            st.success("✅ Video rendered successfully!")
             
-    except Exception as e:
-        import traceback
-        error_details = traceback.format_exc()
-        yield f"❌ Error: {str(e)}\n\nDetails:\n{error_details}", None, None
-
-# Gradio Interface
-def create_interface():
-    with gr.Blocks(title="Manim Video Generator") as demo:
-        gr.Markdown("""
-        # 🎬 Manim Video Generator
-        Create mathematical animations using AI! Describe what you want to see animated.
-        
-        **⚠️ IMPORTANT: You need to set GEMINI_API_KEY in Space Settings → Variables and secrets**
-        
-        Get your free Gemini API key from: https://aistudio.google.com/app/apikey
-        
-        **Examples:**
-        - "Explain Pythagorean theorem with animation"
-        - "Show a sine wave transforming into a cosine wave"
-        - "Animate a circle morphing into a square"
-        - "Show the concept of derivatives with a tangent line"
-        """)
-        
-        with gr.Row():
-            with gr.Column():
-                prompt_input = gr.Textbox(
-                    label="What animation do you want?",
-                    placeholder="e.g., Show a circle morphing into a square",
-                    lines=4
-                )
-                quality_input = gr.Radio(
-                    choices=["low", "medium", "high"],
-                    value="low",
-                    label="Quality (low is faster, recommended)"
+            # Display video
+            st.video(str(video_path))
+            
+            # Download button
+            with open(video_path, 'rb') as f:
+                st.download_button(
+                    label="📥 Download Video",
+                    data=f.read(),
+                    file_name=f"manim_{video_path.stem}.mp4",
+                    mime="video/mp4"
                 )
-                generate_btn = gr.Button("🎬 Generate Video", variant="primary", size="lg")
-                
-            with gr.Column():
-                status_output = gr.Textbox(label="Status", lines=4)
-                video_output = gr.Video(label="Generated Animation")
-        
-        code_output = gr.Code(label="Generated Manim Code", language="python", lines=15)
-        
-        gr.Markdown("""
-        ### 💡 Tips:
-        - Be specific in your description
-        - Start with simple prompts to test
-        - Low quality renders faster (30 seconds to 1 minute)
-        - Medium/High quality may take 2-3 minutes
-        - Complex animations may timeout on free tier
-        
-        ### 🔧 Setup Required:
-        1. Get a free Gemini API key from: https://aistudio.google.com/app/apikey
-        2. In your Space, go to Settings → Variables and secrets
-        3. Add: Name: `GEMINI_API_KEY`, Value: `your-api-key-here`
-        4. Restart the Space
-        """)
-        
-        generate_btn.click(
-            fn=render_video,
-            inputs=[prompt_input, quality_input],
-            outputs=[status_output, code_output, video_output]
-        )
-        
-        # Examples
-        gr.Examples(
-            examples=[
-                ["Show a blue circle morphing into a red square", "low"],
-                ["Animate the Pythagorean theorem: a² + b² = c²", "low"],
-                ["Show a sine wave moving across the screen", "low"],
-                ["Create an animation showing 3 dots forming a triangle", "low"],
-                ["Show a number counting from 0 to 10", "low"],
-            ],
-            inputs=[prompt_input, quality_input]
-        )
-    
-    return demo
-
-if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch()
+            
+            # Show code
+            with st.expander("📝 View Generated Code"):
+                st.code(code, language="python")
+        elif code:
+            st.warning("Video rendering failed, but here's the generated code:")
+            st.code(code, language="python")
+
+st.markdown("---")
+st.markdown("💡 **Tips:** Start with simple prompts • Rendering takes 1-2 minutes • Be specific about what you want")