Tune generation speed and timeouts

Co-authored-by: Codex <codex@openai.com>

app.py

@@ -385,7 +385,7 @@ def get_agent():
     return agent
 
 
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=600)
 def run_visual_lab(prompt: str, retries: int) -> Tuple[str | None, str, str]:
     try:
         visual_agent = get_agent()
@@ -404,8 +404,41 @@ def run_visual_lab(prompt: str, retries: int) -> Tuple[str | None, str, str]:
 def load_example(example: str) -> str:
     examples = {
         "Fourier": "Visualize how a square wave emerges from the first seven Fourier sine components, with each partial sum drawn in neon colors.",
-        "Orbit": "Show a satellite in elliptical orbit around Earth with velocity and gravitational force vectors changing over time.",
+        "Orbit": """Create a beautiful, scientifically accurate 2D physics animation of a satellite in an elliptical orbit around Earth using ManimCE.
+
+Requirements:
+1. Scaling & Geometry (CRITICAL to avoid overlapping):
+- Represent Earth as a clear sphere or circle in the center, but keep its radius small (e.g., radius=0.6) so it doesn't swallow the orbit.
+- Create an explicit, highly elongated elliptical orbit path (`Ellipse(width=6.0, height=3.5)`) shifted slightly so that Earth sits exactly at one of the focal points (Foci) of the ellipse, NOT at the geometric center (Kepler's First Law).
+
+2. Dynamic Vectors & Updaters:
+- Represent the satellite as a distinct, colored Dot (e.g., Cyber Cyan) moving along the elliptical path using a custom tracker or `ValueTracker` for the orbital angle.
+- Create two dynamic arrows (`Vector` or `Arrow`) attached to the satellite dot:
+  * Gravitational Force Vector (Color: Neon Red): Must always point directly from the satellite's current position to Earth's center. Its length must dynamically increase when close to Earth and decrease when far away (Inverse-square law).
+  * Velocity Vector (Color: Bright Green): Must always be perfectly tangent to the elliptical orbit path in the direction of motion. Its length must dynamically represent orbital speed (faster at perigee, slower at apogee).
+- Use `add_updater` on both vectors so their positions, directions, and lengths update smoothly at every single frame based on the satellite's position.
+
+3. Labels & Overlay:
+- Add text labels for "Velocity (v)" and "Gravity (Fg)" matching the vector colors.
+- In the top-right corner, add a LaTeX mathematical text displaying Newton's Law of Universal Gravitation: "F_g = G * (M_1 * M_2) / r²".
+
+4. Animation flow: Run the animation for 2 complete orbital periods (about 10-12 seconds) so the viewer can clearly observe the dramatic speeding up at the close approach and slowing down at the far end. Wrap everything inside a single clean Scene class.""",
         "Eigen": "Explain eigenvectors of a 2D linear transformation using a grid, two vectors, and the transformed plane.",
+        "Simple Pendulum": """Create a high-quality Physics simulation of a Damped Simple Pendulum using ManimCE. 
+
+Requirements:
+1. Physics Setup: Define explicit variables for gravitational acceleration (g=9.81), rod length (L=3.0), initial angle (theta = pi/4), and a damping coefficient (b=0.15) to simulate real-world air resistance. Use a clear numerical integration method (like Euler-Cromer) inside an object updater (`add_updater`) to dynamically recalculate the angular velocity and displacement at every frame (`dt`).
+
+2. Visual Elements (Strict Coordinate Updates):
+- A fixed ceiling line or pivot dot at the top center.
+- A colored dot (e.g., Cyber Cyan or Neon Pink) representing the heavy bob. Make it big enough for viewers to see clearly.
+- A clean line representing the pendulum rod that attaches the pivot to the bob. CRITICAL: You must attach an updater to this rod using `add_updater` so that its end position dynamically calls `rod.put_start_and_end_on(pivot.get_center(), bob.get_center())` at every single frame. The rod must always move dynamically with the bob.
+- Add an elegant fading trace/trail (`TracedPath`) attached to the bob to visually map its decaying sinusoidal path across the screen over time.
+
+3. Mathematical Overlay: In the upper left corner, display the differential equation governing the motion: "d²θ/dt² + (b/m)dθ/dt + (g/L)sin(θ) = 0" rendered beautifully via LaTeX.
+
+4. Animation flow: Let the simulation run smoothly for 12 seconds to clearly demonstrate the kinetic energy converting to thermal energy as the oscillation slowly dampens to a complete halt. Ensure the code is self-contained and wrapped inside a single executable Scene class.
+"""
     }
     return examples.get(example, "")
 
@@ -451,7 +484,7 @@ with gr.Blocks(css=custom_css, title="SciVisual-Agent") as demo:
                 clear_btn = gr.Button("Clear", elem_id="clear-btn")
 
             example_choice = gr.Radio(
-                choices=["Fourier", "Orbit", "Eigen"],
+                choices=["Simple Pendulum", "Orbit", "Eigen"],
                 label="Quick Load",
                 value=None,
                 interactive=True,

executor.py

@@ -47,7 +47,7 @@ def _manim_command_prefix() -> list[str]:
 
 def render_code(
     code: str,
-    timeout_seconds: int = 300,
+    timeout_seconds: int = 600,
     output_dir: Path = OUTPUT_DIR,
     quality: str = "-ql",
 ) -> RenderResult:
@@ -137,7 +137,7 @@ def render_code(
             Path(temp_script_name).unlink(missing_ok=True)
 
 
-def render_manim_scene(code: str, timeout_seconds: int = 300) -> Tuple[bool, str]:
+def render_manim_scene(code: str, timeout_seconds: int = 600) -> Tuple[bool, str]:
     """Compatibility wrapper returning the tuple shape used by the app."""
     result = render_code(code, timeout_seconds=timeout_seconds)
     if result.success and result.video_path:

generator.py

@@ -33,7 +33,7 @@ class LocalModelConfig:
     prompt_mode: str = "chat"
     device_map: str = "auto"
     load_in_4bit: bool = True
-    max_new_tokens: int = 4096
+    max_new_tokens: int = 8192
     temperature: float = 0.0
     top_p: float = 0.9
     timeout_seconds: int = 300
@@ -210,6 +210,9 @@ class LocalManimModel:
         if not self.ready:
             raise RuntimeError(self.error or "Local model is not ready.")
 
+        import torch
+        from transformers import StoppingCriteriaList
+
         if self.no_system_role_support and len(messages) >= 2 and messages[0]["role"] == "system":
             messages = [{"role": "user", "content": messages[0]["content"] + "\n\n" + messages[1]["content"]}]
 
@@ -227,6 +230,9 @@ class LocalManimModel:
             "pad_token_id": self.tokenizer.pad_token_id,
             "eos_token_id": self.tokenizer.eos_token_id,
             "use_cache": True,
+            "stopping_criteria": StoppingCriteriaList(
+                [_StopOnTokenSequence(self.tokenizer.encode("</CODE>", add_special_tokens=False))]
+            ),
         }
         if self.config.temperature and self.config.temperature > 0:
             generate_kwargs.update(
@@ -239,7 +245,8 @@ class LocalManimModel:
         else:
             generate_kwargs["do_sample"] = False
 
-        generated_ids = self.model.generate(**generate_kwargs)
+        with torch.inference_mode():
+            generated_ids = self.model.generate(**generate_kwargs)
         output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()
 
         try:
@@ -276,6 +283,22 @@ class LocalManimModel:
         return "\n\n".join(chunks)
 
 
+class _StopOnTokenSequence:
+    """Stop generation when a specific token suffix appears."""
+
+    def __init__(self, stop_token_ids: List[int]) -> None:
+        self.stop_token_ids = stop_token_ids
+
+    def __call__(self, input_ids, scores, **kwargs) -> bool:
+        del scores, kwargs
+        if not self.stop_token_ids:
+            return False
+        if input_ids.shape[-1] < len(self.stop_token_ids):
+            return False
+        tail = input_ids[0, -len(self.stop_token_ids) :].tolist()
+        return tail == self.stop_token_ids
+
+
 class ManimVisualAgent:
     """Generate Manim code with local inference and repair failures."""
 
@@ -426,9 +449,9 @@ def build_default_agent() -> ManimVisualAgent:
         prompt_mode=os.environ.get("SCIVISUAL_PROMPT_MODE", "chat"),
         device_map=os.environ.get("SCIVISUAL_DEVICE_MAP", "auto"),
         load_in_4bit=os.environ.get("SCIVISUAL_LOAD_IN_4BIT", "1").lower() not in ("0", "false", "no"),
-        max_new_tokens=int(os.environ.get("SCIVISUAL_MAX_NEW_TOKENS", "4096")),
+        max_new_tokens=int(os.environ.get("SCIVISUAL_MAX_NEW_TOKENS", "8192")),
         temperature=float(os.environ.get("SCIVISUAL_TEMPERATURE", "0")),
         top_p=float(os.environ.get("SCIVISUAL_TOP_P", "0.9")),
-        timeout_seconds=int(os.environ.get("SCIVISUAL_RENDER_TIMEOUT", "300")),
+        timeout_seconds=int(os.environ.get("SCIVISUAL_RENDER_TIMEOUT", "600")),
     )
     return ManimVisualAgent(config=config)