react-native-nitro-mlx 0.3.0 → 0.4.0

package/src/models.ts

@@ -5,6 +5,8 @@ export enum ModelFamily {
   Phi = 'Phi',
   SmolLM = 'SmolLM',
   OpenELM = 'OpenELM',
+  PocketTTS = 'PocketTTS',
+  GLMASR = 'GLMASR',
 }
 
 export enum ModelProvider {
@@ -14,9 +16,12 @@ export enum ModelProvider {
   Microsoft = 'Microsoft',
   HuggingFace = 'HuggingFace',
   Apple = 'Apple',
+  Kyutai = 'Kyutai',
 }
 
-export type ModelQuantization = '4bit' | '8bit'
+export type ModelQuantization = '4bit' | '8bit' | 'bf16'
+
+export type ModelType = 'llm' | 'tts' | 'stt'
 
 export interface ModelInfo {
   id: MLXModel
@@ -26,6 +31,7 @@ export interface ModelInfo {
   quantization: ModelQuantization
   displayName: string
   downloadSize: number
+  type: ModelType
 }
 
 export enum MLXModel {
@@ -72,6 +78,14 @@ export enum MLXModel {
   OpenELM_1_1B_8bit = 'mlx-community/OpenELM-1_1B-8bit',
   OpenELM_3B_4bit = 'mlx-community/OpenELM-3B-4bit',
   OpenELM_3B_8bit = 'mlx-community/OpenELM-3B-8bit',
+
+  // PocketTTS (Kyutai) - 44.6M
+  PocketTTS = 'mlx-community/pocket-tts',
+  PocketTTS_8bit = 'mlx-community/pocket-tts-8bit',
+  PocketTTS_4bit = 'mlx-community/pocket-tts-4bit',
+
+  // GLM-ASR (GLMASR) - Speech-to-Text
+  GLM_ASR_Nano_4bit = 'mlx-community/GLM-ASR-Nano-2512-4bit',
 }
 
 export const MLXModels: ModelInfo[] = [
@@ -83,6 +97,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Llama 3.2 1B Instruct (4-bit)',
     downloadSize: 1407777762,
+    type: 'llm',
   },
   {
     id: MLXModel.Llama_3_2_1B_Instruct_8bit,
@@ -92,6 +107,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Llama 3.2 1B Instruct (8-bit)',
     downloadSize: 1313157436,
+    type: 'llm',
   },
   {
     id: MLXModel.Llama_3_2_3B_Instruct_4bit,
@@ -101,6 +117,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Llama 3.2 3B Instruct (4-bit)',
     downloadSize: 2019397474,
+    type: 'llm',
   },
   {
     id: MLXModel.Llama_3_2_3B_Instruct_8bit,
@@ -110,6 +127,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Llama 3.2 3B Instruct (8-bit)',
     downloadSize: 3413784042,
+    type: 'llm',
   },
   {
     id: MLXModel.Qwen2_5_0_5B_Instruct_4bit,
@@ -119,6 +137,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Qwen 2.5 0.5B Instruct (4-bit)',
     downloadSize: 278064920,
+    type: 'llm',
   },
   {
     id: MLXModel.Qwen2_5_0_5B_Instruct_8bit,
@@ -128,6 +147,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Qwen 2.5 0.5B Instruct (8-bit)',
     downloadSize: 525045902,
+    type: 'llm',
   },
   {
     id: MLXModel.Qwen2_5_1_5B_Instruct_4bit,
@@ -137,6 +157,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Qwen 2.5 1.5B Instruct (4-bit)',
     downloadSize: 868628559,
+    type: 'llm',
   },
   {
     id: MLXModel.Qwen2_5_1_5B_Instruct_8bit,
@@ -146,6 +167,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Qwen 2.5 1.5B Instruct (8-bit)',
     downloadSize: 1640414038,
+    type: 'llm',
   },
   {
     id: MLXModel.Qwen2_5_3B_Instruct_4bit,
@@ -155,6 +177,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Qwen 2.5 3B Instruct (4-bit)',
     downloadSize: 1736293090,
+    type: 'llm',
   },
   {
     id: MLXModel.Qwen2_5_3B_Instruct_8bit,
@@ -164,6 +187,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Qwen 2.5 3B Instruct (8-bit)',
     downloadSize: 3279142142,
+    type: 'llm',
   },
   {
     id: MLXModel.Qwen3_1_7B_4bit,
@@ -173,6 +197,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Qwen 3 1.7B (4-bit)',
     downloadSize: 979502864,
+    type: 'llm',
   },
   {
     id: MLXModel.Qwen3_1_7B_8bit,
@@ -182,6 +207,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Qwen 3 1.7B (8-bit)',
     downloadSize: 1839729195,
+    type: 'llm',
   },
   {
     id: MLXModel.Gemma_3_1B_IT_4bit,
@@ -191,6 +217,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Gemma 3 1B IT (4-bit)',
     downloadSize: 770650946,
+    type: 'llm',
   },
   {
     id: MLXModel.Gemma_3_1B_IT_8bit,
@@ -200,6 +227,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Gemma 3 1B IT (8-bit)',
     downloadSize: 1421522471,
+    type: 'llm',
   },
   {
     id: MLXModel.Phi_3_5_Mini_Instruct_4bit,
@@ -209,6 +237,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Phi 3.5 Mini Instruct (4-bit)',
     downloadSize: 2150195856,
+    type: 'llm',
   },
   {
     id: MLXModel.Phi_3_5_Mini_Instruct_8bit,
@@ -218,6 +247,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Phi 3.5 Mini Instruct (8-bit)',
     downloadSize: 4060636056,
+    type: 'llm',
   },
   {
     id: MLXModel.Phi_4_Mini_Instruct_4bit,
@@ -227,6 +257,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'Phi 4 Mini Instruct (4-bit)',
     downloadSize: 2173624891,
+    type: 'llm',
   },
   {
     id: MLXModel.Phi_4_Mini_Instruct_8bit,
@@ -236,6 +267,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'Phi 4 Mini Instruct (8-bit)',
     downloadSize: 4091536167,
+    type: 'llm',
   },
   {
     id: MLXModel.SmolLM_1_7B_Instruct_4bit,
@@ -245,6 +277,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'SmolLM 1.7B Instruct (4-bit)',
     downloadSize: 962855374,
+    type: 'llm',
   },
   {
     id: MLXModel.SmolLM_1_7B_Instruct_8bit,
@@ -254,6 +287,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'SmolLM 1.7B Instruct (8-bit)',
     downloadSize: 1818493993,
+    type: 'llm',
   },
   {
     id: MLXModel.SmolLM2_1_7B_Instruct_4bit,
@@ -263,6 +297,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'SmolLM2 1.7B Instruct (4-bit)',
     downloadSize: 980000000,
+    type: 'llm',
   },
   {
     id: MLXModel.SmolLM2_1_7B_Instruct_8bit,
@@ -272,6 +307,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'SmolLM2 1.7B Instruct (8-bit)',
     downloadSize: 1850000000,
+    type: 'llm',
   },
   {
     id: MLXModel.OpenELM_1_1B_4bit,
@@ -281,6 +317,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'OpenELM 1.1B (4-bit)',
     downloadSize: 608162655,
+    type: 'llm',
   },
   {
     id: MLXModel.OpenELM_1_1B_8bit,
@@ -290,6 +327,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'OpenELM 1.1B (8-bit)',
     downloadSize: 1148048397,
+    type: 'llm',
   },
   {
     id: MLXModel.OpenELM_3B_4bit,
@@ -299,6 +337,7 @@ export const MLXModels: ModelInfo[] = [
     quantization: '4bit',
     displayName: 'OpenELM 3B (4-bit)',
     downloadSize: 1650000000,
+    type: 'llm',
   },
   {
     id: MLXModel.OpenELM_3B_8bit,
@@ -308,5 +347,46 @@ export const MLXModels: ModelInfo[] = [
     quantization: '8bit',
     displayName: 'OpenELM 3B (8-bit)',
     downloadSize: 3100000000,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.PocketTTS,
+    family: ModelFamily.PocketTTS,
+    provider: ModelProvider.Kyutai,
+    parameters: '44.6M',
+    quantization: 'bf16',
+    displayName: 'PocketTTS (bf16)',
+    downloadSize: 180000000,
+    type: 'tts',
+  },
+  {
+    id: MLXModel.PocketTTS_8bit,
+    family: ModelFamily.PocketTTS,
+    provider: ModelProvider.Kyutai,
+    parameters: '44.6M',
+    quantization: '8bit',
+    displayName: 'PocketTTS (8-bit)',
+    downloadSize: 140000000,
+    type: 'tts',
+  },
+  {
+    id: MLXModel.PocketTTS_4bit,
+    family: ModelFamily.PocketTTS,
+    provider: ModelProvider.Kyutai,
+    parameters: '44.6M',
+    quantization: '4bit',
+    displayName: 'PocketTTS (4-bit)',
+    downloadSize: 80000000,
+    type: 'tts',
+  },
+  {
+    id: MLXModel.GLM_ASR_Nano_4bit,
+    family: ModelFamily.GLMASR,
+    provider: ModelProvider.Alibaba,
+    parameters: '1B',
+    quantization: '4bit',
+    displayName: 'GLM-ASR Nano (4-bit)',
+    downloadSize: 600000000,
+    type: 'stt',
   },
 ]

package/src/specs/LLM.nitro.ts

@@ -4,16 +4,81 @@ import type { AnyMap, HybridObject } from 'react-native-nitro-modules'
  * Statistics from the last text generation.
  */
 export interface GenerationStats {
-  /** Total number of tokens generated */
   tokenCount: number
-  /** Generation speed in tokens per second */
   tokensPerSecond: number
-  /** Time in milliseconds until the first token was generated */
   timeToFirstToken: number
-  /** Total generation time in milliseconds */
   totalTime: number
+  toolExecutionTime: number
 }
 
+export interface GenerationStartEvent {
+  type: 'generation_start'
+  timestamp: number
+}
+
+export interface TokenEvent {
+  type: 'token'
+  token: string
+}
+
+export interface ThinkingStartEvent {
+  type: 'thinking_start'
+  timestamp: number
+}
+
+export interface ThinkingChunkEvent {
+  type: 'thinking_chunk'
+  chunk: string
+}
+
+export interface ThinkingEndEvent {
+  type: 'thinking_end'
+  content: string
+  timestamp: number
+}
+
+export interface ToolCallStartEvent {
+  type: 'tool_call_start'
+  id: string
+  name: string
+  arguments: string
+}
+
+export interface ToolCallExecutingEvent {
+  type: 'tool_call_executing'
+  id: string
+}
+
+export interface ToolCallCompletedEvent {
+  type: 'tool_call_completed'
+  id: string
+  result: string
+}
+
+export interface ToolCallFailedEvent {
+  type: 'tool_call_failed'
+  id: string
+  error: string
+}
+
+export interface GenerationEndEvent {
+  type: 'generation_end'
+  content: string
+  stats: GenerationStats
+}
+
+export type StreamEvent =
+  | GenerationStartEvent
+  | TokenEvent
+  | ThinkingStartEvent
+  | ThinkingChunkEvent
+  | ThinkingEndEvent
+  | ToolCallStartEvent
+  | ToolCallExecutingEvent
+  | ToolCallCompletedEvent
+  | ToolCallFailedEvent
+  | GenerationEndEvent
+
 export interface LLMMessage {
   role: string
   content: string
@@ -85,6 +150,11 @@ export interface LLM extends HybridObject<{ ios: 'swift' }> {
     onToolCall?: (toolName: string, args: string) => void,
   ): Promise<string>
 
+  streamWithEvents(
+    prompt: string,
+    onEvent: (eventJson: string) => void,
+  ): Promise<string>
+
   /**
    * Stop the current generation.
    */

package/src/specs/STT.nitro.ts

@@ -0,0 +1,35 @@
+import type { HybridObject } from 'react-native-nitro-modules'
+
+export interface STTLoadOptions {
+  onProgress?: (progress: number) => void
+}
+
+export interface STTTranscriptionInfo {
+  promptTokens: number
+  generationTokens: number
+  tokensPerSecond: number
+  prefillTime: number
+  generateTime: number
+}
+
+export interface STT extends HybridObject<{ ios: 'swift' }> {
+  readonly isLoaded: boolean
+  readonly isTranscribing: boolean
+  readonly isListening: boolean
+  readonly modelId: string
+
+  load(modelId: string, options?: STTLoadOptions): Promise<void>
+
+  transcribe(audio: ArrayBuffer): Promise<string>
+  transcribeStream(
+    audio: ArrayBuffer,
+    onToken: (token: string) => void
+  ): Promise<string>
+
+  startListening(): Promise<void>
+  transcribeBuffer(): Promise<string>
+  stopListening(): Promise<string>
+
+  stop(): void
+  unload(): void
+}

package/src/specs/TTS.nitro.ts

@@ -0,0 +1,30 @@
+import type { HybridObject } from 'react-native-nitro-modules'
+
+export interface TTSLoadOptions {
+  onProgress?: (progress: number) => void
+}
+
+export interface TTSGenerateOptions {
+  voice?: string
+  speed?: number
+}
+
+export interface TTS extends HybridObject<{ ios: 'swift' }> {
+  readonly isLoaded: boolean
+  readonly isGenerating: boolean
+  readonly modelId: string
+  readonly sampleRate: number
+
+  load(modelId: string, options?: TTSLoadOptions): Promise<void>
+  generate(
+    text: string,
+    options?: TTSGenerateOptions
+  ): Promise<ArrayBuffer>
+  stream(
+    text: string,
+    onAudioChunk: (audio: ArrayBuffer) => void,
+    options?: TTSGenerateOptions
+  ): Promise<void>
+  stop(): void
+  unload(): void
+}

package/src/stt.ts

@@ -0,0 +1,67 @@
+import { NitroModules } from 'react-native-nitro-modules'
+import type {
+  STT as STTSpec,
+  STTLoadOptions,
+} from './specs/STT.nitro'
+
+let instance: STTSpec | null = null
+
+function getInstance(): STTSpec {
+  if (!instance) {
+    instance = NitroModules.createHybridObject<STTSpec>('STT')
+  }
+  return instance
+}
+
+export const STT = {
+  load(modelId: string, options?: STTLoadOptions): Promise<void> {
+    return getInstance().load(modelId, options)
+  },
+
+  transcribe(audio: ArrayBuffer): Promise<string> {
+    return getInstance().transcribe(audio)
+  },
+
+  transcribeStream(
+    audio: ArrayBuffer,
+    onToken: (token: string) => void
+  ): Promise<string> {
+    return getInstance().transcribeStream(audio, onToken)
+  },
+
+  startListening(): Promise<void> {
+    return getInstance().startListening()
+  },
+
+  transcribeBuffer(): Promise<string> {
+    return getInstance().transcribeBuffer()
+  },
+
+  stopListening(): Promise<string> {
+    return getInstance().stopListening()
+  },
+
+  stop(): void {
+    getInstance().stop()
+  },
+
+  unload(): void {
+    getInstance().unload()
+  },
+
+  get isLoaded(): boolean {
+    return getInstance().isLoaded
+  },
+
+  get isTranscribing(): boolean {
+    return getInstance().isTranscribing
+  },
+
+  get isListening(): boolean {
+    return getInstance().isListening
+  },
+
+  get modelId(): string {
+    return getInstance().modelId
+  },
+}

package/src/tts.ts

@@ -0,0 +1,60 @@
+import { NitroModules } from 'react-native-nitro-modules'
+import type {
+  TTS as TTSSpec,
+  TTSLoadOptions,
+  TTSGenerateOptions,
+} from './specs/TTS.nitro'
+
+let instance: TTSSpec | null = null
+
+function getInstance(): TTSSpec {
+  if (!instance) {
+    instance = NitroModules.createHybridObject<TTSSpec>('TTS')
+  }
+  return instance
+}
+
+export const TTS = {
+  load(modelId: string, options?: TTSLoadOptions): Promise<void> {
+    return getInstance().load(modelId, options)
+  },
+
+  generate(
+    text: string,
+    options?: TTSGenerateOptions
+  ): Promise<ArrayBuffer> {
+    return getInstance().generate(text, options)
+  },
+
+  stream(
+    text: string,
+    onAudioChunk: (audio: ArrayBuffer) => void,
+    options?: TTSGenerateOptions
+  ): Promise<void> {
+    return getInstance().stream(text, onAudioChunk, options)
+  },
+
+  stop(): void {
+    getInstance().stop()
+  },
+
+  unload(): void {
+    getInstance().unload()
+  },
+
+  get isLoaded(): boolean {
+    return getInstance().isLoaded
+  },
+
+  get isGenerating(): boolean {
+    return getInstance().isGenerating
+  },
+
+  get modelId(): string {
+    return getInstance().modelId
+  },
+
+  get sampleRate(): number {
+    return getInstance().sampleRate
+  },
+}