Update app.py

app.py

@@ -1,11 +1,28 @@
+import os
+import sys
 import re
+import json
+import time
+import threading
+import subprocess
 from pathlib import Path
-from markdown_it import MarkdownIt
-from importlib import import_module
 
-def _make_md():
+import gradio as gr
+import pandas as pd
+import torch
+import spaces
+
+# ---------------------------
+# Markdown rendering (Option A)
+# ---------------------------
+
+def _make_md_markdownit():
+    # Prefer markdown-it-py + mdit-py-plugins if available
+    from importlib import import_module
+    from markdown_it import MarkdownIt
     md = MarkdownIt("gfm-like")
 
+    # Version-agnostic plugin shims
     foot_mod = import_module("mdit_py_plugins.footnote")
     foot = getattr(foot_mod, "footnote", None) or getattr(foot_mod, "footnote_plugin")
     md.use(foot)
@@ -17,42 +34,112 @@ def _make_md():
     cont_mod = import_module("mdit_py_plugins.container")
     container = getattr(cont_mod, "container", None) or getattr(cont_mod, "container_plugin")
     try:
-        md.use(container, "details")            # newer signature
+        md.use(container, "details")
     except TypeError:
-        md.use(lambda m: container(m, name="details"))  # older signature
-
+        md.use(lambda m: container(m, name="details"))
     return md
 
-_md = _make_md()
+def _make_md_pythonmarkdown():
+    # Fallback: Python-Markdown + PyMdown
+    import markdown as md
+    exts = [
+        "extra",            # tables + fenced code
+        "footnotes",
+        "admonition",
+        "toc",
+        "pymdownx.details",
+        "pymdownx.superfences",
+        "pymdownx.tasklist",
+    ]
+    ext_cfg = {"pymdownx.tasklist": {"custom_checkbox": True}, "toc": {"permalink": True}}
+    return ("python-markdown", exts, ext_cfg, md)
 
+try:
+    _md_engine = ("markdown-it", _make_md_markdownit())
+except Exception:
+    _md_engine = _make_md_pythonmarkdown()
 
-
-def md_to_html(text: str) -> str:
-    # Convert common Obsidian patterns to standard Markdown
-    text = re.sub(r'!\[\[([^\]|]+)\]\]', r'![](static/\1)', text)     # image embeds ![[file.png]]
+def _obsidian_rewrites(text: str) -> str:
+    # Obsidian image/file embeds and wiki links
+    text = re.sub(r'!\[\[([^\]|]+)\]\]', r'![](static/\1)', text)     # ![[file.png]]
     text = re.sub(r'\[\[([^\]|]+)\|([^\]]+)\]\]', r'[\2](\1)', text)   # [[file|label]]
     text = re.sub(r'\[\[([^\]]+)\]\]', r'[\1](\1)', text)              # [[file]]
-    return _md.render(text)
+    return text
+
+def md_to_html(text: str) -> str:
+    text = _obsidian_rewrites(text)
+    if _md_engine[0] == "markdown-it":
+        md = _md_engine[1]
+        return md.render(text)
+    else:
+        tag, exts, cfg, md = _md_engine
+        return md.markdown(text, extensions=exts, extension_configs=cfg, output_format="html5")
 
 def render_article(md_path: str, inserts: dict[str, callable]):
-    raw = Path(md_path).read_text(encoding="utf-8")
-    parts = re.split(r"\{\{([A-Z_]+)\}\}", raw)  # split on {{TOKEN}}
+    raw = ""
+    path = Path(md_path)
+    if path.exists():
+        raw = path.read_text(encoding="utf-8")
+    else:
+        raw = f"**Missing article**: `{md_path}` not found.\n\nCreate it in your Space repo."
+
+    # Split on {{TOKEN}} markers (e.g., {{ALLOC_PLOT}})
+    parts = re.split(r"\{\{([A-Z_]+)\}\}", raw)
     with gr.Column():
         for i, part in enumerate(parts):
             if i % 2 == 0:
                 gr.HTML(md_to_html(part))
             else:
                 build = inserts.get(part)
-                (build or (lambda: gr.HTML(f"<p><em>Unknown insert: {part}</em></p>")))()
+                if build is None:
+                    gr.HTML(f"<p><em>Unknown insert: {part}</em></p>")
+                else:
+                    build()
+
+# ---------------------------
+# Terminal (safe, simplified)
+# ---------------------------
+
+def run_shell(cmd: str) -> str:
+    banned = ["|", ">", "<", "&&", "||", "`"]
+    if any(b in cmd for b in banned):
+        return "$ " + cmd + "\nBlocked characters. Use a single command."
+    try:
+        p = subprocess.run(cmd, shell=True, check=False, capture_output=True, text=True, timeout=30)
+        return f"$ {cmd}\n{p.stdout}{p.stderr}"
+    except Exception as e:
+        return f"$ {cmd}\n{e!r}"
 
-# --- Builders that drop your existing widgets in-place ---
 def build_terminal():
     with gr.Group():
         cmd = gr.Textbox(label="Command", value="python -c 'import torch; print(torch.__version__)'")
         run = gr.Button("Run")
-        out = gr.Textbox(label="Output", lines=12)
+        out = gr.Textbox(label="Output", lines=12, interactive=False)
         run.click(run_shell, inputs=cmd, outputs=out)
 
+# ---------------------------------------
+# Attention Mask Visualizer (Transformers)
+# ---------------------------------------
+
+def _import_attention_visualizer():
+    try:
+        from transformers.utils.attention_visualizer import AttentionMaskVisualizer  # type: ignore
+    except Exception as e:
+        raise RuntimeError(
+            "AttentionMaskVisualizer is unavailable in this Transformers version."
+        ) from e
+    return AttentionMaskVisualizer
+
+@spaces.GPU(duration=120)
+def render_attention_mask(model_id: str, prompt: str) -> str:
+    try:
+        AttentionMaskVisualizer = _import_attention_visualizer()
+        vis = AttentionMaskVisualizer(model_id)
+        out = vis(prompt)  # returns embeddable HTML or object with _repr_html_
+        return str(out)
+    except Exception as e:
+        return f"<p>Attention visualizer error: {e}</p>"
+
 def build_attn_vis():
     with gr.Group():
         with gr.Row():
@@ -67,6 +154,69 @@ def build_attn_vis():
         html = gr.HTML()
         go.click(render_attention_mask, inputs=[model, prompt], outputs=html)
 
+# -------------------------------------------------------
+# Transformers caching allocator warmup (time vs MiB plot)
+# -------------------------------------------------------
+
+from transformers import AutoModelForCausalLM, modeling_utils as MU  # noqa: E402
+
+def _measure_load_timeline(model_id: str, disable_warmup: bool):
+    orig = getattr(MU, "caching_allocator_warmup", None)
+    if disable_warmup and orig is not None:
+        MU.caching_allocator_warmup = lambda *a, **k: None  # type: ignore[attr-defined]
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        tl = []
+
+        def sample(start_t, stop_evt):
+            while not stop_evt.is_set():
+                if device == "cuda":
+                    torch.cuda.synchronize()
+                    alloc = torch.cuda.memory_allocated()
+                else:
+                    alloc = 0
+                tl.append({"t": time.perf_counter() - start_t, "MiB": alloc / (1024**2)})
+                time.sleep(0.05)
+
+        if device == "cuda":
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
+
+        start = time.perf_counter()
+        stop_evt = threading.Event()
+        th = threading.Thread(target=sample, args=(start, stop_evt), daemon=True)
+        th.start()
+
+        kwargs = {}
+        if device == "cuda":
+            kwargs.update(dict(torch_dtype=torch.float16, device_map="cuda:0", low_cpu_mem_usage=True))
+        model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
+
+        stop_evt.set()
+        th.join()
+
+        if device == "cuda":
+            torch.cuda.synchronize()
+            tl.append({"t": time.perf_counter() - start, "MiB": torch.cuda.memory_allocated() / (1024**2)})
+
+        del model
+        if device == "cuda":
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+        return tl
+    finally:
+        if orig is not None:
+            MU.caching_allocator_warmup = orig  # restore
+
+@spaces.GPU(duration=240)
+def profile_warmup(model_id: str):
+    on  = _measure_load_timeline(model_id, disable_warmup=False)
+    off = _measure_load_timeline(model_id, disable_warmup=True)
+    rows = [{"t": r["t"], "MiB": r["MiB"], "mode": "warmup ON"} for r in on] + \
+           [{"t": r["t"], "MiB": r["MiB"], "mode": "warmup OFF"} for r in off]
+    return pd.DataFrame(rows)
+
 def build_alloc_plot():
     with gr.Group():
         model = gr.Dropdown(
@@ -78,12 +228,77 @@ def build_alloc_plot():
         go = gr.Button("Run")
         plot = gr.LinePlot(
             x="t", y="MiB", color="mode", overlay_point=True,
-            title="from_pretrained(): time vs CUDA memory_allocated()", width=900, height=420
+            title="from_pretrained() load: time vs CUDA memory_allocated()",
+            tooltip=["t", "MiB", "mode"], width=900, height=420
         )
         go.click(profile_warmup, inputs=[model], outputs=plot)
 
+# ---------------------------
+# Optional FastRTC preview
+# ---------------------------
+
+try:
+    from fastrtc import WebRTC, ReplyOnPause  # type: ignore
+    def _echo_video(frame):
+        yield frame
+    HAS_FASTRTC = True
+except Exception:
+    HAS_FASTRTC = False
+
+def build_fastrtc():
+    if not HAS_FASTRTC:
+        gr.Markdown("Install `fastrtc` to enable this section.")
+        return
+    with gr.Group():
+        gr.Markdown("Camera loopback using FastRTC WebRTC. Extend with streaming handlers later.")
+        rtc = WebRTC(mode="send-receive", modality="video")
+        rtc.stream(ReplyOnPause(_echo_video), inputs=[rtc], outputs=[rtc], time_limit=60)
+
+# ---------------------------
+# Inserts registry
+# ---------------------------
+
 INSERTS = {
     "TERMINAL": build_terminal,
     "ATTN_VIS": build_attn_vis,
     "ALLOC_PLOT": build_alloc_plot,
 }
+
+# ---------------------------
+# Layout / CSS / App
+# ---------------------------
+
+CSS = """
+:root { --toc-w: 280px; }
+#layout { display: grid; grid-template-columns: var(--toc-w) 1fr; gap: 1.25rem; }
+#toc { position: sticky; top: 0.75rem; height: calc(100vh - 1.5rem); overflow: auto; padding-right: .5rem; }
+#toc a { text-decoration: none; display: block; padding: .25rem 0; }
+.section { scroll-margin-top: 72px; }
+.gradio-container { max-width: 1200px !important; margin: 0 auto; }
+hr { border: none; border-top: 1px solid var(--neutral-300); margin: 1.25rem 0; }
+"""
+
+with gr.Blocks(css=CSS, fill_height=True, title="Interactive Blog — Transformers Feature Showcase") as demo:
+    gr.HTML("<h1>Transformers Feature Showcase</h1><p>Interactive, scrollable demo.</p>")
+    with gr.Row(elem_id="layout"):
+        with gr.Column(scale=0):
+            gr.HTML(
+                """
+                <nav id="toc">
+                  <h3>Sections</h3>
+                  <a href="#article">Article</a>
+                  <a href="#rtc">FastRTC (preview)</a>
+                </nav>
+                """
+            )
+        with gr.Column():
+            gr.HTML('<h2 id="article" class="section">Article</h2>')
+            # Author in Obsidian. Put {{ALLOC_PLOT}}, {{ATTN_VIS}}, {{TERMINAL}} where you want widgets.
+            render_article("content/article.md", INSERTS)
+            gr.HTML("<hr/>")
+
+            gr.HTML('<h2 id="rtc" class="section">FastRTC (preview)</h2>')
+            build_fastrtc()
+
+if __name__ == "__main__":
+    demo.launch()