Spaces:
Sleeping
Sleeping
| """ | |
| Code Post-Processor for Manim | |
| Fixes common issues in AI-generated Manim code, particularly around | |
| SurroundingRectangle with indexed MathTex elements. | |
| """ | |
| import re | |
| from typing import List | |
| def fix_surrounding_rectangles(code: str) -> str: | |
| """ | |
| Fix or remove SurroundingRectangle calls that use indexed access to MathTex. | |
| Problem: AI often generates code like: | |
| SurroundingRectangle(equation[0][5], ...) | |
| This doesn't work reliably because MathTex indexing is unpredictable. | |
| Solution: Comment out these problematic lines with an explanation. | |
| Args: | |
| code: Raw generated Manim code | |
| Returns: | |
| Fixed code with problematic SurroundingRectangle calls commented out | |
| """ | |
| lines = code.split('\n') | |
| fixed_lines = [] | |
| skip_next = False # Initialize before the loop | |
| for line in lines: | |
| # Pattern: SurroundingRectangle with indexed access like equation[0][5] | |
| # Matches: SurroundingRectangle(something[X][Y], ...) or SurroundingRectangle(something[X:Y], ...) | |
| if 'SurroundingRectangle(' in line: | |
| # Check if it has indexed or sliced access | |
| # Pattern: variable[number] or variable[number:number] inside SurroundingRectangle() | |
| if re.search(r'SurroundingRectangle\([^,\)]*\[\d+\]', line) or \ | |
| re.search(r'SurroundingRectangle\([^,\)]*\[\d+:\d*\]', line): | |
| # Replace with empty VGroup to prevent NameError in subsequent animations | |
| # while avoiding the rendering crash | |
| indent = len(line) - len(line.lstrip()) | |
| # Extract variable name if present (e.g. "box = SurroundingRectangle(...)") | |
| var_name = "" | |
| if "=" in line: | |
| var_name = line.split("=")[0].strip() | |
| replacement = ' ' * indent + f'{var_name} = VGroup() # Replaced invalid SurroundingRectangle' | |
| else: | |
| # If no assignment, just comment it out as it won't be referenced | |
| replacement = ' ' * indent + '# ' + line.lstrip() + ' # Auto-disabled: indexed SurroundingRectangle' | |
| fixed_lines.append(replacement) | |
| else: | |
| fixed_lines.append(line) | |
| else: | |
| fixed_lines.append(line) | |
| return '\n'.join(fixed_lines) | |
| def remove_problematic_indexing(code: str) -> str: | |
| """ | |
| More aggressive approach: Remove entire blocks that use indexed MathTex highlighting. | |
| This removes the entire voiceover block if it contains indexed SurroundingRectangle. | |
| """ | |
| # For now, use the commenting approach which is safer | |
| return fix_surrounding_rectangles(code) | |
| def fix_undefined_colors(code: str) -> str: | |
| """ | |
| Fix undefined color constants by replacing them with valid Manim colors. | |
| Common issues: | |
| - ORANGE_A, ORANGE_B, etc. -> ORANGE | |
| - RED_A, RED_B, etc. -> RED | |
| - Similar patterns for other colors | |
| Args: | |
| code: Raw generated code | |
| Returns: | |
| Code with undefined colors replaced | |
| """ | |
| # Color mappings: undefined variants -> standard colors | |
| color_replacements = { | |
| # Orange variants | |
| r'\bORANGE_[A-Z]\b': 'ORANGE', | |
| # Red variants | |
| r'\bRED_[A-Z]\b': 'RED', | |
| # Blue variants | |
| r'\bBLUE_[A-Z]\b': 'BLUE', | |
| # Green variants | |
| r'\bGREEN_[A-Z]\b': 'GREEN', | |
| # Yellow variants | |
| r'\bYELLOW_[A-Z]\b': 'YELLOW', | |
| # Purple variants | |
| r'\bPURPLE_[A-Z]\b': 'PURPLE', | |
| # Pink variants | |
| r'\bPINK_[A-Z]\b': 'PINK', | |
| # Teal variants | |
| r'\bTEAL_[A-Z]\b': 'TEAL', | |
| # Gray variants | |
| r'\bGRAY_[A-Z]\b': 'GRAY', | |
| } | |
| for pattern, replacement in color_replacements.items(): | |
| code = re.sub(pattern, replacement, code) | |
| return code | |
| def clean_duplicate_imports(code: str) -> str: | |
| """ | |
| Remove duplicate and incorrect imports. | |
| Keep only the correct imports for our architecture. | |
| """ | |
| lines = code.split('\n') | |
| cleaned_lines = [] | |
| seen_imports = set() | |
| # Imports to remove (old/wrong patterns) | |
| bad_imports = [ | |
| 'from manim_voiceover import VoiceoverScene', | |
| 'from manimator.services import ElevenLabsService', | |
| 'from manim_voiceover.services.gtts import GTTSService', | |
| 'from manim_voiceover.services.elevenlabs import ElevenLabsService', | |
| ] | |
| for line in lines: | |
| stripped = line.strip() | |
| # Skip bad imports | |
| if any(bad_import in line for bad_import in bad_imports): | |
| continue | |
| # Track and deduplicate good imports | |
| if stripped.startswith('from ') or stripped.startswith('import '): | |
| if stripped not in seen_imports: | |
| seen_imports.add(stripped) | |
| cleaned_lines.append(line) | |
| else: | |
| cleaned_lines.append(line) | |
| return '\n'.join(cleaned_lines) | |
| def fix_invalid_manim_parameters(code: str) -> str: | |
| """ | |
| Fix invalid Manim parameters that cause runtime errors. | |
| Common issues: | |
| - rounded_corners parameter in Rectangle (doesn't exist) | |
| - corner_radius parameter in Rectangle (doesn't exist in older Manim versions) | |
| - Invalid parameter names | |
| """ | |
| # Remove corner_radius from Rectangle/RoundedRectangle calls (line by line) | |
| lines = code.split('\n') | |
| fixed_lines = [] | |
| for line in lines: | |
| # Skip lines that are just the parameter | |
| if 'corner_radius' in line and '=' in line: | |
| # Check if this is a standalone parameter line | |
| stripped = line.strip() | |
| if stripped.startswith('corner_radius=') or stripped.startswith('rounded_corners='): | |
| # Skip this line entirely | |
| continue | |
| # Remove inline corner_radius/rounded_corners parameters | |
| line = re.sub(r',\s*corner_radius\s*=\s*[^,)]+', '', line) | |
| line = re.sub(r',\s*rounded_corners\s*=\s*[^,)]+', '', line) | |
| # Remove invalid scale_tips parameter (often hallucinated for scale() method) | |
| if 'scale_tips' in line and '=' in line: | |
| stripped = line.strip() | |
| if stripped.startswith('scale_tips='): | |
| continue | |
| # Handle inline scale_tips removal (both with and without leading comma) | |
| line = re.sub(r',\s*scale_tips\s*=\s*[^,)]+', '', line) | |
| line = re.sub(r'\bscale_tips\s*=\s*[^,)]+\s*,?', '', line) | |
| fixed_lines.append(line) | |
| return '\n'.join(fixed_lines) | |
| def fix_visual_leaks(code: str) -> str: | |
| """ | |
| Prevent 'visual memory leaks' where reassigning a variable leaves the old Mobject on screen. | |
| Also enforces 'Title Exclusivity' to prevent multiple titles from stacking. | |
| Problem 1 (Variable Reuse): | |
| text = Text("A") | |
| self.play(Write(text)) | |
| text = Text("B") # "A" is still on screen! | |
| Problem 2 (Title Stacking): | |
| title1 = Text("Intro") | |
| self.play(Write(title1)) | |
| title2 = Text("Chapter 1") # "Intro" is still on screen! | |
| Solution: | |
| Inject `self.remove(var)` checks. | |
| """ | |
| lines = code.split('\n') | |
| fixed_lines = [] | |
| # Regex to detect Mobject assignments | |
| # Matches: var = ClassName(...) | |
| # We broaden this to catch any Capitalized class instantiation to be safer | |
| assignment_pattern = re.compile(r'^\s*([a-zA-Z_]\w*)\s*=\s*([A-Z]\w*)\(') | |
| # Track variables that have been assigned | |
| assigned_vars = set() | |
| # Track variables that look like titles | |
| active_titles = [] | |
| for line in lines: | |
| match = assignment_pattern.match(line) | |
| if match: | |
| var_name = match.group(1) | |
| class_name = match.group(2) | |
| indent = len(line) - len(line.lstrip()) | |
| indent_str = ' ' * indent | |
| # 1. Handle Variable Reuse (Same variable name) | |
| # We inject this for EVERY assignment to handle loops correctly. | |
| # In a loop, the line `t = Text(...)` appears once but runs multiple times. | |
| # By injecting the check, we ensure the previous iteration's object is removed. | |
| removal = f"{indent_str}if '{var_name}' in locals(): self.remove({var_name})" | |
| fixed_lines.append(removal) | |
| # 2. Handle Title Exclusivity (Different variable names, but both are titles) | |
| # Check if this looks like a title (contains 'title', 'header', 'heading') | |
| # But exclude 'subtitle' or 'sub_title' | |
| is_title = re.search(r'title|header|heading', var_name, re.IGNORECASE) and \ | |
| not re.search(r'sub', var_name, re.IGNORECASE) | |
| if is_title: | |
| # If we are creating a new title, remove ALL previous titles | |
| for old_title in active_titles: | |
| # Don't remove if it's the same variable (handled above) | |
| if old_title != var_name: | |
| removal = f"{indent_str}if '{old_title}' in locals(): self.remove({old_title})" | |
| fixed_lines.append(removal) | |
| # Add this to active titles (if not already there) | |
| if var_name not in active_titles: | |
| active_titles.append(var_name) | |
| assigned_vars.add(var_name) | |
| fixed_lines.append(line) | |
| else: | |
| fixed_lines.append(line) | |
| return '\n'.join(fixed_lines) | |
| def post_process_code(code: str) -> str: | |
| """ | |
| Main entry point for code post-processing. | |
| Applies all fixes to AI-generated Manim code. | |
| Args: | |
| code: Raw generated code | |
| Returns: | |
| Cleaned and fixed code | |
| """ | |
| # Check if we need to add header (before making changes) | |
| has_undefined_colors = bool(re.search(r'\b(ORANGE|RED|BLUE|GREEN|YELLOW|PURPLE|PINK|TEAL|GRAY)_[A-Z]\b', code)) | |
| # Apply fixes in order | |
| code = clean_duplicate_imports(code) | |
| code = fix_undefined_colors(code) | |
| code = fix_invalid_manim_parameters(code) | |
| code = fix_invalid_manim_parameters(code) | |
| code = fix_surrounding_rectangles(code) | |
| code = fix_visual_leaks(code) | |
| # Add header comment explaining post-processing | |
| header = """# NOTE: This code has been automatically post-processed to fix common issues. | |
| # Indexed SurroundingRectangle calls have been disabled as they don't reliably | |
| # highlight the intended equation parts in MathTex objects. | |
| # Undefined color constants have been replaced with standard Manim colors. | |
| # Invalid Manim parameters have been removed or corrected. | |
| """ | |
| # Only add header if we actually made changes | |
| if '# Auto-disabled:' in code or has_undefined_colors or 'rounded_corners' in code: | |
| code = header + code | |
| return code | |
| def validate_code_structure(code: str) -> List[str]: | |
| """ | |
| Validate the generated code for common issues. | |
| Returns: | |
| List of warning messages (empty if no issues) | |
| """ | |
| warnings = [] | |
| # Check for common issues | |
| if 'SurroundingRectangle(' in code: | |
| if re.search(r'SurroundingRectangle\([^,\)]*\[\d+\]', code): | |
| warnings.append("Code contains indexed SurroundingRectangle calls (will be auto-fixed)") | |
| if 'from manim_voiceover.services.gtts import GTTSService' in code: | |
| warnings.append("Code still uses deprecated GTTSService (should use ElevenLabsService)") | |
| return warnings | |