Merge branch 'stakira:master' into master

Author: lottev1991

OpenUtau.Core/Analysis/Game.cs

@@ -49,7 +49,10 @@ public class Game : MidiExtractor<GameOptions> {
     InferenceSession? segmenterSession;
     InferenceSession? estimatorSession;
     InferenceSession? bd2durSession;
+    RunOptions? runOptions;
     bool sessionsLoaded = false;
+    bool disposed = false;
+    volatile bool stopping = false;
     GameConfig config;
     string Location;
 
@@ -97,11 +100,23 @@ public Game(string? location) {
     /// </summary>
     private void EnsureSessionsLoaded() {
         if (sessionsLoaded) return;
+        if (stopping) {
+            throw new OperationCanceledException();
+        }
+        runOptions = new RunOptions();
+        if (stopping) {
+            runOptions.Terminate = true;
+            throw new OperationCanceledException();
+        }
         encoderSession = CreateSession("encoder.onnx", OnnxRunnerChoice.CPUForCoreML);
         segmenterSession = CreateSession("segmenter.onnx", OnnxRunnerChoice.Default);
         estimatorSession = CreateSession("estimator.onnx", OnnxRunnerChoice.Default);
         bd2durSession = CreateSession("bd2dur.onnx", OnnxRunnerChoice.Default);
         sessionsLoaded = true;
+        if (stopping) {
+            runOptions.Terminate = true;
+            throw new OperationCanceledException();
+        }
     }
 
     protected override bool SupportsBatch => true;
@@ -131,59 +146,76 @@ private List<List<TranscribedNote>> RunPipeline(List<float[]> batch, GameOptions
         var waveform = new DenseTensor<float>(waveformData, new[] { B, maxLen });
         var duration = new DenseTensor<float>(durationData, new[] { B });
 
-        // 1. Encoder
-        var (xSeg, xEst, maskT) = RunEncoder(waveform, duration);
+        try {
+            // 1. Encoder
+            var (xSeg, xEst, maskT) = RunEncoder(waveform, duration);
 
-        // 2. Segmentation (D3PM loop)
-        int T = xSeg.Dimensions[1];
-        Tensor<bool> knownBoundaries = new DenseTensor<bool>(new[] { B, T });
-        Tensor<bool> boundaries = new DenseTensor<bool>(new[] { B, T });
+            // 2. Segmentation (D3PM loop)
+            int T = xSeg.Dimensions[1];
+            Tensor<bool> knownBoundaries = new DenseTensor<bool>(new[] { B, T });
+            Tensor<bool> boundaries = new DenseTensor<bool>(new[] { B, T });
 
-        Tensor<long>? language = null;
-        if (config.Languages != null) {
-            int languageId = ResolveLanguageId(options.LanguageCode);
-            language = new DenseTensor<long>(
-                Enumerable.Repeat((long)languageId, B).ToArray(), new[] { B });
-        }
-
-        var segThreshold = new DenseTensor<float>(new[] { options.BoundaryThreshold }, Array.Empty<int>());
-        var radius = new DenseTensor<long>(new long[] { options.BoundaryRadius }, Array.Empty<int>());
+            Tensor<long>? language = null;
+            if (config.Languages != null) {
+                int languageId = ResolveLanguageId(options.LanguageCode);
+                language = new DenseTensor<long>(
+                    Enumerable.Repeat((long)languageId, B).ToArray(), new[] { B });
+            }
 
-        if (config.Loop) {
-            float step = 1.0f / options.SamplingSteps;
-            for (int i = 0; i < options.SamplingSteps; i++) {
-                var t = new DenseTensor<float>(
-                    Enumerable.Repeat(i * step, B).ToArray(), new[] { B });
-                boundaries = RunSegmenter(xSeg, knownBoundaries, boundaries, t, maskT, language, segThreshold, radius);
+            var segThreshold = new DenseTensor<float>(new[] { options.BoundaryThreshold }, Array.Empty<int>());
+            var radius = new DenseTensor<long>(new long[] { options.BoundaryRadius }, Array.Empty<int>());
+
+            if (config.Loop) {
+                float step = 1.0f / options.SamplingSteps;
+                for (int i = 0; i < options.SamplingSteps; i++) {
+                    var t = new DenseTensor<float>(
+                        Enumerable.Repeat(i * step, B).ToArray(), new[] { B });
+                    boundaries = RunSegmenter(xSeg, knownBoundaries, boundaries, t, maskT, language, segThreshold, radius);
+                }
+            } else {
+                boundaries = RunSegmenter(xSeg, knownBoundaries, null, null, maskT, language, segThreshold, radius);
             }
-        } else {
-            boundaries = RunSegmenter(xSeg, knownBoundaries, null, null, maskT, language, segThreshold, radius);
-        }
 
-        // 3. Boundaries to durations
-        var (durations, maskN) = RunBd2Dur(boundaries, maskT);
-        int N = maskN.Dimensions[1];
+            // 3. Boundaries to durations
+            var (durations, maskN) = RunBd2Dur(boundaries, maskT);
+            int N = maskN.Dimensions[1];
 
-        // 4. Estimation
-        var scoreThreshold = new DenseTensor<float>(new[] { options.ScoreThreshold }, Array.Empty<int>());
-        var (presence, scores) = RunEstimator(xEst, boundaries, maskT, maskN, scoreThreshold);
+            // 4. Estimation
+            var scoreThreshold = new DenseTensor<float>(new[] { options.ScoreThreshold }, Array.Empty<int>());
+            var (presence, scores) = RunEstimator(xEst, boundaries, maskT, maskN, scoreThreshold);
 
-        // 5. Split results per batch item
-        var results = new List<List<TranscribedNote>>(B);
-        for (int b = 0; b < B; b++) {
-            var notes = new List<TranscribedNote>(N);
-            for (int i = 0; i < N; i++) {
-                if (!maskN[b, i]) break;
-                notes.Add(new TranscribedNote(durations[b, i], scores[b, i], presence[b, i]));
+            // 5. Split results per batch item
+            var results = new List<List<TranscribedNote>>(B);
+            for (int b = 0; b < B; b++) {
+                var notes = new List<TranscribedNote>(N);
+                for (int i = 0; i < N; i++) {
+                    if (!maskN[b, i]) break;
+                    notes.Add(new TranscribedNote(durations[b, i], scores[b, i], presence[b, i]));
+                }
+
+                results.Add(notes);
             }
 
-            results.Add(notes);
+            return results;
+        } catch (OnnxRuntimeException) {
+            if (runOptions != null && runOptions.Terminate) {
+                throw new OperationCanceledException();
+            }
+            throw;
         }
+    }
 
-        return results;
+    public override void Interrupt() {
+        stopping = true;
+        if (!disposed && runOptions != null) {
+            runOptions.Terminate = true;
+        }
     }
 
     protected override void DisposeManaged() {
+        if (disposed) return;
+        disposed = true;
+        runOptions?.Dispose();
         encoderSession?.Dispose();
         segmenterSession?.Dispose();
         estimatorSession?.Dispose();
@@ -228,7 +260,7 @@ private int ResolveLanguageId(string? languageCode) {
             NamedOnnxValue.CreateFromTensor("duration", duration),
         };
 
-        using var outputs = encoderSession!.Run(inputs);
+        using var outputs = encoderSession!.Run(inputs, encoderSession.OutputNames, runOptions);
 
         var xSeg = outputs.First(o => o.Name == "x_seg").AsTensor<float>().ToDenseTensor();
         var xEst = outputs.First(o => o.Name == "x_est").AsTensor<float>().ToDenseTensor();
@@ -267,7 +299,7 @@ private Tensor<bool> RunSegmenter(
         inputs.Add(NamedOnnxValue.CreateFromTensor("threshold", threshold));
         inputs.Add(NamedOnnxValue.CreateFromTensor("radius", radius));
 
-        using var outputs = segmenterSession!.Run(inputs);
+        using var outputs = segmenterSession!.Run(inputs, segmenterSession.OutputNames, runOptions);
         var boundaries = outputs.First(o => o.Name == "boundaries").AsTensor<bool>().ToDenseTensor();
         return boundaries;
     }
@@ -282,7 +314,7 @@ private Tensor<bool> RunSegmenter(
             NamedOnnxValue.CreateFromTensor("maskT", maskT),
         };
 
-        using var outputs = bd2durSession!.Run(inputs);
+        using var outputs = bd2durSession!.Run(inputs, bd2durSession.OutputNames, runOptions);
         var durations = outputs.First(o => o.Name == "durations").AsTensor<float>().ToDenseTensor();
         var maskN = outputs.First(o => o.Name == "maskN").AsTensor<bool>().ToDenseTensor();
 
@@ -303,7 +335,7 @@ private Tensor<bool> RunSegmenter(
             NamedOnnxValue.CreateFromTensor("threshold", threshold),
         };
 
-        using var outputs = estimatorSession!.Run(inputs);
+        using var outputs = estimatorSession!.Run(inputs, estimatorSession.OutputNames, runOptions);
         var presence = outputs.First(o => o.Name == "presence").AsTensor<bool>().ToDenseTensor();
         var scores = outputs.First(o => o.Name == "scores").AsTensor<float>().ToDenseTensor();
         return (presence, scores);