feat: initialize WhisperKit on app startup (with config flag)

matthiasn · Aug 12, 2024 · 61a5156 · 61a5156
1 parent ad9ed0c
commit 61a5156
Show file tree

Hide file tree

Showing 8 changed files with 78 additions and 6 deletions.
diff --git a/ios/Runner/WhisperKitRunner.swift b/ios/Runner/WhisperKitRunner.swift
@@ -25,6 +25,22 @@ public class WhisperKitRunner: NSObject, FlutterStreamHandler {
 
         transcriptionChannel.setMethodCallHandler { (call, result) in
             switch call.method {
+
+            case "initialize":
+                if (self.whisperKit == nil) {
+                    if (self.eventSink != nil) {
+                        self.eventSink!(["Initializing model...", ""])
+                    }
+                }
+
+                Task {
+                    if (self.whisperKit == nil) {
+                        self.whisperKit = try? await WhisperKit(model: self.model,
+                                                                verbose: true,
+                                                                prewarm: true)
+                    }
+                }
+
             case "transcribe":
                 guard let args = call.arguments as? [String: Any] else { return }
                 let audioFilePath = args["audioFilePath"] as! String

diff --git a/lib/database/journal_db/config_flags.dart b/lib/database/journal_db/config_flags.dart
@@ -42,6 +42,13 @@ Future<void> initConfigFlags(
       status: false,
     ),
   );
+  await db.insertFlagIfNotExists(
+    const ConfigFlag(
+      name: initializeAsrModelOnStartup,
+      description: 'Initialize WhisperKit ASR model on startup',
+      status: false,
+    ),
+  );
   await db.insertFlagIfNotExists(
     const ConfigFlag(
       name: enableMatrixFlag,

diff --git a/lib/pages/settings/flags_page.dart b/lib/pages/settings/flags_page.dart
@@ -25,6 +25,7 @@ class FlagsPage extends StatelessWidget {
           attemptEmbedding,
           enableNotificationsFlag,
           autoTranscribeFlag,
+          initializeAsrModelOnStartup,
           recordLocationFlag,
           allowInvalidCertFlag,
           enableMatrixFlag,

diff --git a/lib/services/asr_service.dart b/lib/services/asr_service.dart
@@ -8,17 +8,21 @@ import 'package:ffmpeg_kit_flutter/return_code.dart';
 import 'package:flutter/cupertino.dart';
 import 'package:flutter/services.dart';
 import 'package:lotti/classes/journal_entities.dart';
+import 'package:lotti/database/database.dart';
 import 'package:lotti/database/logging_db.dart';
 import 'package:lotti/database/settings_db.dart';
 import 'package:lotti/get_it.dart';
 import 'package:lotti/logic/persistence_logic.dart';
 import 'package:lotti/utils/audio_utils.dart';
+import 'package:lotti/utils/consts.dart';
 import 'package:path/path.dart' as p;
 import 'package:path_provider/path_provider.dart';
 
 class AsrService {
   AsrService() {
     _loadSelectedModel();
+    _initialize();
+
     eventChannel.receiveBroadcastStream().listen(
           _onEvent,
           onError: _onError,
@@ -175,6 +179,28 @@ class AsrService {
     running = false;
   }
 
+  Future<void> _initialize() async {
+    final initializeOnStartup = await getIt<JournalDb>().getConfigFlag(
+      initializeAsrModelOnStartup,
+    );
+
+    if (!initializeOnStartup) {
+      return;
+    }
+
+    getIt<LoggingDb>().captureEvent(
+      'initializing WhisperKit',
+      domain: 'ASR',
+      subDomain: 'initialize',
+    );
+
+    try {
+      unawaited(methodChannel.invokeMethod('initialize'));
+    } on PlatformException catch (e) {
+      captureException(e);
+    }
+  }
+
   Future<bool> enqueue({required JournalAudio entry}) async {
     final isQueueEmpty = queue.isEmpty;
     queue.add(entry);

diff --git a/lib/utils/consts.dart b/lib/utils/consts.dart
@@ -4,6 +4,7 @@ const allowInvalidCertFlag = 'allow_invalid_cert';
 const enableSyncFlag = 'enable_sync';
 const recordLocationFlag = 'record_location';
 const autoTranscribeFlag = 'auto_transcribe';
+const initializeAsrModelOnStartup = 'initialize_asr_model_on_startup';
 const enableMatrixFlag = 'enable_matrix';
 const resendAttachments = 'resend_attachments';
 const attemptEmbedding = 'attempt_embedding';
diff --git a/macos/Runner/WhisperKitRunner.swift b/macos/Runner/WhisperKitRunner.swift
@@ -25,26 +25,41 @@ public class WhisperKitRunner: NSObject, FlutterStreamHandler {
 
         transcriptionChannel.setMethodCallHandler { (call, result) in
             switch call.method {
+
+            case "initialize":
+                if (self.whisperKit == nil) {
+                    if (self.eventSink != nil) {
+                        self.eventSink!(["Initializing model...", ""])
+                    }
+                }
+
+                Task {
+                    if (self.whisperKit == nil) {
+                        self.whisperKit = try? await WhisperKit(model: self.model,
+                                                                verbose: true,
+                                                                prewarm: true)
+                    }
+                }
+
             case "transcribe":
                 guard let args = call.arguments as? [String: Any] else { return }
                 let audioFilePath = args["audioFilePath"] as! String
                 let language = args["language"] as! String
                 let detectLanguage = language.isEmpty
-                
+
                 if (self.whisperKit == nil) {
                     if (self.eventSink != nil) {
                         self.eventSink!(["Initializing model...", ""])
                     }
                 }
-
-
+
                 Task {
                     if (self.whisperKit == nil) {
                         self.whisperKit = try? await WhisperKit(model: self.model,
                                                                 verbose: true,
                                                                 prewarm: true)
                     }
-                    
+
                     let transcription = try? await self.whisperKit!.transcribe(
                         audioPath: audioFilePath,
                         decodeOptions: DecodingOptions(
@@ -55,12 +70,13 @@ public class WhisperKitRunner: NSObject, FlutterStreamHandler {
                         ),
                         callback: self.sendTranscriptionProgressEvent
                     )
-                    
+
                     let text : String? = transcription?.first?.text
                     let detectedLanguage = transcription?.first?.language
                     let data = [detectedLanguage, self.model, text]
                     result(data)
                 }
+
             default:
                 result(FlutterMethodNotImplemented)
             }

diff --git a/pubspec.yaml b/pubspec.yaml
@@ -1,7 +1,7 @@
 name: lotti
 description: Achieve your goals and keep your data private with Lotti.
 publish_to: 'none'
-version: 0.9.496+2641
+version: 0.9.497+2642
 
 msix_config:
   display_name: LottiApp

diff --git a/test/database/database_test.dart b/test/database/database_test.dart
@@ -31,6 +31,11 @@ final expectedFlags = <ConfigFlag>{
     description: 'Automatically transcribe audio',
     status: false,
   ),
+  const ConfigFlag(
+    name: initializeAsrModelOnStartup,
+    description: 'Initialize WhisperKit ASR model on startup',
+    status: false,
+  ),
   const ConfigFlag(
     name: recordLocationFlag,
     description: 'Record geolocation?',