KokoroSetup.cs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

using System.Text.Json;
using Lilith.Core.Logging;

namespace Lilith.Core.Kokoro;

public static class KokoroSetup
{
    public const string ModelUrl = "https://github.com/Lyrcaxis/KokoroSharpBinaries/releases/download/v2.0.0/kokoro.onnx";

    public static async Task DownloadAsync(KokoroTtsOptions options, Logger logger,
        CancellationToken ct = default, HttpClient? client = null)
        => await DownloadCoreAsync(options, logger, TimeSpan.FromMinutes(30), ct, client);

    internal static async Task DownloadCoreAsync(KokoroTtsOptions options, Logger logger, TimeSpan timeout,
        CancellationToken ct = default, HttpClient? client = null)
    {
        ct.ThrowIfCancellationRequested();
        string path = options.ResolvedModelPath;
        if (await KokoroModelFiles.IsUsableAsync(path)) return;
        KokoroModelFiles.Quarantine(path);
        Directory.CreateDirectory(Path.GetDirectoryName(path)!);
        string temporary = path + "." + Guid.NewGuid().ToString("N") + ".partial";
        using var ownedClient = client is null ? new HttpClient { Timeout = TimeSpan.FromMinutes(30) } : null;
        using var deadline = CancellationTokenSource.CreateLinkedTokenSource(ct);
        deadline.CancelAfter(timeout);
        ct = deadline.Token;
        try
        {
            using var response = await (client ?? ownedClient!).GetAsync(ModelUrl, HttpCompletionOption.ResponseHeadersRead, ct);
            response.EnsureSuccessStatusCode();
            await using var source = await response.Content.ReadAsStreamAsync(ct);
            long total = response.Content.Headers.ContentLength ?? 0;
            long received = 0;
            await using (var destination = new FileStream(temporary, FileMode.CreateNew, FileAccess.Write, FileShare.None, 81920, true))
            {
                var buffer = new byte[81920];
                int count;
                while ((count = await source.ReadAsync(buffer, ct)) > 0)
                {
                    await destination.WriteAsync(buffer.AsMemory(0, count), ct);
                    received += count;
                    if (total > 0) logger.WriteProgress("TTS", received, total);
                    else logger.WriteProgress("TTS", $"Downloaded {received / 1048576.0:F1} MB");
                }
            }
            if (received == 0 || (total > 0 && received != total)) throw new InvalidDataException("Kokoro download is incomplete.");
            ct.ThrowIfCancellationRequested();
            File.Move(temporary, path, overwrite: true);
            logger.WriteLine("TTS", "Kokoro model download complete.");
        }
        finally { if (File.Exists(temporary)) File.Delete(temporary); }
    }

    /// <summary>Stores preferences before downloading, so an interrupted first boot can resume.</summary>
    public static async Task<KokoroTtsOptions> RunAsync(TextReader input, TextWriter output, Logger logger,
        string settingsPath, CancellationToken ct = default)
        => await RunCoreAsync(input, output, logger, settingsPath, KokoroModelFiles.IsUsableAsync, ct);

    internal static async Task<KokoroTtsOptions> RunCoreAsync(TextReader input, TextWriter output, Logger logger,
        string settingsPath, Func<string, Task<bool>> validateModel, CancellationToken ct = default)
    {
        // Automated or redirected sessions keep their existing input protocol and never download implicitly.
        string? enabled = Environment.GetEnvironmentVariable("LILITH_TTS");
        if (enabled is "0" or "false") return new();
        KokoroTtsOptions? options = null;
        if (File.Exists(settingsPath))
        {
            try { options = JsonSettings.Load<KokoroTtsOptions>(settingsPath, ValidOptions); }
            catch (JsonException) { logger.WriteWarning("TTS", "TTS settings are invalid; configure speech again."); }
        }
        if (options is null)
        {
            if (!await input.ConfirmAsync(output, "Enable Kokoro voice replies? [y/N]: ", ct)) options = new();
            else
            {
                await output.WriteAsync("Voice (af_heart, af_bella, am_adam, bf_emma; or custom ID) [af_heart]: ");
                string voice = (await input.ReadLineAsync(ct))?.Trim() ?? "";
                voice = voice.Length == 0 ? "af_heart" : voice;
                KokoroTts.ValidateVoice(voice);
                bool wait = await input.ConfirmAsync(output, "Wait for speech before accepting the next message? [y/N]: ", ct);
                options = new(true, voice, wait);
            }
            ct.ThrowIfCancellationRequested();
            JsonSettings.Save(settingsPath, options);
        }
        if (!options.Enabled) return options;
        KokoroTts.ValidateVoice(options.VoiceId);
        if (!await validateModel(options.ResolvedModelPath))
        {
            if (File.Exists(options.ResolvedModelPath))
            {
                KokoroModelFiles.Quarantine(options.ResolvedModelPath);
                logger.WriteWarning("TTS", "The Kokoro model is damaged or invalid; download it again to repair speech.");
            }
            if (!await input.ConfirmAsync(output, "Download Kokoro model (about 320 MB)? [y/N]: ", ct))
                return options with { Enabled = false };
            await logger.TimeStageAsync("Kokoro model download", () => DownloadAsync(options, logger, ct));
            if (!await validateModel(options.ResolvedModelPath))
            {
                KokoroModelFiles.Quarantine(options.ResolvedModelPath);
                throw new InvalidDataException("The downloaded Kokoro model is invalid. Relaunch setup to retry.");
            }
        }
        return options;
    }

    private static bool ValidOptions(KokoroTtsOptions options)
    {
        try
        {
            _ = options.ResolvedModelPath;
            KokoroTts.ValidateVoice(options.VoiceId);
            return true;
        }
        catch (Exception ex) when (ex is ArgumentException or NotSupportedException) { return false; }
    }
}