app.py

import os
import sys
from openai import OpenAI
# Set HF_HOME to a writable directory in /tmp
os.environ["HF_HOME"] = "/tmp/huggingface_cache"
os.makedirs("/tmp/huggingface_cache", exist_ok=True)
os.environ["PYTHONUSERBASE"] = "/tmp/.local"
os.makedirs("/tmp/.local", exist_ok=True)

user_site = os.path.join(os.environ["PYTHONUSERBASE"], "lib", "python3.9", "site-packages")

# Insert the user site-packages directory into sys.path if it's not already present.
if user_site not in sys.path:
    sys.path.insert(0, user_site)
# Optionally, you can also set TRANSFORMERS_CACHE for backwards compatibility:
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface_cache"

# Now import the rest of your dependencies
import re
import shutil
import subprocess
import time
import uuid
from threading import Timer


from flask import Flask, render_template, request, url_for, send_from_directory

app = Flask(__name__)

API_KEY = os.environ.get("OPENROUTER")
if not API_KEY:
    raise ValueError("Missing OPENROUTER environment variable.")
client = OpenAI(
  base_url="https://openrouter.ai/api/v1",
  api_key=API_KEY,
)
media_dir = os.path.join("/tmp", "manim_media")
os.makedirs(media_dir, exist_ok=True)

@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        prompt = request.form.get("prompt")
        if not prompt:
            return render_template("index.html")
        
        max_retries = 3
        attempt = 0
        last_error = None
        while attempt < max_retries:
            try:
                # Call the GenAI API to get the Manim code
                completion = client.chat.completions.create(
                    extra_body={},
                    model="anthropic/claude-3.7-sonnet",
                    messages=[
                        {
                        "role": "user",
                        "content": [
                            {
                            "type": "text",
                            "text": 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
                    """
                    },
                ]
                }
            ]
            )
                
                pattern = r"```python\s*(.*?)\s*```"
                match = re.search(pattern, completion.choices[0].message.content, re.DOTALL)
                if not match:
                    raise Exception("No python code block found in the AI response.")
                code = match.group(1)

                
                scene_match = re.search(r"class\s+(\w+)\(.*Scene.*\):", code)
                scene_name = scene_match.group(1) if scene_match else "MyScene"
                
                code_filename = f"generated_video_{uuid.uuid4().hex}.py"
                video_filename = f"output_video_{uuid.uuid4().hex}.mp4"
                
                code_filepath = os.path.join("/tmp", code_filename)
                with open(code_filepath, "w") as f:
                    f.write(code)
                
                
                cmd = [
                    "manim",
                    "-qm",
                    "--media_dir", media_dir,
                    "-o", video_filename,
                    code_filepath,
                    scene_name
                ]
                try:
                    subprocess.run(cmd, check=True, capture_output=True, text=True)
                except subprocess.CalledProcessError as cpe:
                    app.logger.error("Manim error output: %s", cpe.stderr)
                    raise Exception(f"Manim failed: {cpe.stderr}")
                
                expected_dir = os.path.join(media_dir, "videos", code_filename.replace(".py", ""), "720p30")
                video_path_in_media = os.path.join(expected_dir, video_filename)
                if not os.path.exists(video_path_in_media):
                    raise Exception(f"Manim did not produce the expected output file at {video_path_in_media}")
                
                tmp_video_path = os.path.join("/tmp", video_filename)
                shutil.move(video_path_in_media, tmp_video_path)

                def remove_files():
                    for fpath in [tmp_video_path, code_filepath]:
                        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)
                return render_template("result.html", video_url=video_url)
            
            except Exception as e:
                app.logger.error("Attempt %d failed: %s", attempt + 1, e)
                last_error = e
                attempt += 1
                time.sleep(1)
        
        return render_template("result.html", error="An error occurred. Please try again later.")
            
    return render_template("index.html")

@app.route("/video/<filename>")
def get_video(filename):
    return send_from_directory("/tmp", filename)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=False)