app.py

import gradio as gr
import tempfile
import os
import subprocess
import json
import traceback
import shutil
from pathlib import Path
import time
import threading
import queue

class ManimMCPServer:
    def __init__(self):
        self.temp_dir = tempfile.mkdtemp()
        self.max_execution_time = 60  # 60 seconds timeout
        
    def validate_manim_code(self, code):
        """ Basic validation of Manim code """
        dangerous_imports = [
            'subprocess', 'os.system', 'eval', 'exec', 'open',
            '__import__', 'compile', 'globals', 'locals'
        ]
        
        for dangerous in dangerous_imports:
            if dangerous in code:
                return False, f"Dangerous operation detected: {dangerous}"
        
        # Check if it's a valid Python-like structure
        if 'from manim import *' not in code and 'import manim' not in code:
            return False, "Code must import manim"
            
        if 'class' not in code or 'Scene' not in code:
            return False, "Code must contain a Scene class"
            
        return True, "Valid"
    
    def execute_with_timeout(self, cmd, timeout):
        """ Execute command with timeout """
        result_queue = queue.Queue()
        
        def run_command():
            try:
                result = subprocess.run(
                    cmd, 
                    shell=True, 
                    capture_output=True, 
                    text=True,
                    cwd=self.temp_dir
                )
                result_queue.put(('success', result))
            except Exception as e:
                result_queue.put(('error', str(e)))
        
        thread = threading.Thread(target=run_command)
        thread.daemon = True
        thread.start()
        thread.join(timeout)
        
        if thread.is_alive():
            return False, "Execution timed out"
        
        if result_queue.empty():
            return False, "Execution failed"
            
        status, result = result_queue.get()
        if status == 'error':
            return False, result
        
        return True, result
    
    def render_manim_animation(self, code, scene_name=None, quality="medium_quality"):
        """ Render Manim animation and return video path """
        try:
            # Validate code
            is_valid, message = self.validate_manim_code(code)
            if not is_valid:
                return None, f"Code validation failed: {message}"
            
            # Create temporary Python file
            script_path = os.path.join(self.temp_dir, "animation.py")
            with open(script_path, 'w') as f:
                f.write(code)
            
            # Auto-detect scene name if not provided
            if not scene_name:
                lines = code.split('\n')
                for line in lines:
                    if line.strip().startswith('class ') and 'Scene' in line:
                        scene_name = line.split('class ')[1].split('(')[0].strip()
                        break
                
                if not scene_name:
                    return None, "Could not detect scene name. Please specify scene_name parameter."
            
            # Quality settings
            quality_map = {
                "low_quality": "-ql",
                "medium_quality": "-qm", 
                "high_quality": "-qh",
                "production_quality": "-qp"
            }
            
            quality_flag = quality_map.get(quality, "-qm")
            
            # Render animation
            cmd = f"manim {quality_flag} animation.py {scene_name}"
            
            success, result = self.execute_with_timeout(cmd, self.max_execution_time)
            
            if not success:
                return None, f"Rendering failed: {result}"
            
            if result.returncode != 0:
                return None, f"Manim execution failed:\nSTDOUT: {result.stdout}\nSTDERR: {result.stderr}"
            
            # Find generated video file
            media_dir = os.path.join(self.temp_dir, "media", "videos", "animation")
            
            if not os.path.exists(media_dir):
                return None, "No output directory found. Animation may not have been generated."
            
            # Look for MP4 files
            video_files = []
            for root, dirs, files in os.walk(media_dir):
                for file in files:
                    if file.endswith('.mp4'):
                        video_files.append(os.path.join(root, file))
            
            if not video_files:
                return None, "No MP4 files generated. Check your scene code."
                
            # Return the most recent video file
            latest_video = max(video_files, key=os.path.getctime)
            return latest_video, "Animation rendered successfully!"
            
        except Exception as e:
            return None, f"Unexpected error: {str(e)}\n{traceback.format_exc()}"
    
    def cleanup(self):
        """Clean up temporary files"""
        try:
            shutil.rmtree(self.temp_dir)
        except:
            pass

# Initialize the server
mcp_server = ManimMCPServer()

# Example Manim scripts
EXAMPLE_SCRIPTS = {
    "Hello World": '''from manim import *

class HelloWorld(Scene):
    def construct(self):
        text = Text("Hello, World!", font_size=72)
        text.set_color(BLUE)
        self.play(Write(text))
        self.wait(2)
        self.play(FadeOut(text))''',
    
    "Simple Animation": '''from manim import *

class SimpleAnimation(Scene):
    def construct(self):
        circle = Circle()
        circle.set_fill(PINK, opacity=0.5)
        
        square = Square()
        square.set_fill(BLUE, opacity=0.5)
        
        self.play(Create(circle))
        self.play(Transform(circle, square))
        self.play(Rotate(square, PI/2))
        self.wait()''',
        
    "Mathematical Formula": '''from manim import *

class MathFormula(Scene):
    def construct(self):
        formula = MathTex(r"\\frac{d}{dx}\\left(\\frac{1}{x}\\right) = -\\frac{1}{x^2}")
        formula.scale(2)
        formula.set_color(YELLOW)
        
        self.play(Write(formula))
        self.wait(2)
        
        box = SurroundingRectangle(formula, color=WHITE)
        self.play(Create(box))
        self.wait()'''
}

def render_animation(code, scene_name, quality):
    """ Gradio interface function """
    if not code.strip():
        return None, "Please provide Manim code"
    
    try:
        video_path, message = mcp_server.render_manim_animation(code, scene_name, quality)
        
        if video_path:
            return video_path, f"✅ {message}"
        else:
            return None, f"❌ {message}"
            
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def load_example(example_name):
    """ Load example script """
    return EXAMPLE_SCRIPTS.get(example_name, "")

def clear_all():
    """Clear all fields"""
    return "", "", None, ""

# Create Gradio interface
with gr.Blocks(title="Manim MCP Server", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎬 Manim MCP Server
    
    A Gradio-based MCP-like server for rendering Manim Community Edition animations.
    Write your Manim code, specify the scene name, and get an MP4 video output!
    
    ## How to use:
    1. Write or select a Manim script
    2. Specify the scene class name
    3. Choose quality settings
    4. Click "Render Animation"
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown("### 📝 Manim Code")
            
            with gr.Row():
                example_dropdown = gr.Dropdown(
                    choices=list(EXAMPLE_SCRIPTS.keys()),
                    label="Load Example",
                    value=None
                )
                clear_btn = gr.Button("Clear All", size="sm")
            
            code_input = gr.Code(
                label="Manim Script",
                language="python",
                lines=20,
                value=EXAMPLE_SCRIPTS["Hello World"]
            )
            
            with gr.Row():
                scene_name_input = gr.Textbox(
                    label="Scene Class Name",
                    placeholder="e.g., HelloWorld",
                    value="HelloWorld"
                )
                quality_input = gr.Dropdown(
                    choices=["low_quality", "medium_quality", "high_quality", "production_quality"],
                    label="Quality",
                    value="medium_quality"
                )
            
            render_btn = gr.Button("🎬 Render Animation", variant="primary", size="lg")
            
        with gr.Column(scale=1):
            gr.Markdown("### 🎥 Output")
            
            status_output = gr.Textbox(
                label="Status",
                lines=3,
                interactive=False
            )
            
            video_output = gr.Video(
                label="Generated Animation",
                height=400
            )
    
    gr.Markdown("""
    ### 📋 Tips:
    - Make sure your scene class inherits from `Scene`
    - Import manim with `from manim import *`
    - Use `self.play()` to animate objects
    - End with `self.wait()` for a pause
    - Keep animations short to avoid timeouts
    
    ### ⚠️ Limitations:
    - 60-second execution timeout
    - No dangerous operations allowed
    - Medium quality by default to save processing time
    """)
    
    # Event handlers
    example_dropdown.change(
        fn=load_example,
        inputs=[example_dropdown],
        outputs=[code_input]
    )
    
    clear_btn.click(
        fn=clear_all,
        outputs=[code_input, scene_name_input, video_output, status_output]
    )
    
    render_btn.click(
        fn=render_animation,
        inputs=[code_input, scene_name_input, quality_input],
        outputs=[video_output, status_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )