Update app.py

app.py

@@ -1,4 +1,6 @@
-import io, traceback, numpy as np, gradio as gr
+import io, traceback, numpy as np, gradio as gr, matplotlib
+matplotlib.use("Agg")  # headless backend
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from wfgy_sdk import get_engine
 from wfgy_sdk.evaluator import compare_logits, plot_histogram
@@ -8,30 +10,27 @@ tok = AutoTokenizer.from_pretrained(MODEL)
 mdl = AutoModelForCausalLM.from_pretrained(MODEL)
 eng = get_engine()
 
-
 def run(prompt: str):
     prompt = prompt.strip()
     if not prompt:
         return "", "", "no prompt – nothing to show", None
     try:
         ids = tok(prompt, return_tensors="pt").input_ids
-        rawL = mdl(ids).logits[0, -1].detach().cpu().numpy()
+        raw = mdl(ids).logits[0, -1].detach().cpu().numpy()
         G = np.random.randn(256).astype(np.float32)
         I = G + np.random.normal(scale=0.05, size=256).astype(np.float32)
-        modL = eng.run(I, G, rawL)
-        m = compare_logits(rawL, modL)
-        headline = f"▼ var {m['var_drop']*100:4.1f} % | KL {m['kl']:.3f}"
-        fig = plot_histogram(rawL, modL)
+        mod = eng.run(I, G, raw)
+        m = compare_logits(raw, mod)
+        headline = f"▼ var {m['var_drop']*100:4.1f}% | KL {m['kl']:.3f}"
+        fig = plot_histogram(raw, mod)
         buf = io.BytesIO(); fig.savefig(buf, format="png"); buf.seek(0)
-        raw_txt = prompt + tok.decode(int(rawL.argmax()))
-        mod_txt = prompt + tok.decode(int(modL.argmax()))
+        raw_txt = prompt + tok.decode(int(raw.argmax()))
+        mod_txt = prompt + tok.decode(int(mod.argmax()))
         return raw_txt, mod_txt, headline, buf
-    except Exception:
-        # fallback: show traceback as text, return blanks for image
+    except Exception as e:
         tb = traceback.format_exc()
         return "runtime error", tb, "runtime error", None
 
-
 with gr.Blocks(title="WFGY variance gate") as demo:
     gr.Markdown("# 🧠 WFGY simulation demo")
     prompt = gr.Textbox(label="Prompt", value="Explain Schrödinger's cat")
@@ -40,16 +39,12 @@ with gr.Blocks(title="WFGY variance gate") as demo:
     with gr.Row():
         raw_box = gr.Textbox(label="Raw GPT-2")
         mod_box = gr.Textbox(label="After WFGY")
-
     headline = gr.Markdown()
     img = gr.Image(label="Logit histogram")
 
     btn.click(run, prompt, [raw_box, mod_box, headline, img])
 
-    gr.Markdown(
-        "---\n"
-        "### ⭐ Help unlock **WFGY 2.0** — 10 000 stars by **2025-08-01**"
-    )
+    gr.Markdown("---\n### ⭐ 10 000 stars → unlock **WFGY 2.0** by 2025-08-01")
 
 if __name__ == "__main__":
     demo.queue().launch()