Upload app.py with huggingface_hub

app.py

@@ -15,7 +15,7 @@ from typing import Any
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
 
 # Force pandas to fully initialize before transitive imports cause circular import
-import pandas  # noqa: F401
+import pandas  # noqa: F401, I001
 
 import gradio as gr
 import spaces
@@ -192,7 +192,8 @@ def generate_spec(
         )
 
         generated = outputs[0][input_len:]
-        return tokenizer.decode(generated, skip_special_tokens=True)
+        spec = tokenizer.decode(generated, skip_special_tokens=True)
+        return render_mermaid_images(spec)
 
     except Exception as e:
         return (
@@ -204,28 +205,112 @@ def generate_spec(
 # --- Gradio UI ---
 
 
+def render_mermaid_images(spec: str) -> str:
+    """Replace mermaid code blocks with rendered SVG images via mermaid.ink.
+
+    Validates the mermaid syntax first, strips malformed blocks, and
+    converts valid ones to inline images that render reliably regardless
+    of CSS/theme issues.
+    """
+    import base64
+    import re
+
+    valid_starts = (
+        "graph ",
+        "graph\n",
+        "flowchart ",
+        "flowchart\n",
+        "classDiagram",
+        "sequenceDiagram",
+        "stateDiagram",
+        "erDiagram",
+        "gantt",
+        "pie",
+        "gitgraph",
+    )
+
+    def replace_block(match: re.Match) -> str:
+        block = match.group(1).strip()
+        # Must start with a valid diagram type
+        if not any(block.startswith(s) for s in valid_starts):
+            return "*[Mermaid diagram removed: unrecognized diagram type]*"
+        # Check balanced brackets/braces
+        if block.count("[") != block.count("]"):
+            return "*[Mermaid diagram removed: unbalanced brackets]*"
+        if block.count("{") != block.count("}"):
+            return "*[Mermaid diagram removed: unbalanced braces]*"
+        if block.count("(") != block.count(")"):
+            return "*[Mermaid diagram removed: unbalanced parentheses]*"
+        # Encode and return as mermaid.ink image
+        encoded = base64.urlsafe_b64encode(block.encode("utf-8")).decode("ascii")
+        return f"![Dependency Diagram](https://mermaid.ink/svg/{encoded})"
+
+    return re.sub(r"```mermaid\s*\n(.*?)```", replace_block, spec, flags=re.DOTALL)
+
+
 def create_app():
     mermaid_css = """
+    /* Aggressive mermaid text fix — target all SVG text elements */
+    svg[id^="mermaid"] text,
+    svg[id^="mermaid"] tspan,
+    svg[id^="mermaid"] .nodeLabel,
+    svg[id^="mermaid"] .edgeLabel,
+    svg[id^="mermaid"] .label,
+    svg[id^="mermaid"] .cluster-label,
+    .mermaid text,
+    .mermaid tspan,
+    .mermaid .nodeLabel,
+    .mermaid .edgeLabel,
+    .mermaid .label,
+    [data-mermaid] text,
+    [data-mermaid] tspan {
+        color: #1a1a1a !important;
+        fill: #1a1a1a !important;
+    }
+    svg[id^="mermaid"] .node rect,
+    svg[id^="mermaid"] .node polygon,
+    svg[id^="mermaid"] .node circle,
     .mermaid .node rect,
     .mermaid .node polygon,
     .mermaid .node circle {
         fill: #e8f0fe !important;
         stroke: #4a6fa5 !important;
     }
-    .mermaid .nodeLabel,
-    .mermaid .edgeLabel,
-    .mermaid text {
-        color: #1a1a1a !important;
-        fill: #1a1a1a !important;
-    }
+    svg[id^="mermaid"] .edgePath .path,
     .mermaid .edgePath .path {
         stroke: #4a6fa5 !important;
     }
+    svg[id^="mermaid"] .cluster rect,
+    .mermaid .cluster rect {
+        fill: #f0f4ff !important;
+        stroke: #4a6fa5 !important;
+    }
+    """
+
+    mermaid_js = """
+    () => {
+        if (window.mermaid) {
+            window.mermaid.initialize({
+                theme: 'default',
+                themeVariables: {
+                    primaryColor: '#e8f0fe',
+                    primaryTextColor: '#1a1a1a',
+                    primaryBorderColor: '#4a6fa5',
+                    lineColor: '#4a6fa5',
+                    secondaryColor: '#f0f4ff',
+                    tertiaryColor: '#fff',
+                    nodeTextColor: '#1a1a1a',
+                    edgeLabelBackground: '#ffffff',
+                }
+            });
+        }
+    }
     """
 
     with gr.Blocks(
         title="CodeWraith - Module-to-Spec Transformer",
         css=mermaid_css,
+        js=mermaid_js,
     ) as app:
         gr.Markdown(
             "# CodeWraith\n"