remove visual refinement for testing

Dockerfile

@@ -3,10 +3,9 @@ FROM manimcommunity/manim:stable
 WORKDIR /app
 COPY app /app
 
-# Web + Google Gemini (new SDK) + dotenv (optional)
+# Web server + Gemini SDK + dotenv (optional for local dev)
 RUN pip install --no-cache-dir fastapi uvicorn[standard] pydantic python-dotenv google-genai
 
-# Space default port
 ENV PORT=7860
 EXPOSE 7860
 

app/main.py

@@ -1,28 +1,27 @@
-import os, re, uuid, subprocess, sys, base64, traceback
+import os, re, uuid, subprocess, sys, traceback
 from pathlib import Path
-from typing import List, Tuple, Optional
+from typing import Optional, Tuple
 
 from fastapi import FastAPI, HTTPException, Response
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
-# Optional .env support (local dev)
+# Optional .env for local testing
 from dotenv import load_dotenv
 load_dotenv()
 
-# ---------- Gemini (keep your original SDK & calls) ----------
+# -------- Gemini (same SDK style as your Flask app) --------
 from google import genai
 from google.genai import types
 
 API_KEY = os.getenv("GEMINI_API_KEY", "")
-MODEL = os.getenv("GEMINI_MODEL", "gemini-2.5-pro")
-PORT  = int(os.getenv("PORT", "7860"))
+MODEL   = os.getenv("GEMINI_MODEL", "gemini-2.5-pro")
+PORT    = int(os.getenv("PORT", "7860"))
 
-# Configure Gemini client (exactly like your Flask code)
 client = genai.Client(api_key=API_KEY) if API_KEY else None
 
-# ---------- FastAPI ----------
-app = FastAPI(title="Manim Render API (Gemini → Manim → MP4)")
+# -------- FastAPI app --------
+app = FastAPI(title="Manim Render API (error-refine only)")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],   # tighten in prod
@@ -30,122 +29,43 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Work dirs (ephemeral per Space session is fine)
-RUNS   = Path("runs"); RUNS.mkdir(parents=True, exist_ok=True)
+RUNS = Path("runs"); RUNS.mkdir(parents=True, exist_ok=True)
 
-SYSTEM_PROMPT = """You are a Manim CE (0.19.x) code generator and refiner.
+SYSTEM_PROMPT = """You are a Manim CE (0.19.x) code generator/refiner.
 Return ONLY valid Python code (no backticks, no prose).
 Define exactly one class: AutoScene(Scene).
-Keep it short (≤ ~120 s) and renderable.
+Keep it short (preferably ≤ ~60 s) and quickly renderable.
 
 Use: from manim import *
 Allowed imports: manim, math, numpy.
-Forbidden: os, subprocess, sys, requests, pathlib, socket, shutil, psutil, file/network/OS access.
-
-Instructions for frame capture:
-- Capture a frame after each logical scene segment to help with automated evaluation.
-- Use: 
-    frame_image = self.camera.get_image()
-    frame_image.save(FRAME_DIR / "frame_name.png")
-- Place this immediately after each self.play() that introduces or transforms major objects.
-- Choose descriptive names for frames, e.g., "frame_title", "frame_vector_addition".
-- Ensure all frames are saved inside the run-specific folder (FRAME_DIR) to allow automated refinement loops to detect them.
+Forbidden: os, subprocess, sys, requests, pathlib, socket, shutil, psutil, any file/network/OS access.
 """
 
 DEFAULT_SCENE = """from manim import *
 
 class AutoScene(Scene):
     def construct(self):
-        t = Text("Hello").scale(1)
+        t = Text("Hello from Manim").scale(1)
         self.play(Write(t))
-        frame_image = self.camera.get_image()
-        frame_image.save(FRAME_DIR / "hello.png")
         self.wait(0.5)
 """
 
-# ---------- helpers ----------
+# ---------------- helpers ----------------
 def _clean_code(text: str) -> str:
-    """Strip code fences and whitespace."""
     if not text:
         return ""
-    code = re.sub(r"^```(?:python)?\n|```$", "", text, flags=re.MULTILINE).strip()
-    return code
-
-def _prepare_frame_prefixes(frame_paths: List[str], prefix_len=300):
-    """(Kept from original) Produce small b64 prefixes (not sent unless you want)."""
-    mapping = []
-    for p in frame_paths or []:
-        try:
-            b = Path(p).read_bytes()
-            b64 = base64.b64encode(b).decode("utf-8")
-            short = b64[:prefix_len]
-            mapping.append((Path(p).name, short))
-        except Exception as e:
-            print("Failed reading frame", p, e, file=sys.stderr)
-    return mapping
-
-def llm_generate_manim_code(prompt: str, previous_code: Optional[str] = None, frame_paths: Optional[List[str]] = None) -> str:
-    """Generate or refine Manim AutoScene code using Gemini (multimodal with frames)."""
-    if not client:
-        return DEFAULT_SCENE
-
-    try:
-        contents = []
-        if previous_code and frame_paths:
-            # Refinement mode (use captured frames as image Parts)
-            user_prompt = f"""Refine the following Manim CE (0.19.x) code to fix layout/overlap/readability issues.
-Keep the math concept and animations intact; make minimal, clear fixes.
-
-Original code:
-{previous_code}
-
-Notes:
-- Use the captured frames to identify overlapping objects or text outside the viewport.
-- Adjust positions, font sizes, or layout as needed.
-- Return ONLY the full corrected Python code that defines class AutoScene(Scene) and uses 'from manim import *'.
-"""
-            contents.append(user_prompt)
-            for path in frame_paths:
-                try:
-                    image_bytes = Path(path).read_bytes()
-                    contents.append(types.Part.from_bytes(image_bytes, mime_type="image/png"))
-                except Exception as e:
-                    print("Skipping frame include:", path, e, file=sys.stderr)
-        else:
-            # First generation
-            contents.append(f"Create AutoScene for: {prompt}")
-
-        response = client.models.generate_content(
-            model=MODEL,
-            config=types.GenerateContentConfig(system_instruction=SYSTEM_PROMPT),
-            contents=contents,
-        )
-        code = _clean_code(getattr(response, "text", str(response)))
-
-        if "class AutoScene" not in code:
-            code = previous_code or DEFAULT_SCENE
-        return code
-
-    except Exception:
-        print("LLM error:", file=sys.stderr)
-        traceback.print_exc()
-        return previous_code or DEFAULT_SCENE
+    # strip code fences if present
+    return re.sub(r"^```(?:python)?\\n|```$", "", text, flags=re.MULTILINE).strip()
 
-def _run_manim(scene_code: str, run_id: Optional[str] = None) -> Tuple[bytes, List[str]]:
-    """Render scene with Manim, return (mp4_bytes, captured_frame_paths)."""
+def _run_manim(scene_code: str, run_id: Optional[str] = None) -> bytes:
+    """Render with Manim (low quality, fast) and return MP4 bytes."""
     run_id = run_id or str(uuid.uuid4())[:8]
     work = RUNS / run_id; work.mkdir(parents=True, exist_ok=True)
     media = work / "media"; media.mkdir(parents=True, exist_ok=True)
     scene_path = work / "scene.py"
 
-    # Inject FRAME_DIR so user code can save frames there
-    scene_with_path = f"""
-import os
-from pathlib import Path
-FRAME_DIR = Path(r"{work.resolve()}")
-{scene_code}
-"""
-    scene_path.write_text(scene_with_path, encoding="utf-8")
+    # Write scene code
+    scene_path.write_text(scene_code, encoding="utf-8")
 
     env = os.environ.copy()
     env["PYTHONPATH"] = str(work)
@@ -158,27 +78,42 @@ FRAME_DIR = Path(r"{work.resolve()}")
     ]
     proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, env=env)
     if proc.returncode != 0:
-        print("Manim stdout/stderr:\n", proc.stdout, file=sys.stderr)
+        print("Manim stdout/stderr:\\n", proc.stdout, file=sys.stderr)
         raise RuntimeError("Manim render failed")
 
-    # Find output video
-    mp4_path = None
+    # Locate output mp4
+    mp4 = None
     for p in media.rglob(f"{run_id}.mp4"):
-        mp4_path = p; break
-    if not mp4_path:
+        mp4 = p; break
+    if not mp4:
         for p in media.rglob("*.mp4"):
-            mp4_path = p; break
-    if not mp4_path:
+            mp4 = p; break
+    if not mp4:
         raise RuntimeError("Rendered video not found")
 
-    mp4_bytes = Path(mp4_path).read_bytes()
+    return mp4.read_bytes()
 
-    # Gather any saved frames by the scene
-    frames = [str(p) for p in work.rglob("*.png")]
-    return mp4_bytes, frames
+def llm_generate_manim_code(prompt: str, previous_code: Optional[str] = None) -> str:
+    """First-pass generation (no visual refinement)."""
+    if not client:
+        return DEFAULT_SCENE
+    try:
+        response = client.models.generate_content(
+            model=MODEL,
+            config=types.GenerateContentConfig(system_instruction=SYSTEM_PROMPT),
+            contents=f"Create AutoScene for: {prompt}",
+        )
+        code = _clean_code(getattr(response, "text", str(response)))
+        if "class AutoScene" not in code:
+            code = previous_code or DEFAULT_SCENE
+        return code
+    except Exception:
+        print("LLM generate error:", file=sys.stderr)
+        traceback.print_exc()
+        return previous_code or DEFAULT_SCENE
 
 def llm_refine_from_error(previous_code: str, error_message: str) -> str:
-    """Refine code using error text if Manim failed."""
+    """Only used when Manim fails; no frames / visual step."""
     if not client:
         return previous_code or DEFAULT_SCENE
     try:
@@ -187,14 +122,14 @@ def llm_refine_from_error(previous_code: str, error_message: str) -> str:
 Code:
 {previous_code}
 
-Error message / stack trace:
+Error / stack trace:
 {error_message}
 
 Task:
 - Fix the bug while preserving the math logic and animations.
-- Common mistakes: using .animate inside layout methods, invalid method chaining, etc.
+- Common mistakes: invalid method chaining, wrong .animate usage, bad mobject usage, etc.
 - Ensure the corrected code defines exactly one class AutoScene(Scene).
-- Return ONLY the corrected Python code.
+- Return ONLY the corrected Python code (no backticks).
 """
         response = client.models.generate_content(
             model=MODEL,
@@ -206,45 +141,26 @@ Task:
             return previous_code or DEFAULT_SCENE
         return code
     except Exception:
-        print("LLM error during error refinement:", file=sys.stderr)
+        print("LLM refine error:", file=sys.stderr)
         traceback.print_exc()
         return previous_code or DEFAULT_SCENE
 
-def refine_loop(prompt: str, max_iters: int = 2, max_error_refines: int = 1) -> bytes:
-    """Pipeline: generate → render → visual refine (frames) → final render; returns final MP4 bytes."""
+def refine_loop(prompt: str, max_error_refines: int = 1) -> bytes:
+    """Generate → render; if render fails, LLM-refine once → re-render."""
     code = llm_generate_manim_code(prompt)
-    last_video_bytes: Optional[bytes] = None
-
-    for i in range(max_iters):
-        try:
-            video_bytes, frames = _run_manim(code, run_id=f"iter{i}")
-            last_video_bytes = video_bytes
-        except Exception:
-            print("Render error (iter):", file=sys.stderr)
-            traceback.print_exc()
-
-            if i < max_error_refines:
-                error_msg = traceback.format_exc()
-                code = llm_refine_from_error(previous_code=code, error_message=error_msg)
-                continue
-            else:
-                break
-
-        if not frames:
-            print("No captured frames detected; stopping refinement loop.")
-            break
-
-        # Visual refinement with frames
-        code = llm_generate_manim_code(prompt, previous_code=code, frame_paths=frames)
-
-    # Final render
     try:
-        final_video, _ = _run_manim(code, run_id="final")
-        return final_video
+        return _run_manim(code, run_id="iter0")
     except Exception:
-        return last_video_bytes or b""
+        print("Render failed (iter0), attempting error-based refinement...", file=sys.stderr)
+        traceback.print_exc()
+        if max_error_refines <= 0:
+            raise
 
-# ---------- API models & routes ----------
+        refined = llm_refine_from_error(previous_code=code, error_message=traceback.format_exc())
+        # Try a second render
+        return _run_manim(refined, run_id="iter1")
+
+# ---------------- API ----------------
 class PromptIn(BaseModel):
     prompt: str
 
@@ -254,16 +170,14 @@ def health():
 
 @app.post("/generate-and-render")
 def generate_and_render(inp: PromptIn):
-    """
-    Send: {"prompt": "..."}  ->  Returns: MP4 bytes (video/mp4)
-    """
     if not inp.prompt or not inp.prompt.strip():
         raise HTTPException(400, "Missing prompt")
-    video = refine_loop(inp.prompt.strip(), max_iters=2, max_error_refines=1)
-    if not video:
-        raise HTTPException(500, "Failed to produce video")
+    try:
+        mp4 = refine_loop(inp.prompt.strip(), max_error_refines=1)
+    except Exception:
+        raise HTTPException(500, "Failed to produce video after refinement")
     return Response(
-        content=video,
+        content=mp4,
         media_type="video/mp4",
         headers={"Content-Disposition": 'inline; filename="result.mp4"'}
     )