Add model prewarm service on system wake and app launch

This commit is contained in:
Beingpax 2025-12-17 21:23:30 +05:45
parent 078e02c503
commit a2f19b04c6
4 changed files with 187 additions and 8 deletions

View File

@ -0,0 +1,158 @@
import Foundation
import SwiftData
import os
import AppKit
@MainActor
final class ModelPrewarmService: ObservableObject {
private let whisperState: WhisperState
private let modelContext: ModelContext
private let logger = Logger(subsystem: "com.prakashjoshipax.voiceink", category: "ModelPrewarm")
// Services (initialized lazily)
private var localTranscriptionService: LocalTranscriptionService?
private var parakeetTranscriptionService: ParakeetTranscriptionService?
private let nativeAppleTranscriptionService = NativeAppleTranscriptionService()
private let cloudTranscriptionService = CloudTranscriptionService()
// Sample audio for prewarming
private let prewarmAudioURL = Bundle.main.url(forResource: "esc", withExtension: "wav")
// User preference key
private let prewarmEnabledKey = "PrewarmModelOnWake"
init(whisperState: WhisperState, modelContext: ModelContext) {
self.whisperState = whisperState
self.modelContext = modelContext
setupNotifications()
schedulePrewarmOnAppLaunch()
}
// MARK: - Notification Setup
private func setupNotifications() {
let center = NSWorkspace.shared.notificationCenter
// Trigger on wake from sleep
center.addObserver(
self,
selector: #selector(schedulePrewarm),
name: NSWorkspace.didWakeNotification,
object: nil
)
logger.notice("🌅 ModelPrewarmService initialized - listening for wake and app launch")
}
// MARK: - Trigger Handlers
/// Trigger on app launch (cold start)
private func schedulePrewarmOnAppLaunch() {
logger.notice("🌅 App launched, scheduling prewarm")
Task {
try? await Task.sleep(for: .seconds(3))
await performPrewarm()
}
}
/// Trigger on wake from sleep or screen unlock
@objc private func schedulePrewarm() {
logger.notice("🌅 Mac activity detected (wake/unlock), scheduling prewarm")
Task {
try? await Task.sleep(for: .seconds(3))
await performPrewarm()
}
}
// MARK: - Core Prewarming Logic
private func performPrewarm() async {
guard shouldPrewarm() else { return }
guard let audioURL = prewarmAudioURL else {
logger.error("❌ Prewarm audio file (esc.wav) not found")
return
}
guard let currentModel = whisperState.currentTranscriptionModel else {
logger.notice("🌅 No model selected, skipping prewarm")
return
}
logger.notice("🌅 Prewarming \(currentModel.displayName)")
let startTime = Date()
do {
// Initialize services lazily
if localTranscriptionService == nil {
localTranscriptionService = LocalTranscriptionService(
modelsDirectory: whisperState.modelsDirectory,
whisperState: whisperState
)
}
if parakeetTranscriptionService == nil {
parakeetTranscriptionService = ParakeetTranscriptionService()
}
// Run transcription to trigger model loading and ANE compilation
let transcribedText: String
switch currentModel.provider {
case .local:
transcribedText = try await localTranscriptionService!.transcribe(audioURL: audioURL, model: currentModel)
case .parakeet:
transcribedText = try await parakeetTranscriptionService!.transcribe(audioURL: audioURL, model: currentModel)
case .nativeApple:
transcribedText = try await nativeAppleTranscriptionService.transcribe(audioURL: audioURL, model: currentModel)
default:
transcribedText = try await cloudTranscriptionService.transcribe(audioURL: audioURL, model: currentModel)
}
let duration = Date().timeIntervalSince(startTime)
// Save for telemetry
let transcription = Transcription(
text: "[PREWARM] \(transcribedText)",
duration: 1.0,
audioFileURL: audioURL.absoluteString,
transcriptionModelName: currentModel.displayName,
transcriptionDuration: duration
)
modelContext.insert(transcription)
try? modelContext.save()
logger.notice("🌅 Prewarm completed in \(String(format: "%.2f", duration))s")
} catch {
logger.error("❌ Prewarm failed: \(error.localizedDescription)")
}
}
// MARK: - Validation
private func shouldPrewarm() -> Bool {
// Check if user has enabled prewarming
let isEnabled = UserDefaults.standard.object(forKey: prewarmEnabledKey) as? Bool ?? true
guard isEnabled else {
logger.notice("🌅 Prewarm disabled by user")
return false
}
// Only prewarm local models (Parakeet and Whisper need ANE compilation)
guard let model = whisperState.currentTranscriptionModel else {
return false
}
switch model.provider {
case .local, .parakeet:
return true
default:
logger.notice("🌅 Skipping prewarm - cloud models don't need it")
return false
}
}
deinit {
NSWorkspace.shared.notificationCenter.removeObserver(self)
logger.notice("🌅 ModelPrewarmService deinitialized")
}
}

View File

@ -6,6 +6,7 @@ struct ModelSettingsView: View {
@AppStorage("IsTextFormattingEnabled") private var isTextFormattingEnabled = true @AppStorage("IsTextFormattingEnabled") private var isTextFormattingEnabled = true
@AppStorage("IsVADEnabled") private var isVADEnabled = true @AppStorage("IsVADEnabled") private var isVADEnabled = true
@AppStorage("AppendTrailingSpace") private var appendTrailingSpace = true @AppStorage("AppendTrailingSpace") private var appendTrailingSpace = true
@AppStorage("PrewarmModelOnWake") private var prewarmModelOnWake = true
@State private var customPrompt: String = "" @State private var customPrompt: String = ""
@State private var isEditing: Bool = false @State private var isEditing: Bool = false
@ -96,13 +97,25 @@ struct ModelSettingsView: View {
Text("Voice Activity Detection (VAD)") Text("Voice Activity Detection (VAD)")
} }
.toggleStyle(.switch) .toggleStyle(.switch)
InfoTip( InfoTip(
title: "Voice Activity Detection", title: "Voice Activity Detection",
message: "Detect speech segments and filter out silence to improve accuracy of local models." message: "Detect speech segments and filter out silence to improve accuracy of local models."
) )
} }
HStack {
Toggle(isOn: $prewarmModelOnWake) {
Text("Prewarm model (Experimental)")
}
.toggleStyle(.switch)
InfoTip(
title: "Prewarm Model (Experimental)",
message: "Turn this on if transcriptions with local models are taking longer than expected. Runs silent background transcription on app launch and wake to trigger optimization."
)
}
} }
.padding() .padding()
.background(Color(NSColor.controlBackgroundColor)) .background(Color(NSColor.controlBackgroundColor))

View File

@ -381,7 +381,7 @@ struct SettingsView: View {
.controlSize(.large) .controlSize(.large)
} }
} }
SettingsSection( SettingsSection(
icon: "lock.shield", icon: "lock.shield",
title: "Data & Privacy", title: "Data & Privacy",

View File

@ -22,12 +22,15 @@ struct VoiceInkApp: App {
@AppStorage("hasCompletedOnboarding") private var hasCompletedOnboarding = false @AppStorage("hasCompletedOnboarding") private var hasCompletedOnboarding = false
@AppStorage("enableAnnouncements") private var enableAnnouncements = true @AppStorage("enableAnnouncements") private var enableAnnouncements = true
@State private var showMenuBarIcon = true @State private var showMenuBarIcon = true
// Audio cleanup manager for automatic deletion of old audio files // Audio cleanup manager for automatic deletion of old audio files
private let audioCleanupManager = AudioCleanupManager.shared private let audioCleanupManager = AudioCleanupManager.shared
// Transcription auto-cleanup service for zero data retention // Transcription auto-cleanup service for zero data retention
private let transcriptionAutoCleanupService = TranscriptionAutoCleanupService.shared private let transcriptionAutoCleanupService = TranscriptionAutoCleanupService.shared
// Model prewarm service for optimizing model on wake from sleep
@StateObject private var prewarmService: ModelPrewarmService
init() { init() {
// Configure FluidAudio logging subsystem // Configure FluidAudio logging subsystem
@ -101,21 +104,26 @@ struct VoiceInkApp: App {
let hotkeyManager = HotkeyManager(whisperState: whisperState) let hotkeyManager = HotkeyManager(whisperState: whisperState)
_hotkeyManager = StateObject(wrappedValue: hotkeyManager) _hotkeyManager = StateObject(wrappedValue: hotkeyManager)
let menuBarManager = MenuBarManager() let menuBarManager = MenuBarManager()
_menuBarManager = StateObject(wrappedValue: menuBarManager) _menuBarManager = StateObject(wrappedValue: menuBarManager)
appDelegate.menuBarManager = menuBarManager
let activeWindowService = ActiveWindowService.shared let activeWindowService = ActiveWindowService.shared
activeWindowService.configure(with: enhancementService) activeWindowService.configure(with: enhancementService)
activeWindowService.configureWhisperState(whisperState) activeWindowService.configureWhisperState(whisperState)
_activeWindowService = StateObject(wrappedValue: activeWindowService) _activeWindowService = StateObject(wrappedValue: activeWindowService)
let prewarmService = ModelPrewarmService(whisperState: whisperState, modelContext: container.mainContext)
_prewarmService = StateObject(wrappedValue: prewarmService)
appDelegate.menuBarManager = menuBarManager
// Ensure no lingering recording state from previous runs // Ensure no lingering recording state from previous runs
Task { Task {
await whisperState.resetOnLaunch() await whisperState.resetOnLaunch()
} }
AppShortcuts.updateAppShortcutParameters() AppShortcuts.updateAppShortcutParameters()
} }