487 lines
15 KiB
Dart
487 lines
15 KiB
Dart
import 'dart:io';
|
|
import 'dart:convert';
|
|
import 'dart:math' as math;
|
|
import 'dart:typed_data';
|
|
import 'package:flutter/services.dart' show rootBundle;
|
|
import 'package:flutter_onnxruntime/flutter_onnxruntime.dart';
|
|
import 'package:logger/logger.dart';
|
|
import 'package:path_provider/path_provider.dart';
|
|
|
|
final logger = Logger(
|
|
printer: PrettyPrinter(methodCount: 0, errorMethodCount: 5, lineLength: 80),
|
|
);
|
|
|
|
class UnicodeProcessor {
|
|
final Map<int, int> indexer;
|
|
|
|
UnicodeProcessor._(this.indexer);
|
|
|
|
static Future<UnicodeProcessor> load(String path) async {
|
|
final json = jsonDecode(
|
|
path.startsWith('assets/')
|
|
? await rootBundle.loadString(path)
|
|
: File(path).readAsStringSync(),
|
|
);
|
|
|
|
final indexer = json is List
|
|
? {
|
|
for (var i = 0; i < json.length; i++)
|
|
if (json[i] is int && json[i] >= 0) i: json[i] as int
|
|
}
|
|
: (json as Map<String, dynamic>)
|
|
.map((k, v) => MapEntry(int.parse(k), v as int));
|
|
|
|
return UnicodeProcessor._(indexer);
|
|
}
|
|
|
|
Map<String, dynamic> call(List<String> textList) {
|
|
final lengths = textList.map((t) => t.length).toList();
|
|
final maxLen = lengths.reduce(math.max);
|
|
|
|
final textIds = textList.map((text) {
|
|
final row = List<int>.filled(maxLen, 0);
|
|
final runes = text.runes.toList();
|
|
for (var i = 0; i < runes.length; i++) {
|
|
row[i] = indexer[runes[i]] ?? 0;
|
|
}
|
|
return row;
|
|
}).toList();
|
|
|
|
return {'textIds': textIds, 'textMask': _lengthToMask(lengths)};
|
|
}
|
|
|
|
List<List<List<double>>> _lengthToMask(List<int> lengths, [int? maxLen]) {
|
|
maxLen ??= lengths.reduce(math.max);
|
|
return lengths
|
|
.map((len) => [List.generate(maxLen!, (i) => i < len ? 1.0 : 0.0)])
|
|
.toList();
|
|
}
|
|
}
|
|
|
|
class Style {
|
|
final OrtValue ttl, dp;
|
|
final List<int> ttlShape, dpShape;
|
|
Style(this.ttl, this.dp, this.ttlShape, this.dpShape);
|
|
}
|
|
|
|
class TextToSpeech {
|
|
final Map<String, dynamic> cfgs;
|
|
final UnicodeProcessor textProcessor;
|
|
final OrtSession dpOrt, textEncOrt, vectorEstOrt, vocoderOrt;
|
|
final int sampleRate, baseChunkSize, chunkCompressFactor, ldim;
|
|
|
|
TextToSpeech(this.cfgs, this.textProcessor, this.dpOrt, this.textEncOrt,
|
|
this.vectorEstOrt, this.vocoderOrt)
|
|
: sampleRate = cfgs['ae']['sample_rate'],
|
|
baseChunkSize = cfgs['ae']['base_chunk_size'],
|
|
chunkCompressFactor = cfgs['ttl']['chunk_compress_factor'],
|
|
ldim = cfgs['ttl']['latent_dim'];
|
|
|
|
Future<Map<String, dynamic>> call(String text, Style style, int totalStep,
|
|
{double speed = 1.05, double silenceDuration = 0.3}) async {
|
|
final chunks = _chunkText(text);
|
|
List<double>? wavCat;
|
|
double durCat = 0;
|
|
|
|
for (final chunk in chunks) {
|
|
final result = await _infer([chunk], style, totalStep, speed: speed);
|
|
final wav = _safeCast<double>(result['wav']);
|
|
final duration = _safeCast<double>(result['duration']);
|
|
|
|
if (wavCat == null) {
|
|
wavCat = wav;
|
|
durCat = duration[0];
|
|
} else {
|
|
wavCat = [
|
|
...wavCat,
|
|
...List<double>.filled((silenceDuration * sampleRate).floor(), 0.0),
|
|
...wav
|
|
];
|
|
durCat += duration[0] + silenceDuration;
|
|
}
|
|
}
|
|
|
|
return {
|
|
'wav': wavCat,
|
|
'duration': [durCat]
|
|
};
|
|
}
|
|
|
|
Future<Map<String, dynamic>> _infer(
|
|
List<String> textList, Style style, int totalStep,
|
|
{double speed = 1.05}) async {
|
|
final bsz = textList.length;
|
|
final result = textProcessor.call(textList);
|
|
|
|
final textIdsRaw = result['textIds'];
|
|
final textIds = textIdsRaw is List<List<int>>
|
|
? textIdsRaw
|
|
: (textIdsRaw as List).map((row) => (row as List).cast<int>()).toList();
|
|
|
|
final textMaskRaw = result['textMask'];
|
|
final textMask = textMaskRaw is List<List<List<double>>>
|
|
? textMaskRaw
|
|
: (textMaskRaw as List)
|
|
.map((batch) => (batch as List)
|
|
.map((row) => (row as List).cast<double>())
|
|
.toList())
|
|
.toList();
|
|
|
|
final textIdsShape = [bsz, textIds[0].length];
|
|
final textMaskShape = [bsz, 1, textMask[0][0].length];
|
|
final textMaskTensor = await _toTensor(textMask, textMaskShape);
|
|
|
|
final dpResult = await dpOrt.run({
|
|
'text_ids': await _intToTensor(textIds, textIdsShape),
|
|
'style_dp': style.dp,
|
|
'text_mask': textMaskTensor,
|
|
});
|
|
final durOnnx = _safeCast<double>(await dpResult.values.first.asList());
|
|
final scaledDur = durOnnx.map((d) => d / speed).toList();
|
|
|
|
final textEncResult = await textEncOrt.run({
|
|
'text_ids': await _intToTensor(textIds, textIdsShape),
|
|
'style_ttl': style.ttl,
|
|
'text_mask': textMaskTensor,
|
|
});
|
|
|
|
final latentData = _sampleNoisyLatent(scaledDur);
|
|
final noisyLatentRaw = latentData['noisyLatent'];
|
|
var noisyLatent = noisyLatentRaw is List<List<List<double>>>
|
|
? noisyLatentRaw
|
|
: (noisyLatentRaw as List)
|
|
.map((batch) => (batch as List)
|
|
.map((row) => (row as List).cast<double>())
|
|
.toList())
|
|
.toList();
|
|
|
|
final latentMaskRaw = latentData['latentMask'];
|
|
final latentMask = latentMaskRaw is List<List<List<double>>>
|
|
? latentMaskRaw
|
|
: (latentMaskRaw as List)
|
|
.map((batch) => (batch as List)
|
|
.map((row) => (row as List).cast<double>())
|
|
.toList())
|
|
.toList();
|
|
|
|
final latentShape = [bsz, noisyLatent[0].length, noisyLatent[0][0].length];
|
|
final latentMaskTensor =
|
|
await _toTensor(latentMask, [bsz, 1, latentMask[0][0].length]);
|
|
|
|
final totalStepTensor =
|
|
await _scalarToTensor(List.filled(bsz, totalStep.toDouble()), [bsz]);
|
|
|
|
// Denoising loop
|
|
for (var step = 0; step < totalStep; step++) {
|
|
final result = await vectorEstOrt.run({
|
|
'noisy_latent': await _toTensor(noisyLatent, latentShape),
|
|
'text_emb': textEncResult.values.first,
|
|
'style_ttl': style.ttl,
|
|
'text_mask': textMaskTensor,
|
|
'latent_mask': latentMaskTensor,
|
|
'total_step': totalStepTensor,
|
|
'current_step':
|
|
await _scalarToTensor(List.filled(bsz, step.toDouble()), [bsz]),
|
|
});
|
|
|
|
final denoisedRaw = await result.values.first.asList();
|
|
final denoised = denoisedRaw is List<double>
|
|
? denoisedRaw
|
|
: _safeCast<double>(denoisedRaw);
|
|
var idx = 0;
|
|
for (var b = 0; b < noisyLatent.length; b++) {
|
|
for (var d = 0; d < noisyLatent[b].length; d++) {
|
|
for (var t = 0; t < noisyLatent[b][d].length; t++) {
|
|
noisyLatent[b][d][t] = denoised[idx++];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
final vocoderResult = await vocoderOrt
|
|
.run({'latent': await _toTensor(noisyLatent, latentShape)});
|
|
final wavRaw = await vocoderResult.values.first.asList();
|
|
final wav = wavRaw is List<double> ? wavRaw : _safeCast<double>(wavRaw);
|
|
|
|
return {'wav': wav, 'duration': scaledDur};
|
|
}
|
|
|
|
Map<String, dynamic> _sampleNoisyLatent(List<double> duration) {
|
|
final wavLenMax = duration.reduce(math.max) * sampleRate;
|
|
final wavLengths = duration.map((d) => (d * sampleRate).floor()).toList();
|
|
final chunkSize = baseChunkSize * chunkCompressFactor;
|
|
final latentLen = ((wavLenMax + chunkSize - 1) / chunkSize).floor();
|
|
final latentDim = ldim * chunkCompressFactor;
|
|
|
|
final random = math.Random();
|
|
final noisyLatent = List.generate(
|
|
duration.length,
|
|
(_) => List.generate(
|
|
latentDim,
|
|
(_) => List.generate(latentLen, (_) {
|
|
final u1 = math.max(1e-10, random.nextDouble());
|
|
final u2 = random.nextDouble();
|
|
return math.sqrt(-2.0 * math.log(u1)) * math.cos(2.0 * math.pi * u2);
|
|
}),
|
|
),
|
|
);
|
|
|
|
final latentMask = _getLatentMask(wavLengths);
|
|
|
|
for (var b = 0; b < noisyLatent.length; b++) {
|
|
for (var d = 0; d < noisyLatent[b].length; d++) {
|
|
for (var t = 0; t < noisyLatent[b][d].length; t++) {
|
|
noisyLatent[b][d][t] *= latentMask[b][0][t];
|
|
}
|
|
}
|
|
}
|
|
|
|
return {'noisyLatent': noisyLatent, 'latentMask': latentMask};
|
|
}
|
|
|
|
List<List<List<double>>> _getLatentMask(List<int> wavLengths) {
|
|
final latentSize = baseChunkSize * chunkCompressFactor;
|
|
final latentLengths = wavLengths
|
|
.map((len) => ((len + latentSize - 1) / latentSize).floor())
|
|
.toList();
|
|
final maxLen = latentLengths.reduce(math.max);
|
|
return latentLengths
|
|
.map((len) => [List.generate(maxLen, (i) => i < len ? 1.0 : 0.0)])
|
|
.toList();
|
|
}
|
|
|
|
List<String> _chunkText(String text, {int maxLen = 300}) {
|
|
final paragraphs = text
|
|
.trim()
|
|
.split(RegExp(r'\n\s*\n+'))
|
|
.where((p) => p.trim().isNotEmpty)
|
|
.toList();
|
|
|
|
final chunks = <String>[];
|
|
for (var paragraph in paragraphs) {
|
|
paragraph = paragraph.trim();
|
|
if (paragraph.isEmpty) continue;
|
|
|
|
final sentences = paragraph.split(RegExp(
|
|
r'(?<!Mr\.|Mrs\.|Ms\.|Dr\.|Prof\.)(?<!\b[A-Z]\.)(?<=[.!?])\s+'));
|
|
|
|
var currentChunk = '';
|
|
for (final sentence in sentences) {
|
|
if (currentChunk.length + sentence.length + 1 <= maxLen) {
|
|
currentChunk += (currentChunk.isNotEmpty ? ' ' : '') + sentence;
|
|
} else {
|
|
if (currentChunk.isNotEmpty) chunks.add(currentChunk.trim());
|
|
currentChunk = sentence;
|
|
}
|
|
}
|
|
if (currentChunk.isNotEmpty) chunks.add(currentChunk.trim());
|
|
}
|
|
|
|
return chunks;
|
|
}
|
|
|
|
List<T> _safeCast<T>(dynamic raw) {
|
|
if (raw is List<T>) return raw;
|
|
if (raw is List) {
|
|
if (raw.isNotEmpty && raw.first is List) {
|
|
return _flattenList<T>(raw);
|
|
}
|
|
if (T == double) {
|
|
return raw
|
|
.map((e) => e is num ? e.toDouble() : double.parse(e.toString()))
|
|
.toList() as List<T>;
|
|
}
|
|
return raw.cast<T>();
|
|
}
|
|
throw Exception('Cannot convert $raw to List<$T>');
|
|
}
|
|
|
|
List<T> _flattenList<T>(dynamic list) {
|
|
if (list is List) {
|
|
return list.expand((e) => _flattenList<T>(e)).toList();
|
|
}
|
|
if (T == double && list is num) {
|
|
return [list.toDouble()] as List<T>;
|
|
}
|
|
return [list as T];
|
|
}
|
|
|
|
Future<OrtValue> _toTensor(dynamic array, List<int> dims) async {
|
|
final flat = _flattenList<double>(array);
|
|
return await OrtValue.fromList(Float32List.fromList(flat), dims);
|
|
}
|
|
|
|
Future<OrtValue> _scalarToTensor(List<double> array, List<int> dims) async {
|
|
return await OrtValue.fromList(Float32List.fromList(array), dims);
|
|
}
|
|
|
|
Future<OrtValue> _intToTensor(List<List<int>> array, List<int> dims) async {
|
|
final flat = array.expand((row) => row).toList();
|
|
return await OrtValue.fromList(Int64List.fromList(flat), dims);
|
|
}
|
|
}
|
|
|
|
Future<TextToSpeech> loadTextToSpeech(String onnxDir,
|
|
{bool useGpu = false}) async {
|
|
if (useGpu) throw Exception('GPU mode not supported yet');
|
|
|
|
logger.i('Loading TTS models from $onnxDir');
|
|
|
|
final cfgs = await _loadCfgs(onnxDir);
|
|
final sessions = await _loadOnnxAll(onnxDir);
|
|
final textProcessor =
|
|
await UnicodeProcessor.load('$onnxDir/unicode_indexer.json');
|
|
|
|
logger.i('TTS models loaded successfully');
|
|
|
|
return TextToSpeech(
|
|
cfgs,
|
|
textProcessor,
|
|
sessions['dpOrt']!,
|
|
sessions['textEncOrt']!,
|
|
sessions['vectorEstOrt']!,
|
|
sessions['vocoderOrt']!,
|
|
);
|
|
}
|
|
|
|
Future<Style> loadVoiceStyle(List<String> paths) async {
|
|
final bsz = paths.length;
|
|
|
|
final firstJson = jsonDecode(
|
|
paths[0].startsWith('assets/')
|
|
? await rootBundle.loadString(paths[0])
|
|
: File(paths[0]).readAsStringSync(),
|
|
);
|
|
|
|
final ttlDims = List<int>.from(firstJson['style_ttl']['dims']);
|
|
final dpDims = List<int>.from(firstJson['style_dp']['dims']);
|
|
|
|
final ttlFlat = Float32List(bsz * ttlDims[1] * ttlDims[2]);
|
|
final dpFlat = Float32List(bsz * dpDims[1] * dpDims[2]);
|
|
|
|
for (var i = 0; i < bsz; i++) {
|
|
final json = jsonDecode(
|
|
paths[i].startsWith('assets/')
|
|
? await rootBundle.loadString(paths[i])
|
|
: File(paths[i]).readAsStringSync(),
|
|
);
|
|
|
|
final ttlData = _flattenToDouble(json['style_ttl']['data']);
|
|
final dpData = _flattenToDouble(json['style_dp']['data']);
|
|
|
|
ttlFlat.setRange(i * ttlDims[1] * ttlDims[2],
|
|
(i + 1) * ttlDims[1] * ttlDims[2], ttlData);
|
|
dpFlat.setRange(
|
|
i * dpDims[1] * dpDims[2], (i + 1) * dpDims[1] * dpDims[2], dpData);
|
|
}
|
|
|
|
final ttlShape = [bsz, ttlDims[1], ttlDims[2]];
|
|
final dpShape = [bsz, dpDims[1], dpDims[2]];
|
|
|
|
return Style(
|
|
await OrtValue.fromList(ttlFlat, ttlShape),
|
|
await OrtValue.fromList(dpFlat, dpShape),
|
|
ttlShape,
|
|
dpShape,
|
|
);
|
|
}
|
|
|
|
Future<Map<String, dynamic>> _loadCfgs(String onnxDir) async {
|
|
final path = '$onnxDir/tts.json';
|
|
final json = jsonDecode(await rootBundle.loadString(path));
|
|
return json as Map<String, dynamic>;
|
|
}
|
|
|
|
Future<String> copyModelToFile(String path) async {
|
|
final byteData = await rootBundle.load(path);
|
|
final tempDir = await getApplicationCacheDirectory();
|
|
final modelPath = '${tempDir.path}/${path.split("/").last}';
|
|
|
|
final file = File(modelPath);
|
|
await file.writeAsBytes(byteData.buffer.asUint8List());
|
|
return modelPath;
|
|
}
|
|
|
|
Future<Map<String, OrtSession>> _loadOnnxAll(String dir) async {
|
|
final ort = OnnxRuntime();
|
|
final models = [
|
|
'duration_predictor',
|
|
'text_encoder',
|
|
'vector_estimator',
|
|
'vocoder'
|
|
];
|
|
|
|
final sessions = await Future.wait(models.map((name) async {
|
|
final path = await copyModelToFile('$dir/$name.onnx');
|
|
logger.d('Loading $name.onnx');
|
|
return ort.createSessionFromAsset(path);
|
|
}));
|
|
|
|
return {
|
|
'dpOrt': sessions[0],
|
|
'textEncOrt': sessions[1],
|
|
'vectorEstOrt': sessions[2],
|
|
'vocoderOrt': sessions[3],
|
|
};
|
|
}
|
|
|
|
List<double> _flattenToDouble(dynamic list) {
|
|
if (list is List) return list.expand((e) => _flattenToDouble(e)).toList();
|
|
return [list is num ? list.toDouble() : double.parse(list.toString())];
|
|
}
|
|
|
|
void writeWavFile(String filename, List<double> audioData, int sampleRate) {
|
|
const numChannels = 1;
|
|
const bitsPerSample = 16;
|
|
final dataSize = audioData.length * 2;
|
|
|
|
final buffer = ByteData(44 + dataSize);
|
|
var offset = 0;
|
|
|
|
// RIFF header
|
|
for (var byte in [0x52, 0x49, 0x46, 0x46]) {
|
|
buffer.setUint8(offset++, byte);
|
|
}
|
|
buffer.setUint32(offset, 36 + dataSize, Endian.little);
|
|
offset += 4;
|
|
|
|
// WAVE
|
|
for (var byte in [0x57, 0x41, 0x56, 0x45]) {
|
|
buffer.setUint8(offset++, byte);
|
|
}
|
|
|
|
// fmt chunk
|
|
for (var byte in [0x66, 0x6D, 0x74, 0x20]) {
|
|
buffer.setUint8(offset++, byte);
|
|
}
|
|
buffer.setUint32(offset, 16, Endian.little);
|
|
offset += 4;
|
|
buffer.setUint16(offset, 1, Endian.little);
|
|
offset += 2;
|
|
buffer.setUint16(offset, numChannels, Endian.little);
|
|
offset += 2;
|
|
buffer.setUint32(offset, sampleRate, Endian.little);
|
|
offset += 4;
|
|
buffer.setUint32(offset, sampleRate * numChannels * 2, Endian.little);
|
|
offset += 4;
|
|
buffer.setUint16(offset, numChannels * 2, Endian.little);
|
|
offset += 2;
|
|
buffer.setUint16(offset, bitsPerSample, Endian.little);
|
|
offset += 2;
|
|
|
|
// data chunk
|
|
for (var byte in [0x64, 0x61, 0x74, 0x61]) {
|
|
buffer.setUint8(offset++, byte);
|
|
}
|
|
buffer.setUint32(offset, dataSize, Endian.little);
|
|
offset += 4;
|
|
|
|
// Write audio samples
|
|
for (var i = 0; i < audioData.length; i++) {
|
|
final sample = (audioData[i].clamp(-1.0, 1.0) * 32767).round();
|
|
buffer.setInt16(offset + i * 2, sample, Endian.little);
|
|
}
|
|
|
|
File(filename).writeAsBytesSync(buffer.buffer.asUint8List());
|
|
}
|