fix am2 and add vvad

Author: yujonglee

.github/actions/argmax_sdk_setup/action.yaml

@@ -1,4 +1,4 @@
-# https://docs.argmaxinc.com/guides/upgrading-to-pro-sdk
+# https://app.argmaxinc.com/docs/guides/upgrading-to-pro-sdk
 # Sets up Argmax Pro SDK registry access for Swift Package Manager
 
 name: "Setup Argmax SDK"
@@ -10,12 +10,8 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: Setup Argmax Registry
-      shell: bash
-      run: |
-        swift package-registry set --global --scope argmaxinc https://api.argmaxinc.com/v1/sdk
+    - shell: bash
+      run: swift package-registry set --global --scope argmaxinc https://api.argmaxinc.com/v1/sdk
 
-    - name: Authenticate with Argmax
-      shell: bash
-      run: |
-        swift package-registry login https://api.argmaxinc.com/v1/sdk/login --token "${{ inputs.api-token }}"
+    - shell: bash
+      run: swift package-registry login https://api.argmaxinc.com/v1/sdk/login --token "${{ inputs.api-token }}"

.github/workflows/desktop_cd.yaml

@@ -99,6 +99,8 @@ jobs:
         uses: ./.github/actions/argmax_sdk_setup
         with:
           api-token: ${{ secrets.AM_SECRET_TOKEN }}
+      - if: ${{ matrix.include_am }}
+        run: cargo test -p am2
       - run: pnpm -F ui build
       - if: ${{ matrix.include_am }}
         run: |

Cargo.lock

@@ -79,6 +79,7 @@ hypr-turso = { path = "crates/turso", package = "turso" }
 hypr-vad-ext = { path = "crates/vad-ext", package = "vad-ext" }
 hypr-vad2 = { path = "crates/vad2", package = "vad2" }
 hypr-vad3 = { path = "crates/vad3", package = "vad3" }
+hypr-vvad = { path = "crates/vvad", package = "vvad" }
 hypr-whisper = { path = "crates/whisper", package = "whisper" }
 hypr-whisper-local = { path = "crates/whisper-local", package = "whisper-local" }
 hypr-whisper-local-model = { path = "crates/whisper-local-model", package = "whisper-local-model" }
@@ -256,7 +257,7 @@ objc2-user-notifications = "0.3"
 
 tokenizers = "0.21.4"
 
-swift-rs = { git = "https://github.com/yujonglee/swift-rs", branch = "fix-framework" }
+swift-rs = { git = "https://github.com/yujonglee/swift-rs", rev = "8c39ff4" }
 sysinfo = "0.37.0"
 
 [patch.crates-io]

Cargo.toml

@@ -3,6 +3,7 @@ fn main() {
     {
         swift_rs::SwiftLinker::new("13.0")
             .with_package("swift-lib", "./swift-lib/")
+            .with_framework("ArgmaxSDK")
             .link();
     }
 }

crates/am2/build.rs

@@ -1,26 +1,107 @@
-use swift_rs::swift;
+use std::sync::OnceLock;
 
-swift!(fn initialize_am2_sdk());
+use swift_rs::{swift, SRArray, SRObject, SRString};
 
-swift!(fn check_am2_ready() -> bool);
+swift!(fn initialize_am2_sdk(api_key: &SRString));
+swift!(fn am2_vad_init() -> bool);
+swift!(fn am2_vad_detect(samples_ptr: *const f32, samples_len: i64) -> SRObject<VadResultArray>);
+swift!(fn am2_vad_index_to_seconds(index: i64) -> f32);
+
+static SDK_INITIALIZED: OnceLock<()> = OnceLock::new();
+
+#[repr(C)]
+pub struct VadResultArray {
+    pub data: SRArray<bool>,
+}
 
 pub fn init() {
-    unsafe {
-        initialize_am2_sdk();
-    }
+    SDK_INITIALIZED.get_or_init(|| {
+        let api_key = std::env::var("AM_API_KEY").unwrap_or_default();
+        let api_key = SRString::from(api_key.as_str());
+        unsafe {
+            initialize_am2_sdk(&api_key);
+        }
+    });
 }
 
 pub fn is_ready() -> bool {
-    unsafe { check_am2_ready() }
+    SDK_INITIALIZED.get().is_some()
+}
+
+pub mod vad {
+    use std::sync::OnceLock;
+
+    use super::*;
+
+    static VAD_INITIALIZED: OnceLock<bool> = OnceLock::new();
+
+    pub fn init() -> bool {
+        *VAD_INITIALIZED.get_or_init(|| unsafe { am2_vad_init() })
+    }
+
+    pub fn is_ready() -> bool {
+        VAD_INITIALIZED.get().copied().unwrap_or(false)
+    }
+
+    pub fn detect(samples: &[f32]) -> Vec<bool> {
+        let result = unsafe { am2_vad_detect(samples.as_ptr(), samples.len() as i64) };
+        result.data.as_slice().to_vec()
+    }
+
+    pub fn index_to_seconds(index: usize) -> f32 {
+        unsafe { am2_vad_index_to_seconds(index as i64) }
+    }
+
+    #[derive(Debug, Clone)]
+    pub struct VoiceSegment {
+        pub start_seconds: f32,
+        pub end_seconds: f32,
+    }
+
+    pub fn detect_segments(samples: &[f32]) -> Vec<VoiceSegment> {
+        let voice_activity = detect(samples);
+        let mut segments = Vec::new();
+        let mut in_voice = false;
+        let mut segment_start = 0.0;
+
+        for (i, &is_voice) in voice_activity.iter().enumerate() {
+            if is_voice && !in_voice {
+                segment_start = index_to_seconds(i);
+                in_voice = true;
+            } else if !is_voice && in_voice {
+                segments.push(VoiceSegment {
+                    start_seconds: segment_start,
+                    end_seconds: index_to_seconds(i),
+                });
+                in_voice = false;
+            }
+        }
+
+        if in_voice {
+            segments.push(VoiceSegment {
+                start_seconds: segment_start,
+                end_seconds: index_to_seconds(voice_activity.len()),
+            });
+        }
+
+        segments
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_am2_swift_compilation() {
+    fn test_am2_sdk_init() {
         init();
         assert!(is_ready());
     }
+
+    #[test]
+    fn test_am2_vad_init() {
+        init();
+        assert!(vad::init());
+        assert!(vad::is_ready());
+    }
 }

crates/am2/src/lib.rs

@@ -2,16 +2,69 @@ import Argmax
 import Foundation
 import SwiftRs
 
-private var isAM2Ready = false
+private var vadInstance: VoiceActivityDetector?
 
 @_cdecl("initialize_am2_sdk")
-public func initialize_am2_sdk() {
-  isAM2Ready = true
-  print("AM2 SDK initialized successfully")
+public func initialize_am2_sdk(apiKey: SRString) {
+  let semaphore = DispatchSemaphore(value: 0)
+
+  Task {
+    let key = apiKey.toString()
+    if !key.isEmpty {
+      await ArgmaxSDK.with(ArgmaxConfig(apiKey: key))
+    }
+    semaphore.signal()
+  }
+
+  semaphore.wait()
+}
+
+public class VadResultArray: NSObject {
+  var data: SRArray<Bool>
+
+  init(_ data: [Bool]) {
+    self.data = SRArray(data)
+  }
+}
+
+@_cdecl("am2_vad_init")
+public func am2_vad_init() -> Bool {
+  let semaphore = DispatchSemaphore(value: 0)
+  var success = false
+
+  Task {
+    do {
+      vadInstance = try await VoiceActivityDetector.modelVAD()
+      success = true
+    } catch {
+      success = false
+    }
+    semaphore.signal()
+  }
+
+  semaphore.wait()
+  return success
 }
 
-@_cdecl("check_am2_ready")
-public func check_am2_ready() -> Bool {
-  return isAM2Ready
+@_cdecl("am2_vad_detect")
+public func am2_vad_detect(
+  samplesPtr: UnsafePointer<Float>,
+  samplesLen: Int
+) -> VadResultArray {
+  guard let vad = vadInstance else {
+    return VadResultArray([])
+  }
+
+  let audioArray = Array(UnsafeBufferPointer(start: samplesPtr, count: samplesLen))
+  let voiceSegments = vad.voiceActivity(in: audioArray)
+
+  return VadResultArray(voiceSegments)
 }
 
+@_cdecl("am2_vad_index_to_seconds")
+public func am2_vad_index_to_seconds(index: Int) -> Float {
+  guard let vad = vadInstance else {
+    return 0.0
+  }
+  return vad.voiceActivityIndexToSeconds(index)
+}

crates/am2/swift-lib/src/lib.swift

@@ -0,0 +1,10 @@
+[package]
+name = "vvad"
+version = "0.1.0"
+edition = "2021"
+
+[target.'cfg(all(target_os = "macos", target_arch = "aarch64"))'.dependencies]
+hypr-am2 = { workspace = true }
+
+[target.'cfg(not(all(target_os = "macos", target_arch = "aarch64")))'.dependencies]
+earshot = { workspace = true }