Update index.js

index.js

@@ -1,81 +1,51 @@
-import { env, AutoTokenizer, RawImage, Tensor } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers';
-import { getModelJSON, getModelFile } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.0.2/src/utils/hub.js";
-import * as ort from "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.0/dist/ort.webgpu.mjs";
+import { env, AutoTokenizer, RawImage, Tensor } from '@huggingface/transformers';
+import { getModelJSON, getModelFile } from "@huggingface/transformers/utils/hub.js";
+import * as ort from "onnxruntime-web/webgpu";
 
-// Since we will download the model from the Hugging Face Hub, we can skip the local model check
-env.allowLocalModels = false;
-
-// Reference the elements that we will need
-const status = document.getElementById('status');
-const fileUpload = document.getElementById('upload');
-const imageContainer = document.getElementById('container');
-const example = document.getElementById('example');
-
-const EXAMPLE_URL = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg";
 const INPUT_IMAGE_SIZE = [960, 960];
 const HEIGHT_FACTOR = 10;
 const WIDTH_FACTOR = 10;
 const IMAGE_EMBED_SIZE = WIDTH_FACTOR * HEIGHT_FACTOR;
 const MAX_SEQ_LENGTH = 1024;
-const ONNX_URL = "http://localhost:3004/onnx";
 const BASE_MODEL = "Qwen/Qwen2-VL-2B-Instruct";
 const ONNX_MODEL = "pdufour/Qwen2-VL-2B-Instruct-ONNX-Q4-F16";
 const QUANT = "q4f16";
 const MAX_SINGLE_CHAT_LENGTH = 10;
 
-status.textContent = 'Loading model...';
-status.textContent = 'Ready';
-
-example.addEventListener('click', (e) => {
-    e.preventDefault();
-    parse(EXAMPLE_URL, 'Describe this image.');
-});
-
-fileUpload.addEventListener('change', function (e) {
-    const file = e.target.files[0];
-    if (!file) {
-        return;
-    }
-
-    const reader = new FileReader();
+let ortSessionA, ortSessionB, ortSessionC;
 
-    // Set up a callback when the file is loaded
-    reader.onload = e2 => parse(e2.target.result, '');
-
-    reader.readAsDataURL(file);
-});
+async function initializeSessions() {
+  ortSessionA = await ort.InferenceSession.create(
+    await getModelFile(ONNX_MODEL, `onnx/QwenVL_A_${QUANT}.onnx`),
+    { executionProviders: ["webgpu"] }
+  );
 
-async function parse(img, txt) {
-    imageContainer.innerHTML = '';
-    imageContainer.style.backgroundImage = `url(${img})`;
+  ortSessionB = await ort.InferenceSession.create(
+    await getModelFile(ONNX_MODEL, `onnx/QwenVL_B_${QUANT}.onnx`),
+    { executionProviders: ["webgpu"] }
+  );
 
-    status.textContent = 'Analysing...';
-    const output = await imageTextToText(img, txt);
-    status.textContent = '';
-    output.forEach(renderBox);
+  ortSessionC = await ort.InferenceSession.create(
+    await getModelFile(ONNX_MODEL, `onnx/QwenVL_C_${QUANT}.onnx`),
+    { executionProviders: ["webgpu"] }
+  );
 }
 
-
-export async function imageTextToText(
-  imagePath,
-  query,
-  vision = true
-) {
-  const config = (await getModelJSON(BASE_MODEL, "config.json"))
-
+export async function imageTextToText(imagePath, query, vision = true) {
+  const config = await getModelJSON(BASE_MODEL, "config.json");
+  
   const prompt_head_len = new Tensor("int64", new BigInt64Array([5n]), [1]);
-
-  let position_ids;
-  let num_decode = 0;
   let history_len = new Tensor("int64", new BigInt64Array([0n]), [1]);
-
+  let pos_factor = new Tensor("float16", new Uint16Array([0]), [1]);
+  let attention_mask = new ort.Tensor("float16", new Uint16Array([0xfbff]), [1]);
+  
   let past_key_states = new ort.Tensor(
     "float16",
     new Uint16Array(
       config.num_hidden_layers *
-        config.num_key_value_heads *
-        MAX_SEQ_LENGTH *
-        (config.hidden_size / config.num_attention_heads)
+      config.num_key_value_heads *
+      MAX_SEQ_LENGTH *
+      (config.hidden_size / config.num_attention_heads)
     ).fill(0),
     [
       config.num_hidden_layers,
@@ -84,17 +54,8 @@ export async function imageTextToText(
       config.hidden_size / config.num_attention_heads,
     ]
   );
-
   let past_value_states = past_key_states;
 
-  let attention_mask = new ort.Tensor(
-    "float16",
-    new Uint16Array([0xfbff]), // -65504.0 in float16
-    [1]
-  );
-
-  let pos_factor = new Tensor("float16", new Uint16Array([0]), [1]);
-
   const tokenizer = await AutoTokenizer.from_pretrained(BASE_MODEL);
   const prompt = `\n<|im_start|>user\n<|vision_start|><|vision_end|>${query}<|im_end|>\n<|im_start|>assistant\n`;
   const token = await tokenizer(prompt, {
@@ -103,18 +64,17 @@ export async function imageTextToText(
     tokenize: true,
   }).input_ids;
 
-  const seq_length = token.dims[1];
-  let ids_len = new Tensor("int64", new BigInt64Array([BigInt(seq_length)]), [
-    1,
-  ]);
-
+  let ids_len = new Tensor("int64", new BigInt64Array([BigInt(token.dims[1])]), [1]);
   let input_ids = new ort.Tensor(
     "int32",
     new Int32Array(MAX_SEQ_LENGTH).fill(0),
     [MAX_SEQ_LENGTH]
   );
+  input_ids.data.set(Array.from(token.data.slice(0, token.dims[1]), Number));
 
-  input_ids.data.set(Array.from(token.data.slice(0, seq_length), Number));
+  // Get position IDs from Session C
+  const dummy = new ort.Tensor("int32", new Int32Array([0]), []);
+  let { position_ids } = await ortSessionC.run({ dummy });
 
   if (vision) {
     let image = await RawImage.fromURL(imagePath);
@@ -122,51 +82,40 @@ export async function imageTextToText(
     image = image.rgb().toTensor("CHW").to("float32").div_(255.0);
     const pixel_values = image.unsqueeze(0);
 
-    const ortSessionA = await ort.InferenceSession.create(
-      await getModelFile(ONNX_MODEL, `onnx/QwenVL_A_${QUANT}.onnx`),
-      { executionProviders: ["webgpu"] }
-    );
-
     const { image_embed } = await ortSessionA.run({ pixel_values });
-
     ids_len = ids_len.add(BigInt(IMAGE_EMBED_SIZE));
 
     const ortSessionD = await ort.InferenceSession.create(
-      `${BASE_URL}/QwenVL_D${suffix}.onnx`,
+      await getModelFile(ONNX_MODEL, `onnx/QwenVL_D_${QUANT}.onnx`),
       { executionProviders: ["webgpu"] }
     );
 
-    ({ hidden_states: past_key_states, position_ids } =
-      await ortSessionD.run({
-        "hidden_states.1": past_key_states,
-        image_embed,
-        ids_len,
-        "ids_len_minus": new Tensor(
-          "int32",
-          new Int32Array([Number(ids_len.item()) - Number(prompt_head_len.item())]),
-          [1]
-        ),
-        "split_factor": new Tensor(
-          "int32",
-          new Int32Array([
-            MAX_SEQ_LENGTH - Number(ids_len.item()) - IMAGE_EMBED_SIZE,
-          ]),
-          [1]
-        ),
-      }));
-  }
+    const result = await ortSessionD.run({
+      "hidden_states.1": past_key_states,
+      image_embed,
+      ids_len,
+      "ids_len_minus": new Tensor(
+        "int32",
+        new Int32Array([Number(ids_len.item()) - Number(prompt_head_len.item())]),
+        [1]
+      ),
+      "split_factor": new Tensor(
+        "int32",
+        new Int32Array([MAX_SEQ_LENGTH - Number(ids_len.item()) - IMAGE_EMBED_SIZE]),
+        [1]
+      ),
+    });
 
-  const ortSessionB = await ort.InferenceSession.create(
-    `${BASE_URL}/QwenVL_B${suffix}.onnx`,
-    { executionProviders: ["webgpu"] }
-  );
+    past_key_states = result.hidden_states;
+    position_ids = result.position_ids;
+  }
 
-  while (
-    num_decode < MAX_SINGLE_CHAT_LENGTH &&
-    Number(history_len.data[0]) < MAX_SEQ_LENGTH
-  ) {
+  let num_decode = 0;
+  let output = '';
+  
+  while (num_decode < MAX_SINGLE_CHAT_LENGTH && Number(history_len.data[0]) < MAX_SEQ_LENGTH) {
     const ortSessionE = await ort.InferenceSession.create(
-      `${BASE_URL}/QwenVL_E_q4f16.onnx`,
+      await getModelFile(ONNX_MODEL, `onnx/QwenVL_E_${QUANT}.onnx`),
       { executionProviders: ["wasm"] }
     );
 
@@ -184,8 +133,9 @@ export async function imageTextToText(
     const token_id = result.max_logit_ids;
     if (token_id === 151643 || token_id === 151645) break;
 
+    output += tokenizer.decode([...token_id.data]);
+    
     num_decode++;
-
     history_len = history_len.add(BigInt(1));
     pos_factor = new Tensor(
       "float16",
@@ -204,6 +154,8 @@ export async function imageTextToText(
 
     past_key_states = hidden_states;
   }
-}
 
+  return output;
+}
 
+await initializeSessions();