Spaces:
Running
Running
File size: 3,407 Bytes
ae0648d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | /* PXG-Tiny playground worker — runs the real NumPy pipeline in-browser
via Pyodide. The model weights (~1.8 MB) load from this same site. */
let pyodideReady = null;
async function boot() {
importScripts("https://cdn.jsdelivr.net/pyodide/v0.26.2/full/pyodide.js");
const py = await loadPyodide();
self.postMessage({ type: "status", msg: "loading numpy…" });
await py.loadPackage("numpy");
self.postMessage({ type: "status", msg: "fetching model files…" });
// write runtime modules + weights into the Pyodide FS
const files = [
["pxg_tiny/__init__.py", "/lib/pxg_tiny/__init__.py"],
["pxg_tiny/config.py", "/lib/pxg_tiny/config.py"],
["pxg_tiny/tokenizer.py", "/lib/pxg_tiny/tokenizer.py"],
["pxg_tiny/quality.py", "/lib/pxg_tiny/quality.py"],
["pxg_tiny/runtime_pipeline.py", "/lib/pxg_tiny/runtime_pipeline.py"],
["webgen.py", "/lib/webgen.py"],
["weights/runtime.json", "/lib/bundle/runtime.json"],
["weights/tokenizer.json", "/lib/bundle/tokenizer.json"],
];
py.FS.mkdirTree("/lib/pxg_tiny");
py.FS.mkdirTree("/lib/bundle");
for (const [src, dst] of files) {
const buf = await (await fetch(src)).arrayBuffer();
py.FS.writeFile(dst, new Uint8Array(buf));
}
for (const npz of ["gen_int8.npz", "vq_int8.npz"]) {
const buf = await (await fetch("weights/" + npz)).arrayBuffer();
py.FS.writeFile("/lib/bundle/" + npz, new Uint8Array(buf));
self.postMessage({ type: "status", msg: "fetched " + npz });
}
self.postMessage({ type: "status", msg: "warming up the model…" });
await py.runPythonAsync(`
import sys
sys.path.insert(0, "/lib")
from webgen import WebPipeline
PIPE = WebPipeline("/lib/bundle")
_g, _m = PIPE.generate("a golden sword", seed=0)
"ready"
`);
self.postMessage({ type: "ready" });
return py;
}
const PYTHON_GEN = `
import json
import numpy as np
from webgen import WebPipeline
def _run(prompt, seed, verify):
grid, meta = PIPE.generate(prompt, seed=int(seed), enforce_quality=bool(verify))
if grid is None:
return json.dumps({"gate": meta["gate"], "message": meta.get("message",""), "rgba": None,
"seed": 0, "attempt": 0})
rgba = PIPE.grid_rgba(grid)
return json.dumps({"gate": meta.get("gate",""), "message": meta.get("message",""),
"rgba": rgba.reshape(-1).tolist(),
"seed": int(meta.get("seed", seed)),
"attempt": int(meta.get("attempt", 0))})
`;
onmessage = async (e) => {
const { type, prompt, seed, verify } = e.data;
if (type === "init") {
try {
if (!pyodideReady) pyodideReady = boot();
await pyodideReady;
const py = await pyodideReady;
if (!py.globals.has("_run")) py.runPython(PYTHON_GEN);
} catch (err) {
self.postMessage({ type: "error", msg: String(err) });
}
} else if (type === "generate") {
try {
if (!pyodideReady) pyodideReady = boot();
const py = await pyodideReady;
if (!py.globals.has("_run")) py.runPython(PYTHON_GEN);
py.globals.set("_prompt", prompt);
py.globals.set("_seed", seed);
py.globals.set("_verify", verify);
const out = py.runPython(
"_run(_prompt, _seed, _verify)");
self.postMessage({ type: "result", data: JSON.parse(out) });
} catch (err) {
self.postMessage({ type: "error", msg: String(err) });
}
}
};
|