Upload example_web.js with huggingface_hub

example_web.js

@@ -0,0 +1,170 @@
+// Minimal browser inference for Bartholomheow/Supra2-IMG-ONNX — no build step.
+// Serve this folder over http(s) and open index.html (modules + WebGPU need secure context).
+//
+// npm i onnxruntime-web @huggingface/transformers   (or use the import map below in index.html)
+/*
+<script type="importmap">
+{"imports":{
+  "onnxruntime-web": "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/ort.web.min.mjs",
+  "@huggingface/transformers": "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1"
+}}
+</script>
+*/
+
+const REPO = 'Bartholomheow/Supra2-IMG-ONNX';
+const fileUrl = (path) => `https://huggingface.co/${REPO}/resolve/main/${path}`;
+
+function mulberry32(seed) {
+  let a = seed >>> 0;
+  return () => {
+    a = (a + 0x6d2b79f5) | 0;
+    let t = Math.imul(a ^ (a >>> 15), 1 | a);
+    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+}
+
+function gaussian(seed, n) {
+  const rand = mulberry32(seed);
+  const out = new Float32Array(n);
+  for (let i = 0; i < n; i += 2) {
+    const r = Math.sqrt(-2 * Math.log(Math.max(rand(), 1e-12)));
+    const a = 2 * Math.PI * rand();
+    out[i] = r * Math.cos(a);
+    if (i + 1 < n) out[i + 1] = r * Math.sin(a);
+  }
+  return out;
+}
+
+async function download(url, file, onProgress) {
+  const cache = await caches.open('supra2-img-v1');
+  const hit = await cache.match(url);
+  if (hit) return hit.arrayBuffer();
+  const res = await fetch(url);
+  if (!res.ok) throw new Error(`download failed (${res.status}): ${file}`);
+  const total = Number(res.headers.get('content-length')) || null;
+  const reader = res.body.getReader();
+  const chunks = [];
+  let loaded = 0;
+  for (;;) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    chunks.push(value);
+    loaded += value.byteLength;
+    onProgress?.({ file, loaded, total });
+  }
+  const buf = new Uint8Array(loaded);
+  let off = 0;
+  for (const c of chunks) { buf.set(c, off); off += c.byteLength; }
+  await cache.put(url, new Response(buf.slice(0)));
+  return buf.buffer;
+}
+
+function halfToFloat(h) {
+  const s = (h & 0x8000) << 16;
+  const e = (h >> 10) & 0x1f;
+  const m = h & 0x3ff;
+  const f = new Float32Array(1);
+  const u = new Uint32Array(f.buffer);
+  if (e === 0) {
+    if (m === 0) { u[0] = s; return f[0]; }
+    let mm = m, ee = -14;
+    while ((mm & 0x400) === 0) { mm <<= 1; ee -= 1; }
+    u[0] = s | ((ee + 127) << 23) | ((mm & 0x3ff) << 13);
+    return f[0];
+  }
+  if (e === 31) { u[0] = s | 0x7f800000 | (m << 13); return f[0]; }
+  u[0] = s | ((e + 112) << 23) | (m << 13);
+  return f[0];
+}
+
+export async function loadSupra({ onProgress } = {}) {
+  const ort = await import('onnxruntime-web');
+  const { AutoTokenizer } = await import('@huggingface/transformers');
+  const cfg = await (await fetch(fileUrl('pipeline_config.json'))).json();
+  const [ditBuf, encBuf, vaeBuf] = await Promise.all([
+    download(fileUrl(cfg.dit), cfg.dit, onProgress),
+    download(fileUrl(cfg.text_encoder), cfg.text_encoder, onProgress),
+    download(fileUrl(cfg.vae_decoder), cfg.vae_decoder, onProgress),
+  ]);
+  const opts = { executionProviders: ['webgpu'] };
+  const [dit, enc, vae] = await Promise.all([
+    ort.InferenceSession.create(ditBuf, opts),
+    ort.InferenceSession.create(encBuf, opts),
+    ort.InferenceSession.create(vaeBuf, opts),
+  ]);
+  const tokenize = await AutoTokenizer.from_pretrained(REPO);
+  return { ort, cfg, dit, enc, vae, tokenize };
+}
+
+export async function generate(model, prompt, { seed = 1, steps, cfg: guide } = {}) {
+  const { ort, cfg, dit, enc, vae, tokenize } = model;
+  steps ??= cfg.default_steps;
+  guide ??= cfg.default_cfg;
+  const N = cfg.latent_ch * cfg.latent_size ** 2;
+
+  async function encode(text) {
+    const t = await tokenize([text], { padding: 'max_length', truncation: true, max_length: cfg.ctx_len, return_attention_mask: true });
+    const toBig = (a) => BigInt64Array.from(a.data, (v) => BigInt(v));
+    const ids = toBig(t.input_ids);
+    const am = toBig(t.attention_mask);
+    const out = await enc.run({
+      input_ids: new ort.Tensor('int64', ids, [1, cfg.ctx_len]),
+      attention_mask: new ort.Tensor('int64', am, [1, cfg.ctx_len]),
+    });
+    const hidden = Object.values(out)[0];
+    // Encoder output may be fp16 halves or fp32 — normalize to float32.
+    const raw = hidden.data;
+    const data = raw instanceof Uint16Array ? Float32Array.from(raw, halfToFloat) : Float32Array.from(raw);
+    const mask = new Float32Array(cfg.ctx_len);
+    for (let i = 0; i < cfg.ctx_len; i++) mask[i] = am[i] === 0n ? 0 : 1;
+    return { ctx: data, mask };
+  }
+
+  // NOTE: encoder output here is decoded as fp32 for clarity. The shipped
+  // encoder_model.onnx is fp16 — for exact halves use a proper f16->f32 cast.
+  const cond = await encode(prompt);
+  const uncond = await encode('');
+  const shape = [1, cfg.latent_ch, cfg.latent_size, cfg.latent_size];
+  const toCtx = (c) => new ort.Tensor('float32', c.ctx, [1, cfg.ctx_len, c.ctx.length / cfg.ctx_len]);
+  let z = gaussian([...prompt].reduce((a, c) => a + c.codePointAt(0), seed), N);
+  const dt = 1 / steps;
+  for (let i = 0; i < steps; i++) {
+    const t = new ort.Tensor('float32', new Float32Array([i * dt]), [1]);
+    const vs = [];
+    for (const c of [cond, uncond]) {
+      const out = await dit.run({
+        z: new ort.Tensor('float32', z, shape), t,
+        ctx: toCtx(c), ctx_mask: new ort.Tensor('float32', c.mask, [1, cfg.ctx_len]),
+      });
+      vs.push(Object.values(out)[0].data);
+    }
+    const [vc, vu] = vs;
+    const next = new Float32Array(N);
+    for (let j = 0; j < N; j++) next[j] = z[j] + dt * (vu[j] + guide * (vc[j] - vu[j]));
+    z = next;
+  }
+  const scaled = new Float32Array(N);
+  for (let i = 0; i < N; i++) scaled[i] = z[i] / cfg.vae_scale;
+  const img = await vae.run({ z: new ort.Tensor('float32', scaled, shape) });
+  return { pixels: Object.values(img)[0].data, size: cfg.image_size }; // float NCHW in [-1, 1]
+}
+
+export function paint(pixels, size, canvas) {
+  canvas.width = size;
+  canvas.height = size;
+  const ctx = canvas.getContext('2d');
+  const img = ctx.createImageData(size, size);
+  for (let i = 0; i < size * size; i++) {
+    img.data[i * 4] = Math.round(Math.min(1, Math.max(0, (pixels[i] + 1) / 2)) * 255);
+    img.data[i * 4 + 1] = Math.round(Math.min(1, Math.max(0, (pixels[size * size + i] + 1) / 2)) * 255);
+    img.data[i * 4 + 2] = Math.round(Math.min(1, Math.max(0, (pixels[2 * size * size + i] + 1) / 2)) * 255);
+    img.data[i * 4 + 3] = 255;
+  }
+  ctx.putImageData(img, 0, 0);
+}
+
+// Usage:
+// const model = await loadSupra({ onProgress: (p) => console.log(p.file, p.loaded) });
+// const { pixels, size } = await generate(model, 'a lighthouse above violet clouds at dusk');
+// paint(pixels, size, document.querySelector('canvas'));