Update js/agate.js

js/agate.js

@@ -24,6 +24,12 @@ export const VERSIONS = {
 const countUrl = (v) => `https://huggingface.co/Logolabs/agate-preview-${v}/resolve/main/config.json`;
 const FILES_T2 = ["tokenizer.json", "tokenizer_config.json", "text_encoder.onnx", "generator.onnx", "taesd_decoder.onnx"];
 
+// Image-to-image: the SD 1.5 VAE encoder lives in its own Cache Storage bucket. It must NOT share a model
+// cache: Agate.load() deletes every entry of a release's cache that is not one of that release's files.
+const IMG_ENCODER_URL = "./models/vae_encoder/model_sd15.onnx";
+const IMG_ENCODER_CACHE = "agate-img-encoder";
+const SD_LATENT_SCALE = 0.18215;
+
 export async function webgpuStatus() {
   if (!("gpu" in navigator)) return { ok: false, why: "This browser does not expose WebGPU (navigator.gpu is missing)." };
   try {
@@ -75,6 +81,7 @@ export async function hasCachedModel(version = "001") {
 export async function clearCache() {
   let ok = true;
   for (const v of Object.values(VERSIONS)) { try { ok = (await caches.delete(v.cache)) && ok; } catch { ok = false; } }
+  try { await caches.delete(IMG_ENCODER_CACHE); } catch { ok = false; }
   return ok;
 }
 
@@ -244,6 +251,80 @@ export class Agate {
     return out;
   }
 
+  // Loads the SD 1.5 VAE encoder on first use (own cache bucket, see IMG_ENCODER_CACHE).
+  async loadImageEncoder() {
+    if (this.sessions.imgEncoder) return this.sessions.imgEncoder;
+    const key = new Request(new URL(IMG_ENCODER_URL, location.href));
+    let buf;
+    try {
+      ({ buf } = await fetchCached(IMG_ENCODER_CACHE, IMG_ENCODER_URL, key, 0, () => {}));
+    } catch (e) {
+      throw new Error(`Could not load the image encoder (${IMG_ENCODER_URL}): ${e.message || e}`);
+    }
+    this.sessions.imgEncoder = await ort.InferenceSession.create(buf, {
+      executionProviders: this.ep === "webgpu" ? ["webgpu", "wasm"] : ["wasm"],
+    });
+    return this.sessions.imgEncoder;
+  }
+
+  // File -> scaled SD-VAE latent (Float32Array, C * (res/8)^2) at the CURRENT resolution. Re-run it after a
+  // resolution switch: a 512 px latent is 64x64, a 256 px one is 32x32.
+  async encodeReferenceImage(file) {
+    if (!file) throw new Error("Please select a reference image.");
+    const encoder = await this.loadImageEncoder();
+
+    // Read the image and center-crop it to a square.
+    const bitmap = await createImageBitmap(file);
+    const size = this.resolution;
+    const canvas = document.createElement("canvas");
+    canvas.width = size;
+    canvas.height = size;
+
+    const ctx = canvas.getContext("2d", { willReadFrequently: true });
+    const scale = Math.max(size / bitmap.width, size / bitmap.height);
+    const sw = size / scale;
+    const sh = size / scale;
+    const sx = (bitmap.width - sw) / 2;
+    const sy = (bitmap.height - sh) / 2;
+
+    ctx.drawImage(bitmap, sx, sy, sw, sh, 0, 0, size, size);
+    bitmap.close();
+
+    const pixels = ctx.getImageData(0, 0, size, size).data;
+    const plane = size * size;
+    const rgb = new Float32Array(3 * plane);
+
+    // RGB, channel-first, normalized to [-1, 1].
+    for (let i = 0; i < plane; i++) {
+      rgb[i] = pixels[i * 4] / 127.5 - 1;
+      rgb[plane + i] = pixels[i * 4 + 1] / 127.5 - 1;
+      rgb[2 * plane + i] = pixels[i * 4 + 2] / 127.5 - 1;
+    }
+
+    const inputName = encoder.inputNames[0];
+    const outputName = encoder.outputNames[0];
+    const result = await encoder.run({
+      [inputName]: new ort.Tensor("float32", rgb, [1, 3, size, size]),
+    });
+
+    const latent = result[outputName];
+    const C = this.manifest.latent_ch;
+    const perLatent = C * (size / 8) * (size / 8);
+    let data = latent.data;
+
+    // Some SD VAE exports return "moments" (mean + logvar, 2*C channels) instead of a sampled latent.
+    // The mean is the first C channels, which are contiguous in NCHW layout.
+    if (data.length === 2 * perLatent) data = data.subarray(0, perLatent);
+    if (data.length !== perLatent) {
+      throw new Error(`Unexpected encoder output: ${latent.dims.join(" × ")} (expected ${perLatent} values)`);
+    }
+
+    // Agate samples in scaled SD-VAE latent space.
+    const out = Float32Array.from(data, (value) => value * SD_LATENT_SCALE);
+    latent.dispose?.();
+    return out;
+  }
+
   // SD3's resolution shift in Agate's convention (t = 0 noise), as agate/pipeline.py shift_t
   static shiftT(t, shift) {
     if (shift === 1) return t;
@@ -304,7 +385,16 @@ export class Agate {
 
   // onPreview({ step, steps, plan, x1, hw }) after every step, when given and the graph outputs the plan: plan =
   // the thinker output for the conditional branch (Float32Array 640 x 16 x 16), x1 = z + (1 - t) v (guided).
-  async generate({ prompt, negative = "", seed = 0, steps = 50, cfg = 3.0, noise = null, watermark = true, onStep = () => {}, onPreview = null, previewReady = () => true, shouldStop = () => false }) {
+  //
+  // Image-to-image: pass referenceLatent (from encodeReferenceImage) and denoise in [0, 1]. denoise = 1 is plain
+  // text-to-image noise; lower values start closer to the reference. The sampler then runs round(steps * denoise)
+  // steps (at least 1; none at all for denoise 0, which returns the reference latent decoded), so the step size and the cost per unit of change match a full text-to-image run.
+  async generate({
+    prompt, negative = "", seed = 0, steps = 50, cfg = 3.0,
+    noise = null, referenceLatent = null, denoise = 1.0,
+    watermark = true, onStep = () => {}, onPreview = null,
+    previewReady = () => true, shouldStop = () => false
+  }) {
     const C = this.manifest.latent_ch, D = this.manifest.ctx_dim;
     const rc = this.mr ? this.manifest.resolutions[String(this.resolution)] : { latent_hw: this.manifest.latent_hw, shift: 1 };
     const hw = rc.latent_hw, shift = Number(rc.shift || 1);
@@ -329,42 +419,82 @@ export class Agate {
       feeds.counts = new ort.Tensor("float32", cnt, [2, L]);
     }
 
-    let z = noise ? Float32Array.from(noise) : Agate.gaussianNoise(seed, per);
-    const grid = Array.from({ length: steps + 1 }, (_, i) => Agate.shiftT(i / steps, shift));
+    // ---- starting point: pure noise (text-to-image) or a noised reference latent (image-to-image) ----
+    const dn = Number(denoise);
+    denoise = Number.isFinite(dn) ? Math.max(0, Math.min(1, dn)) : 1;
+
+    if (referenceLatent && referenceLatent.length !== per) {
+      throw new Error(
+        `Reference latent has ${referenceLatent.length} values; expected ${per}. ` +
+        `Re-encode the reference image after changing resolution.`
+      );
+    }
+
+    const zNoise = noise
+      ? Float32Array.from(noise)
+      : Agate.gaussianNoise(seed, per);
+
+    let z = zNoise;
+    const startT = referenceLatent ? 1 - denoise : 0;
+
+    if (referenceLatent && denoise === 0) {
+      z = Float32Array.from(referenceLatent);
+    } else if (referenceLatent) {
+      const ts = Agate.shiftT(startT, shift);
+      z = new Float32Array(per);
+
+      for (let k = 0; k < per; k++) {
+        z[k] = (1 - ts) * zNoise[k] + ts * referenceLatent[k];
+      }
+    }
+
+    // denoise 0 means "return the reference as is": no sampling. Otherwise at least one step.
+    const nSteps = !referenceLatent ? steps : denoise === 0 ? 0 : Math.max(1, Math.round(steps * denoise));
+    const grid = nSteps === 0 ? [1] : Array.from(
+      { length: nSteps + 1 },
+      (_, i) => Agate.shiftT(
+        startT + (1 - startT) * i / nSteps,
+        shift
+      )
+    );
+
     t0 = performance.now();
     const stepMs = [];
-    if (this.fast) {
+    if (nSteps === 0) {
+      // z is already the reference latent: skip the sampler
+    } else if (this.fast) {
       z = await this.sampleFast({ z, grid, cfg, hw, C, L, D, ctx, mask, cnt: feeds.counts?.data, onStep, onPreview, previewReady, shouldStop, stepMs });
     } else {
-    const zz = new Float32Array(2 * per), tt = new Float32Array(2);
-    for (let i = 0; i < steps; i++) {
-      if (shouldStop()) throw new Error("stopped");
-      const ts = performance.now();
-      zz.set(z, 0); zz.set(z, per);
-      tt[0] = tt[1] = grid[i];
-      const want = onPreview && this.hasPlan ? ["v", "plan"] : ["v"];
-      const out = await this.sessions.gen.run({ z: new ort.Tensor("float32", zz, [2, C, hw, hw]), t: new ort.Tensor("float32", tt, [2]), ...feeds }, want);
-      const v = out.v.data;
-      const dt = grid[i + 1] - grid[i];
-      const zn = new Float32Array(per);
-      let x1 = null;
-      if (out.plan) x1 = new Float32Array(per);
-      for (let k = 0; k < per; k++) {
-        const vc = v[k], vu = v[per + k], g = vu + cfg * (vc - vu);
-        zn[k] = z[k] + dt * g;
-        if (x1) x1[k] = z[k] + (1 - grid[i]) * g;
+      const zz = new Float32Array(2 * per), tt = new Float32Array(2);
+      for (let i = 0; i < nSteps; i++) {
+        if (shouldStop()) throw new Error("stopped");
+        const ts = performance.now();
+        zz.set(z, 0); zz.set(z, per);
+        tt[0] = tt[1] = grid[i];
+        const want = onPreview && this.hasPlan ? ["v", "plan"] : ["v"];
+        const out = await this.sessions.gen.run({ z: new ort.Tensor("float32", zz, [2, C, hw, hw]), t: new ort.Tensor("float32", tt, [2]), ...feeds }, want);
+        const v = out.v.data;
+        const dt = grid[i + 1] - grid[i];
+        const zn = new Float32Array(per);
+        let x1 = null;
+        if (out.plan) x1 = new Float32Array(per);
+        for (let k = 0; k < per; k++) {
+          const vc = v[k], vu = v[per + k], g = vu + cfg * (vc - vu);
+          zn[k] = z[k] + dt * g;
+          if (x1) x1[k] = z[k] + (1 - grid[i]) * g;
+        }
+        out.v.dispose?.();
+        if (out.plan) { const plan = out.plan.data.slice(); out.plan.dispose?.(); onPreview({ step: i + 1, steps: nSteps, plan, x1, hw }); }
+        z = zn;
+        stepMs.push(performance.now() - ts);
+        onStep(i + 1, nSteps, z);
+        await new Promise((r) => setTimeout(r, 0));          // let the progress bar paint
       }
-      out.v.dispose?.();
-      if (out.plan) { const plan = out.plan.data.slice(); out.plan.dispose?.(); onPreview({ step: i + 1, steps, plan, x1, hw }); }
-      z = zn;
-      stepMs.push(performance.now() - ts);
-      onStep(i + 1, steps, z);
-      await new Promise((r) => setTimeout(r, 0));          // let the progress bar paint
-    }
     }
     T.samplerMs = performance.now() - t0;
-    T.firstStepMs = stepMs[0];
-    T.stepMs = stepMs.length > 1 ? (T.samplerMs - stepMs[0]) / (stepMs.length - 1) : stepMs[0];
+    T.steps = nSteps;
+    T.firstStepMs = stepMs[0] ?? 0;
+    T.stepMs = stepMs.length > 1 ? (T.samplerMs - stepMs[0]) / (stepMs.length - 1) : (stepMs[0] ?? 0);
 
     t0 = performance.now();
     const img = await this.sessions.vae.run({ latent: new ort.Tensor("float32", z, [1, C, hw, hw]) });
@@ -385,4 +515,4 @@ export class Agate {
     T.bucket = L;
     return { rgba, width: W, height: H, latent: z, timings: T, prepared: { ...prep, countsNonzero }, watermarked: watermark };
   }
-}
+}