Update app.py

app.py

@@ -1,6 +1,7 @@
 import io, traceback, numpy as np, gradio as gr, matplotlib
-matplotlib.use("Agg")  # headless backend
+matplotlib.use("Agg")
 
+from PIL import Image
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from wfgy_sdk import get_engine
 from wfgy_sdk.evaluator import compare_logits, plot_histogram
@@ -10,6 +11,7 @@ tok = AutoTokenizer.from_pretrained(MODEL)
 mdl = AutoModelForCausalLM.from_pretrained(MODEL)
 eng = get_engine()
 
+
 def run(prompt: str):
     prompt = prompt.strip()
     if not prompt:
@@ -22,15 +24,21 @@ def run(prompt: str):
         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)
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png")
+        buf.seek(0)
+        img = Image.open(buf)
+
         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 as e:
+        return raw_txt, mod_txt, headline, img
+    except Exception:
         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,7 +48,7 @@ with gr.Blocks(title="WFGY variance gate") as demo:
         raw_box = gr.Textbox(label="Raw GPT-2")
         mod_box = gr.Textbox(label="After WFGY")
     headline = gr.Markdown()
-    img = gr.Image(label="Logit histogram")
+    img = gr.Image(label="Logit histogram", type="pil")
 
     btn.click(run, prompt, [raw_box, mod_box, headline, img])