Merge pull request #4 from Jefreesujit/feat/update-models

Update Models

Author: Jefreesujit

src/app/audio-helpers.ts

@@ -1,4 +1,4 @@
-import type { AudioTranscriptionChunk } from "../types";
+import type { AudioTranscriptionChunk, ModelDescriptor } from "../types";
 
 export const EMPTY_TRANSCRIPTION_MESSAGE =
   "No speech was detected in that clip. Try recording closer to the microphone or upload a clearer file.";
@@ -21,3 +21,17 @@ export const normalizeTranscriptionText = (
 
   return chunkText;
 };
+
+export const resolveVoiceSpeakerEmbeddings = (
+  model: Pick<ModelDescriptor, "runtime">,
+  voiceId: string | undefined,
+) => {
+  if (voiceId) {
+    const matchedVoice = model.runtime.voices?.find((voice) => voice.id === voiceId);
+    if (matchedVoice?.speakerEmbeddingsUrl) {
+      return matchedVoice.speakerEmbeddingsUrl;
+    }
+  }
+
+  return model.runtime.speakerEmbeddingsUrl;
+};

src/models.ts

@@ -383,6 +383,68 @@ export const CURATED_MODELS: ModelDescriptor[] = [
       visionLoaderKind: "lfm2_5_vl",
     },
   }),
+  createVisionModel({
+    id: "onnx-community/gemma-4-E2B-it-qat-mobile-ONNX",
+    label: "Gemma 4 E2B Mobile",
+    summary:
+      "QAT mobile Gemma 4 export for experimental multimodal chat on stronger browsers.",
+    publisher: "onnx-community",
+    paramsLabel: "2B params",
+    parameterTier: "L",
+    estimatedDownloadLabel: "~2.4 GB multimodal download",
+    category: "desktop_experimental",
+    starter: false,
+    tested: false,
+    hf: {
+      modelId: "onnx-community/gemma-4-E2B-it-qat-mobile-ONNX",
+      pipelineTag: "any-to-any",
+      libraryName: "transformers.js",
+      tags: ["transformers.js", "any-to-any", "conversational", "gemma4"],
+      baseModel: "google/gemma-4-E2B-it-qat-mobile-transformers",
+      hasChatTemplate: true,
+    },
+    runtime: {
+      contextWindowTokens: 32768,
+      visionLoaderKind: "gemma_any_to_any",
+      visionDtypes: {
+        audio_encoder: "q2f16",
+        decoder_model_merged: "q2f16",
+        embed_tokens: "q2f16",
+        vision_encoder: "fp16",
+      },
+    },
+  }),
+  createVisionModel({
+    id: "onnx-community/gemma-3n-E2B-it-ONNX",
+    label: "Gemma 3n E2B",
+    summary:
+      "Gemma 3n's experimental any-to-any browser export with text, image, and audio towers.",
+    publisher: "onnx-community",
+    paramsLabel: "2B params",
+    parameterTier: "L",
+    estimatedDownloadLabel: "Large multimodal browser download",
+    category: "desktop_experimental",
+    starter: false,
+    tested: false,
+    hf: {
+      modelId: "onnx-community/gemma-3n-E2B-it-ONNX",
+      pipelineTag: "image-text-to-text",
+      libraryName: "transformers.js",
+      tags: ["transformers.js", "image-text-to-text", "conversational", "gemma3n"],
+      baseModel: "google/gemma-3n-E2B-it",
+      hasChatTemplate: true,
+    },
+    runtime: {
+      contextWindowTokens: 32768,
+      visionLoaderKind: "gemma_any_to_any",
+      visionDtypes: {
+        audio_encoder: "q4f16",
+        decoder_model_merged: "q4f16",
+        embed_tokens: "q4f16",
+        vision_encoder: "fp16",
+      },
+    },
+  }),
   createTextModel({
     id: "HuggingFaceTB/SmolLM3-3B-ONNX",
     label: "SmolLM3 3B",
@@ -432,29 +494,34 @@ export const CURATED_MODELS: ModelDescriptor[] = [
       contextWindowTokens: 32768,
     },
   }),
-  createTextModel({
-    id: "onnx-community/gemma-4-E2B-it-ONNX",
-    label: "Gemma 4 2B",
-    summary: "Experimental new Gemma model with strong reasoning capabilities.",
+  createVisionModel({
+    id: "onnx-community/gemma-4-E4B-it-qat-mobile-ONNX",
+    label: "Gemma 4 E4B Mobile",
+    summary:
+      "Larger QAT mobile Gemma 4 variant for desktop-class multimodal experiments in-browser.",
     publisher: "onnx-community",
-    paramsLabel: "2B params",
+    paramsLabel: "4B params",
     parameterTier: "L",
-    estimatedDownloadLabel: "Desktop experimental download",
+    estimatedDownloadLabel: "~3.0 GB multimodal download",
     category: "desktop_experimental",
     starter: false,
     tested: false,
     hf: {
-      modelId: "onnx-community/gemma-4-E2B-it-ONNX",
-      pipelineTag: "text-generation",
+      modelId: "onnx-community/gemma-4-E4B-it-qat-mobile-ONNX",
+      pipelineTag: "any-to-any",
       libraryName: "transformers.js",
-      tags: ["transformers.js", "text-generation", "conversational"],
-      baseModel: "google/gemma-4-E2B-it",
+      tags: ["transformers.js", "any-to-any", "conversational", "gemma4"],
+      baseModel: "google/gemma-4-E4B-it-qat-mobile-transformers",
       hasChatTemplate: true,
     },
     runtime: {
-      preferredDtype: "q4f16",
-      fallbackDtype: "q4",
       contextWindowTokens: 32768,
+      visionLoaderKind: "gemma_any_to_any",
+      visionDtypes: {
+        decoder_model_merged: "q2f16",
+        embed_tokens: "q2f16",
+        vision_encoder: "fp16",
+      },
     },
   }),
 ];
@@ -640,6 +707,59 @@ export const CURATED_AUDIO_MODELS: ModelDescriptor[] = [
       voices: [{ id: "default", label: "Default" }],
     },
   }),
+  createTtsModel({
+    id: "onnx-community/Kokoro-82M-v1.0-ONNX",
+    label: "Kokoro 82M",
+    summary:
+      "Compact browser TTS model with stronger natural voices and multiple speaker presets.",
+    publisher: "onnx-community",
+    paramsLabel: "82M params",
+    parameterTier: "S",
+    estimatedDownloadLabel: "~95 MB download",
+    category: "audio_smaller",
+    tested: false,
+    hf: {
+      modelId: "onnx-community/Kokoro-82M-v1.0-ONNX",
+      pipelineTag: "text-to-speech",
+      libraryName: "transformers.js",
+      tags: ["transformers.js", "text-to-speech", "kokoro"],
+      baseModel: "hexgrad/Kokoro-82M",
+    },
+    runtime: {
+      contextWindowTokens: 0,
+      preferredDtype: "q4f16",
+      fallbackDtype: "q4",
+      audioLoaderKind: "pipeline_tts",
+      audioSampleRate: 24000,
+      defaultVoice: "af_bella",
+      voices: [
+        {
+          id: "af_bella",
+          label: "Bella",
+          speakerEmbeddingsUrl:
+            "https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main/voices/af_bella.bin",
+        },
+        {
+          id: "af_sarah",
+          label: "Sarah",
+          speakerEmbeddingsUrl:
+            "https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main/voices/af_sarah.bin",
+        },
+        {
+          id: "am_michael",
+          label: "Michael",
+          speakerEmbeddingsUrl:
+            "https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main/voices/am_michael.bin",
+        },
+        {
+          id: "bf_emma",
+          label: "Emma",
+          speakerEmbeddingsUrl:
+            "https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main/voices/bf_emma.bin",
+        },
+      ],
+    },
+  }),
 ];
 
 export const CURATED_AUDIO_MODELS_BY_ID = Object.fromEntries(

src/test/audio-helpers.test.ts

@@ -3,7 +3,9 @@ import { describe, expect, it } from "vitest";
 import {
   EMPTY_TRANSCRIPTION_MESSAGE,
   normalizeTranscriptionText,
+  resolveVoiceSpeakerEmbeddings,
 } from "../app/audio-helpers";
+import { createModelDescriptor } from "./factories";
 
 describe("audio transcription helpers", () => {
   it("keeps direct transcription text when available", () => {
@@ -27,4 +29,32 @@ describe("audio transcription helpers", () => {
     expect(EMPTY_TRANSCRIPTION_MESSAGE).toMatch(/No speech was detected/i);
     expect(normalizeTranscriptionText("", [])).toBe("");
   });
+
+  it("resolves voice-specific speaker embeddings when a named voice is selected", () => {
+    const model = createModelDescriptor({
+      task: "tts",
+      runtime: {
+        voices: [
+          {
+            id: "voice-a",
+            label: "Voice A",
+            speakerEmbeddingsUrl: "https://example.com/voice-a.bin",
+          },
+          {
+            id: "voice-b",
+            label: "Voice B",
+            speakerEmbeddingsUrl: "https://example.com/voice-b.bin",
+          },
+        ],
+        speakerEmbeddingsUrl: "https://example.com/default.bin",
+      },
+    });
+
+    expect(resolveVoiceSpeakerEmbeddings(model, "voice-b")).toBe(
+      "https://example.com/voice-b.bin",
+    );
+    expect(resolveVoiceSpeakerEmbeddings(model, "missing")).toBe(
+      "https://example.com/default.bin",
+    );
+  });
 });

src/test/factories.ts

@@ -17,7 +17,7 @@ export const createDeviceCapabilities = (
 });
 
 export const createModelDescriptor = (
-  overrides: Partial<ModelDescriptor> & {
+  overrides: Partial<Omit<ModelDescriptor, "hf" | "runtime">> & {
     hf?: Partial<ModelDescriptor["hf"]>;
     runtime?: Partial<ModelDescriptor["runtime"]>;
   } = {},

src/test/model-helpers.test.ts

@@ -8,7 +8,7 @@ import {
   getFallbackThreadModel,
   getRecommendedModel,
 } from "../app/model-helpers";
-import { CURATED_MODELS_BY_ID } from "../models";
+import { CURATED_AUDIO_MODELS_BY_ID, CURATED_MODELS_BY_ID } from "../models";
 import { saveLastModel } from "../storage";
 import { createDeviceCapabilities, createModelDescriptor } from "./factories";
 
@@ -44,10 +44,25 @@ describe("chat model defaults", () => {
     );
 
     expect(modelIds).toContain("onnx-community/Supertonic-TTS-ONNX");
+    expect(modelIds).toContain("onnx-community/Kokoro-82M-v1.0-ONNX");
     expect(modelIds).toContain("Xenova/speecht5_tts");
     expect(modelIds).not.toContain("onnx-community/moonshine-base-ONNX");
   });
 
+  it("keeps the lightweight Gemma default while exposing the new Gemma mobile variants", () => {
+    expect(getRecommendedModel(createDeviceCapabilities(), {})?.id).toBe(
+      "onnx-community/gemma-3-270m-it-ONNX",
+    );
+    expect(CURATED_MODELS_BY_ID["onnx-community/gemma-4-E2B-it-qat-mobile-ONNX"])
+      .toBeDefined();
+    expect(CURATED_MODELS_BY_ID["onnx-community/gemma-4-E4B-it-qat-mobile-ONNX"])
+      .toBeDefined();
+    expect(CURATED_MODELS_BY_ID["onnx-community/gemma-3n-E2B-it-ONNX"])
+      .toBeDefined();
+    expect(CURATED_AUDIO_MODELS_BY_ID["onnx-community/Kokoro-82M-v1.0-ONNX"])
+      .toBeDefined();
+  });
+
   it("filters recent audio models by task", () => {
     const recentModels = [
       createModelDescriptor({

src/types.ts

@@ -81,12 +81,24 @@ export type CuratedCategoryKey =
   | "audio_smaller"
   | "audio_desktop_experimental";
 export type Dtype = "q4f16" | "q4" | "q8" | "int8" | "uint8" | "fp16" | "fp32";
-export type VisionLoaderKind = "qwen3_5" | "lfm2_5_vl";
+export type VisionLoaderKind = "qwen3_5" | "lfm2_5_vl" | "gemma_any_to_any";
+export type VisionRuntimeDtype = Dtype | "auto" | "bnb4" | "q2f16" | "q8f16";
 export type AudioLoaderKind =
   | "pipeline_asr"
   | "pipeline_tts"
   | "speecht5_tts"
   | "supertonic_tts";
+export type VoiceDescriptor = {
+  id: string;
+  label: string;
+  speakerEmbeddingsUrl?: string;
+};
+export type VisionDtypeConfig = Partial<
+  Record<
+    "audio_encoder" | "decoder_model_merged" | "embed_tokens" | "vision_encoder",
+    VisionRuntimeDtype
+  >
+>;
 export type ChatPersistenceStatus =
   | "ready"
   | "fallback_local_storage"
@@ -127,12 +139,13 @@ export type ModelRuntimeConfig = {
   contextWindowTokens: number;
   chatTemplateOptions?: Record<string, boolean | number | string>;
   visionLoaderKind?: VisionLoaderKind;
+  visionDtypes?: VisionDtypeConfig;
   audioLoaderKind?: AudioLoaderKind;
   supportsTimestamps?: boolean;
   audioSampleRate?: number;
   defaultLanguage?: string;
   defaultVoice?: string;
-  voices?: Array<{ id: string; label: string }>;
+  voices?: VoiceDescriptor[];
   speakerEmbeddingsUrl?: string;
 };
 

src/worker/generation.ts

@@ -5,6 +5,7 @@ import {
   TextStreamer,
 } from "@huggingface/transformers";
 
+import { resolveVoiceSpeakerEmbeddings } from "../app/audio-helpers";
 import type {
   GenerationOptions,
   ModelDescriptor,
@@ -389,7 +390,7 @@ export const generateSpeech = async ({
   model,
   requestId,
   text,
-  voice: _voice,
+  voice,
   speed,
 }: GenerateSpeechParams) => {
   await session.ensureModelReady(model);
@@ -409,8 +410,9 @@ export const generateSpeech = async ({
   }
 
   const output = await speechSynthesizer(text, {
-    speaker_embeddings: model.runtime.speakerEmbeddingsUrl,
+    speaker_embeddings: resolveVoiceSpeakerEmbeddings(model, voice),
     speed,
+    voice,
   });
 
   return {

src/worker/model-session.ts

@@ -429,6 +429,32 @@ export const createModelSession = (postMessageToUi: WorkerMessagePoster) => {
       } satisfies LoadResources;
     }
 
+    if (model.runtime.visionLoaderKind === "gemma_any_to_any") {
+      const [nextProcessor, nextVisionModel] = await Promise.all([
+        AutoProcessor.from_pretrained(model.hf.modelId, {
+          progress_callback: progressHandler,
+        }),
+        AutoModelForImageTextToText.from_pretrained(model.hf.modelId, {
+          dtype: (model.runtime.visionDtypes ?? {
+            embed_tokens: "q4f16",
+            vision_encoder: "fp16",
+            decoder_model_merged: "q4f16",
+            audio_encoder: "q4f16",
+          }) as any,
+          device: "webgpu",
+          progress_callback: progressHandler,
+        }),
+      ]);
+
+      return {
+        textGenerator: null,
+        processor: nextProcessor,
+        visionModel: nextVisionModel as VisionModelInstance,
+        speechRecognizer: null,
+        speechSynthesizer: null,
+      } satisfies LoadResources;
+    }
+
     throw new Error(
       "This vision loader is not supported in the browser worker yet.",
     );