KokoroModelFiles.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

using Microsoft.ML.OnnxRuntime;

namespace Lilith.Core.Kokoro;

internal static class KokoroModelFiles
{
    public static async Task<bool> IsUsableAsync(string path)
    {
        if (!File.Exists(path) || new FileInfo(path).Length == 0) return false;
        try
        {
            return await Task.Run(() =>
            {
                using var session = new InferenceSession(path);
                return session.InputMetadata.ContainsKey("tokens") && session.InputMetadata.ContainsKey("style") &&
                    session.InputMetadata.ContainsKey("speed") && session.OutputMetadata.Count > 0;
            });
        }
        // ONNX invalid protobuf (7) or invalid graph (10) identifies model data damage.
        // Other runtime failures propagate: they must not quarantine a usable model.
        catch (OnnxRuntimeException ex) when (ex.Message.StartsWith("[ErrorCode:InvalidProtobuf]", StringComparison.Ordinal) ||
            ex.Message.StartsWith("[ErrorCode:InvalidGraph]", StringComparison.Ordinal)) { return false; }
    }

    public static void Quarantine(string path)
    {
        if (File.Exists(path)) File.Move(path, path + ".invalid-" + Guid.NewGuid().ToString("N"));
    }
}