main
dart 172 lines 5.34 KB
Raw
1 import 'dart:async';
2 import 'dart:io';
3
4 import 'package:dio/dio.dart';
5 import 'package:path_provider/path_provider.dart';
6 import 'package:wakelock_plus/wakelock_plus.dart';
7
8 import '../../../core/logger/app_logger.dart';
9 import 'kitten_tts_model.dart';
10 import 'sherpa_bundle_downloader.dart';
11
12 /// Downloads sherpa-onnx Kitten TTS bundles.
13 class KittenTtsDownloader {
14 final Dio _dio = Dio();
15 final Map<String, CancelToken> _activeDownloads = {};
16
17 Future<Directory> _getEngineDir() async {
18 final supportDir = await getApplicationSupportDirectory();
19 final dir = Directory('${supportDir.path}/tts_models/kitten');
20 if (!await dir.exists()) {
21 await dir.create(recursive: true);
22 }
23 return dir;
24 }
25
26 Future<Directory> _getVariantDir(KittenTtsModelVariant variant) async {
27 final engineDir = await _getEngineDir();
28 return Directory('${engineDir.path}/${variant.bundleDirName}');
29 }
30
31 Future<bool> isVariantDownloaded(KittenTtsModelVariant variant) async {
32 final dir = await _getVariantDir(variant);
33 return await File('${dir.path}/${variant.modelFileName}').exists() &&
34 await File('${dir.path}/${variant.voicesFileName}').exists() &&
35 await File('${dir.path}/${variant.tokensFileName}').exists() &&
36 await Directory('${dir.path}/${variant.dataDirName}').exists();
37 }
38
39 Stream<KittenTtsFileProgress> downloadModel(KittenTtsModel model) {
40 final variant = model.variant;
41 final variantKey = variant.name;
42
43 return Stream<KittenTtsFileProgress>.multi((controller) async {
44 final engineDir = await _getEngineDir();
45 final variantDir = await _getVariantDir(variant);
46 await variantDir.parent.create(recursive: true);
47
48 await WakelockPlus.enable();
49
50 try {
51 final cancelToken = CancelToken();
52 _activeDownloads[variantKey] = cancelToken;
53
54 final tempDir = await getTemporaryDirectory();
55 final archivePath = '${tempDir.path}/${variant.bundleDirName}.tar.bz2';
56 final archiveFile = File(archivePath);
57 final fileName = '${variant.bundleDirName}.tar.bz2';
58
59 controller.add(
60 KittenTtsFileProgress(
61 fileName: fileName,
62 variant: variant,
63 receivedBytes: 0,
64 totalBytes: model.totalSizeBytes,
65 isComplete: false,
66 ),
67 );
68
69 await _dio.download(
70 variant.tarballUrl,
71 archiveFile.path,
72 cancelToken: cancelToken,
73 onReceiveProgress: (received, total) {
74 final totalBytes = total > 0 ? total : model.totalSizeBytes;
75 controller.add(
76 KittenTtsFileProgress(
77 fileName: fileName,
78 variant: variant,
79 receivedBytes: received,
80 totalBytes: totalBytes,
81 isComplete: false,
82 ),
83 );
84 },
85 );
86
87 if (cancelToken.isCancelled) {
88 return;
89 }
90
91 await extractTarBz2(
92 archivePath: archiveFile.path,
93 targetDir: engineDir,
94 );
95
96 if (await archiveFile.exists()) {
97 await archiveFile.delete();
98 }
99
100 controller.add(
101 KittenTtsFileProgress(
102 fileName: fileName,
103 variant: variant,
104 receivedBytes: model.totalSizeBytes,
105 totalBytes: model.totalSizeBytes,
106 isComplete: true,
107 ),
108 );
109 Log.info('Kitten ${variant.displayName} downloaded');
110 } catch (e, st) {
111 if (e is DioException && CancelToken.isCancel(e)) {
112 Log.info('Kitten download cancelled for $variantKey');
113 } else {
114 controller.addError(e, st);
115 }
116 } finally {
117 _activeDownloads.remove(variantKey);
118 if (_activeDownloads.isEmpty) {
119 await WakelockPlus.disable();
120 }
121 await controller.close();
122 }
123 });
124 }
125
126 void cancelDownload(KittenTtsModelVariant variant) {
127 _activeDownloads[variant.name]?.cancel();
128 }
129
130 Future<void> deleteVariant(KittenTtsModelVariant variant) async {
131 final dir = await _getVariantDir(variant);
132 await deleteDirectoryIfExists(dir);
133 Log.info('KittenTTS variant ${variant.displayName} deleted');
134 }
135
136 Future<Set<KittenTtsModelVariant>> getDownloadedVariants() async {
137 final result = <KittenTtsModelVariant>{};
138 for (final variant in KittenTtsModelVariant.values) {
139 if (await isVariantDownloaded(variant)) {
140 result.add(variant);
141 }
142 }
143 return result;
144 }
145
146 Future<String?> getModelPath(KittenTtsModelVariant variant) async {
147 final dir = await _getVariantDir(variant);
148 final file = File('${dir.path}/${variant.modelFileName}');
149 if (await file.exists()) return file.path;
150 return null;
151 }
152
153 Future<String?> getVoicesPath(KittenTtsModelVariant variant) async {
154 final dir = await _getVariantDir(variant);
155 final file = File('${dir.path}/${variant.voicesFileName}');
156 if (await file.exists()) return file.path;
157 return null;
158 }
159
160 Future<String?> getTokensPath(KittenTtsModelVariant variant) async {
161 final dir = await _getVariantDir(variant);
162 final file = File('${dir.path}/${variant.tokensFileName}');
163 if (await file.exists()) return file.path;
164 return null;
165 }
166
167 Future<Directory> getDataDir(KittenTtsModelVariant variant) async {
168 return Directory(
169 '${(await _getVariantDir(variant)).path}/${variant.dataDirName}',
170 );
171 }
172 }