diffusionpi 0.2.0

package/app/extensions/smooth-scroll.ts

@@ -0,0 +1,159 @@
+// diffusionpi hack: smooth viewport scroll.
+//
+// Pi's TUI renders the whole chat in one frame and its differential renderer
+// appends every new line at once, so a diffusion commit (~a whole canvas,
+// often 15+ lines) shoves the viewport up in a single jump. There is no
+// public extension API over the chat container or the viewport, so this
+// extension deliberately reaches into the live TUI instance (grabbed through
+// a throwaway widget factory) and wraps the chat container's render() to
+// reveal appended lines at a bounded rate instead of all at once. Real
+// content, paced only at the render layer.
+//
+// This is version-sensitive by design: it checks that the component tree
+// looks like Pi 0.8x (nine TUI children, chat at index 2) and no-ops
+// otherwise, so a Pi upgrade degrades to the stock jumpy behavior rather
+// than breaking.
+//
+// Tuning:
+//   DIFFUSIONPI_SMOOTH_SCROLL=0        disable the hack
+//   DIFFUSIONPI_SCROLL_SPEED=40        reveal rate in lines per second
+//   DIFFUSIONPI_SCROLL_DEBUG=<path>    append per-frame pacing decisions to a file
+import { appendFileSync } from "node:fs";
+
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+
+type RenderComponent = {
+  render(width: number): string[];
+  invalidate?(): void;
+};
+
+type TuiLike = {
+  children: RenderComponent[];
+  requestRender(force?: boolean): void;
+};
+
+type WidgetUi = {
+  setWidget(
+    key: string,
+    content: ((tui: TuiLike, theme: unknown) => RenderComponent) | undefined
+  ): void;
+};
+
+const enabled = process.env["DIFFUSIONPI_SMOOTH_SCROLL"] !== "0";
+const linesPerSecond = positiveNumber(process.env["DIFFUSIONPI_SCROLL_SPEED"], 40);
+const debugPath = process.env["DIFFUSIONPI_SCROLL_DEBUG"];
+
+function debugLog(message: string): void {
+  if (debugPath === undefined || debugPath === "") {
+    return;
+  }
+  try {
+    appendFileSync(debugPath, `${String(Date.now())} ${message}\n`);
+  } catch {
+    // Debug logging must never break rendering.
+  }
+}
+// A growth this large is a session restore or theme redraw, not streaming:
+// show it instantly instead of animating through it.
+const snapThresholdLines = 300;
+// Upper bound on the time credited between renders, so the first paced frame
+// after an idle period does not instantly spend seconds of accumulated budget.
+const maxTickSeconds = 0.1;
+const catchUpDelayMs = 33;
+const expectedTuiChildren = 9;
+const chatChildIndex = 2;
+
+function positiveNumber(raw: string | undefined, fallback: number): number {
+  const value = Number(raw);
+  return Number.isFinite(value) && value > 0 ? value : fallback;
+}
+
+export default function diffusionpiSmoothScroll(pi: ExtensionAPI): void {
+  if (!enabled) {
+    return;
+  }
+  let patched = false;
+  pi.on("session_start", (_event, ctx) => {
+    if (ctx.mode !== "tui" || patched) {
+      return;
+    }
+    patched = tryPatch(ctx);
+  });
+}
+
+// The widget factory is invoked synchronously with the live TUI instance;
+// registering and immediately removing a probe widget is the only way an
+// extension can get its hands on the TUI outside a dialog.
+function tryPatch(ctx: ExtensionContext): boolean {
+  let tui: TuiLike | undefined;
+  const ui = ctx.ui as unknown as WidgetUi;
+  if (typeof ui.setWidget !== "function") {
+    return false;
+  }
+  ui.setWidget("diffusionpi-smooth-scroll-probe", (liveTui) => {
+    tui = liveTui;
+    return { render: () => [], invalidate: () => undefined };
+  });
+  ui.setWidget("diffusionpi-smooth-scroll-probe", undefined);
+  if (tui === undefined || !Array.isArray(tui.children)) {
+    return false;
+  }
+  if (tui.children.length !== expectedTuiChildren) {
+    return false;
+  }
+  const chat = tui.children[chatChildIndex];
+  if (chat === undefined || typeof chat.render !== "function") {
+    return false;
+  }
+  installPacer(tui, chat);
+  debugLog("pacer installed");
+  return true;
+}
+
+function installPacer(tui: TuiLike, chat: RenderComponent): void {
+  const original = chat.render.bind(chat);
+  let revealed = Number.POSITIVE_INFINITY; // first render snaps to full
+  let lastWidth = -1;
+  let lastTickAt = 0;
+  let carry = 0;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+
+  const scheduleCatchUp = (): void => {
+    if (timer !== undefined) {
+      return;
+    }
+    timer = setTimeout(() => {
+      timer = undefined;
+      tui.requestRender();
+    }, catchUpDelayMs);
+  };
+
+  chat.render = (width: number): string[] => {
+    const lines = original(width);
+    const now = Date.now();
+    if (width !== lastWidth) {
+      // Rewrap after a resize changes every line count; animating the reflow
+      // would look like scrolling through the whole session.
+      lastWidth = width;
+      revealed = lines.length;
+      return lines;
+    }
+    if (lines.length <= revealed || lines.length - revealed > snapThresholdLines) {
+      revealed = lines.length;
+      lastTickAt = now;
+      carry = 0;
+      return lines;
+    }
+    const dt = Math.min(maxTickSeconds, Math.max(0, (now - lastTickAt) / 1000));
+    lastTickAt = now;
+    carry += linesPerSecond * dt;
+    const step = Math.floor(carry);
+    carry -= step;
+    revealed = Math.min(lines.length, revealed + step);
+    debugLog(`pacing revealed=${String(revealed)} total=${String(lines.length)} step=${String(step)}`);
+    if (revealed < lines.length) {
+      scheduleCatchUp();
+    }
+    return lines.slice(0, revealed);
+  };
+}

package/app/pi-factory.toml

@@ -0,0 +1,31 @@
+id = "diffusionpi"
+name = "DiffusionPi"
+version = "0.1.0"
+schema_version = 1
+description = "Pi + DiffusionGemma with a live diffusion canvas"
+state_dir = "~/.local/state/diffusionpi"
+# Stock Pi via pi-factory's default pi_command (npx @earendil-works/pi-coding-agent).
+thinking = "medium"
+
+[provider]
+id = "vllm"
+base_url = "http://127.0.0.1:8000/v1"
+api = "openai-completions"
+
+[model]
+id = "nvidia/diffusiongemma-26B-A4B-it-NVFP4"
+name = "DiffusionGemma 26B-A4B NVFP4"
+context_window = 32768
+max_tokens = 8192
+reasoning = true
+
+[[extensions]]
+path = "extensions/diffusion-canvas.ts"
+
+# Shared self-driving demo extension; inert unless PI_DEMO_MODE=1 (the demo
+# manifest sets it), so it is safe to keep loaded in interactive sessions.
+[[extensions]]
+path = "../node_modules/pi-demo-mode/extensions/demo-mode.ts"
+
+[[extensions]]
+path = "extensions/smooth-scroll.ts"

package/app/prompts/demo-followup.md

@@ -0,0 +1 @@
+Continue. Try to write as long as possible.

package/app/prompts/demo-initial.md

@@ -0,0 +1 @@
+You are narrating a never-ending sci-fi adventure. Continue in short paragraphs. Whenever the user sends a message, treat it as a live director note and incorporate it immediately. Never end the story.

package/bin/diffusionpi

@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+# Launches the diffusionpi Pi app bundle via pi-factory (stock Pi, no fork).
+# Subcommands: run (default), demo, plan, validate.
+set -euo pipefail
+
+root="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")/.." && pwd)"
+app_dir="$root/app"
+
+pi_factory() {
+  if command -v pi-factory >/dev/null 2>&1; then
+    pi-factory "$@"
+  else
+    npx -y @dutifuldev/pi-factory@latest "$@"
+  fi
+}
+
+# Shared Pi extension packages (pi-demo-mode) are npm dependencies of the
+# bundle; fetch them on first run so the TOML extension paths resolve.
+ensure_deps() {
+  if [ ! -d "$root/node_modules/pi-demo-mode" ]; then
+    npm install --prefix "$root" --no-fund --no-audit >&2
+  fi
+}
+
+ensure_deps
+
+case "${1:-run}" in
+  run) pi_factory run --app-dir "$app_dir" ;;
+  demo)
+    PI_DEMO_INITIAL_PROMPT_FILE="$app_dir/prompts/demo-initial.md" \
+      PI_DEMO_FOLLOWUP_PROMPT_FILE="$app_dir/prompts/demo-followup.md" \
+      pi_factory run --app-file "$app_dir/demo.pi-factory.toml"
+    ;;
+  plan) pi_factory plan --app-dir "$app_dir" ;;
+  validate) pi_factory validate "$app_dir" ;;
+  *)
+    echo "usage: diffusionpi [run|demo|plan|validate]" >&2
+    exit 2
+    ;;
+esac

package/docs/diffusion-canvas-repro.md

@@ -0,0 +1,104 @@
+# Reproducing the live diffusion canvas
+
+This documents every piece needed to reproduce the truthful diffusion canvas
+visualization: DiffusionGemma served by a patched vLLM that streams its
+intermediate denoising states, rendered live in the Pi TUI.
+
+## The pieces
+
+| Piece                 | Where it lives                                                                                                       | What it does                                                                                                                                                                                                                                                                                                                                                             |
+| --------------------- | -------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| vLLM canvas streaming | [`osolmaz/vllm`](https://github.com/osolmaz/vllm), branch `diffusion-canvas-events`, release tag `canvas-v0.23.1rc4` | Adds the opt-in `--diffusion-stream-canvas` server flag and the `GET /v1/diffusion/events` SSE side channel that emits the detokenized canvas on every denoising step. See the fork's [`DIFFUSION_CANVAS.md`](https://github.com/osolmaz/vllm/blob/diffusion-canvas-events/DIFFUSION_CANVAS.md). Tracked as [PR #1 on the fork](https://github.com/osolmaz/vllm/pull/1). |
+| Pi widget             | [`app/extensions/diffusion-canvas.ts`](../app/extensions/diffusion-canvas.ts) in this repo                           | Self-contained Pi extension that renders the canvas above the editor. Not published separately; this repo is its home.                                                                                                                                                                                                                                                   |
+| diffusionpi app       | [`app/`](../app/) bundle + [`bin/diffusionpi`](../bin/diffusionpi) launcher                                          | pi-factory bundle that launches stock Pi with the widget against the local vLLM server.                                                                                                                                                                                                                                                                                  |
+
+The vLLM changes are deliberately not upstreamed; they are maintained as
+tagged fork releases (`canvas-v<upstream version>rc<N>`, PEP 440-parseable so
+vLLM's version detection accepts them).
+
+## 1. Install the patched vLLM
+
+All fork changes are pure Python, so the install overlays them on the
+official precompiled kernels — no compilation:
+
+```bash
+VLLM_USE_PRECOMPILED=1 \
+VLLM_PRECOMPILED_WHEEL_COMMIT=4e5ca89cfe98121642d76b40e32a006f4d0fbf3b \
+pip install git+https://github.com/osolmaz/vllm@canvas-v0.23.1rc4
+```
+
+`VLLM_PRECOMPILED_WHEEL_COMMIT` pins the official wheel that provides the
+binary artifacts. It must match the fork release's upstream base commit,
+which is recorded in the fork's `DIFFUSION_CANVAS.md`.
+
+## 2. Serve DiffusionGemma with canvas streaming
+
+```bash
+vllm serve nvidia/diffusiongemma-26B-A4B-it-NVFP4 \
+  --host 127.0.0.1 --port 8000 \
+  --max-model-len 32768 --max-num-seqs 16 --max-num-batched-tokens 8192 \
+  --kv-cache-dtype fp8 \
+  --enable-auto-tool-choice --tool-call-parser gemma4 \
+  --diffusion-stream-canvas
+```
+
+`--diffusion-stream-canvas` requires the single-process FastAPI frontend (no
+`--grpc`, headless mode, or external/hybrid/multi-port data-parallel load
+balancing); the server rejects incompatible combinations at startup.
+
+To sanity-check the side channel independently of any client:
+
+```bash
+curl -N http://127.0.0.1:8000/v1/diffusion/events
+```
+
+then run any chat completion against the server; per-step JSON events
+(`request_id`, `step`, `block`, `text`) should stream out.
+
+## 3. Run the visualizer
+
+Through the diffusionpi launcher (from the repo root):
+
+```bash
+bin/diffusionpi         # interactive session
+bin/diffusionpi demo    # self-driving story demo
+```
+
+Or in plain Pi, load the self-contained extension file directly (copy it into
+your Pi extensions directory or pass it with `--extension`):
+
+```bash
+pi --extension ./diffusionpi/app/extensions/diffusion-canvas.ts
+```
+
+The widget needs no configuration: it derives the events and metrics URLs
+from the active model's `baseUrl`, with `PI_DIFFUSION_CANVAS_EVENTS_URL` and
+`PI_DIFFUSION_CANVAS_METRICS_URL` as overrides. Against a server without the
+side channel it falls back to a clearly labeled simulation paced by the real
+commit bursts.
+
+## How the truthful path works
+
+- Each denoising step, vLLM's diffusion sampler produces an `argmax_canvas`
+  (its current best guess for every canvas position). The fork carries these
+  through the draft-token path into `EngineCoreOutput`, detokenizes them in
+  the API server, and broadcasts them on `/v1/diffusion/events`.
+- The Pi extension assigns each completion a request id up front via the
+  `X-Request-Id` header (vLLM derives the `chatcmpl-*` id from it), then
+  subscribes to the side channel scoped to that id
+  (`?request_id=...`). Uncorrelated canvases are never rendered, so a shared
+  server cannot leak another client's text into the TUI.
+- The widget renders committed text and the resolving canvas as one
+  continuous document windowed at the settled/canvas seam, keeps a stable
+  height, discards canvas snapshots of already-committed blocks (the `block`
+  ordinal), and substitutes glyphs of ambiguous terminal width, so the TUI's
+  differential renderer never leaves stale frames in scrollback.
+
+## Upgrading the fork to a newer vLLM
+
+1. Rebase `diffusion-canvas-events` onto the new upstream base commit.
+2. Run the fork's test suite: `pytest tests/v1/engine/test_diffusion_events.py`.
+3. Update the base commit in the fork's `DIFFUSION_CANVAS.md` and in the
+   install one-liners here and in the package README.
+4. Tag the next `canvas-v<upstream version>rc<N>` and push the tag.
+5. Re-verify the install one-liner in a scratch venv and the live TUI flow.

package/package.json

@@ -0,0 +1,33 @@
+{
+  "name": "diffusionpi",
+  "version": "0.2.0",
+  "description": "Watch a diffusion LLM denoise its answer live in the Pi TUI",
+  "keywords": [
+    "pi",
+    "diffusion",
+    "vllm",
+    "tui"
+  ],
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/osolmaz/diffusionpi.git"
+  },
+  "bugs": {
+    "url": "https://github.com/osolmaz/diffusionpi/issues"
+  },
+  "homepage": "https://github.com/osolmaz/diffusionpi#readme",
+  "bin": {
+    "diffusionpi": "bin/diffusionpi"
+  },
+  "files": [
+    "bin/**",
+    "app/**",
+    "docs/**",
+    "README.md",
+    "LICENSE"
+  ],
+  "dependencies": {
+    "pi-demo-mode": "github:osolmaz/pi-demo-mode#v0.1.0"
+  }
+}