Merge pull request #983 from oxygen-dioxide/free-memory

DiffSinger: Free memory for singers no longer in use

Author: stakira

OpenUtau.Core/DiffSinger/DiffSingerPitch.cs

@@ -14,7 +14,7 @@
 
 namespace OpenUtau.Core.DiffSinger
 {
-    public class DsPitch
+    public class DsPitch : IDisposable
     {
         string rootPath;
         DsConfig dsConfig;
@@ -73,6 +73,7 @@ void SetRange<T>(T[] list, T value, int startIndex, int endIndex){
                 list[i] = value;
             }
         }
+        
         public RenderPitchResult Process(RenderPhrase phrase){
             var startMs = Math.Min(phrase.notes[0].positionMs, phrase.phones[0].positionMs) - headMs;
             var endMs = phrase.notes[^1].endMs + tailMs;
@@ -280,5 +281,22 @@ public RenderPitchResult Process(RenderPhrase phrase){
                 };
             }
         }
+
+        private bool disposedValue;
+        
+        protected virtual void Dispose(bool disposing) {
+            if (!disposedValue) {
+                if (disposing) {
+                    linguisticModel?.Dispose();
+                    pitchModel?.Dispose();
+                }
+                disposedValue = true;
+            }
+        }
+
+        public void Dispose() {
+            Dispose(disposing: true);
+            GC.SuppressFinalize(this);
+        }
     }
 }

OpenUtau.Core/DiffSinger/DiffSingerRenderer.cs

@@ -101,10 +101,12 @@ public Task<RenderResult> Render(RenderPhrase phrase, Progress progress, int tra
                         }
                     }
                     if (result.samples == null) {
-                        result.samples = InvokeDiffsinger(phrase, depth, speedup);
-                        var source = new WaveSource(0, 0, 0, 1);
-                        source.SetSamples(result.samples);
-                        WaveFileWriter.CreateWaveFile16(wavPath, new ExportAdapter(source).ToMono(1, 0));
+                        result.samples = InvokeDiffsinger(phrase, depth, speedup, cancellation);
+                        if (result.samples != null) {
+                            var source = new WaveSource(0, 0, 0, 1);
+                            source.SetSamples(result.samples);
+                            WaveFileWriter.CreateWaveFile16(wavPath, new ExportAdapter(source).ToMono(1, 0));
+                        }
                     }
                     if (result.samples != null) {
                         Renderers.ApplyDynamics(phrase, result);
@@ -120,7 +122,7 @@ public Task<RenderResult> Render(RenderPhrase phrase, Progress progress, int tra
         leadingMs、positionMs、estimatedLengthMs: timeaxis layout in Ms, double
          */
 
-        float[] InvokeDiffsinger(RenderPhrase phrase, int depth, int speedup) {
+        float[] InvokeDiffsinger(RenderPhrase phrase, int depth, int speedup, CancellationTokenSource cancellation) {
             var singer = phrase.singer as DiffSingerSinger;
             //Check if dsconfig.yaml is correct
             if(String.IsNullOrEmpty(singer.dsConfig.vocoder) ||
@@ -130,6 +132,7 @@ float[] InvokeDiffsinger(RenderPhrase phrase, int depth, int speedup) {
             }
 
             var vocoder = singer.getVocoder();
+            var acousticModel = singer.getAcousticSession();
             var frameMs = vocoder.frameMs();
             var frameSec = frameMs / 1000;
             int headFrames = (int)Math.Round(headMs / frameMs);
@@ -218,8 +221,16 @@ float[] InvokeDiffsinger(RenderPhrase phrase, int depth, int speedup) {
             }
 
             //Variance: Energy and Breathiness
+            
             if(singer.dsConfig.useBreathinessEmbed || singer.dsConfig.useEnergyEmbed){
-                var varianceResult = singer.getVariancePredictor().Process(phrase);
+                var variancePredictor = singer.getVariancePredictor();
+                VarianceResult varianceResult;
+                lock(variancePredictor){
+                    if(cancellation.IsCancellationRequested) {
+                        return null;
+                    }
+                    varianceResult = singer.getVariancePredictor().Process(phrase);
+                }
                 //TODO: let user edit variance curves
                 if(singer.dsConfig.useEnergyEmbed){
                     var energyCurve = phrase.curves.FirstOrDefault(curve => curve.Item1 == ENE);
@@ -246,26 +257,36 @@ float[] InvokeDiffsinger(RenderPhrase phrase, int depth, int speedup) {
                         .Reshape(new int[] { 1, breathiness.Length })));
                 }
             }
-
-            var acousticModel = singer.getAcousticSession();
-            Onnx.VerifyInputNames(acousticModel, acousticInputs);
             Tensor<float> mel;
-            var acousticOutputs = acousticModel.Run(acousticInputs);
-            mel = acousticOutputs.First().AsTensor<float>().Clone();
-            
+            lock(acousticModel){
+                if(cancellation.IsCancellationRequested) {
+                    return null;
+                }
+                Onnx.VerifyInputNames(acousticModel, acousticInputs);
+                var acousticOutputs = acousticModel.Run(acousticInputs);
+                mel = acousticOutputs.First().AsTensor<float>().Clone();
+            }
             //vocoder
             //waveform = session.run(['waveform'], {'mel': mel, 'f0': f0})[0]
             var vocoderInputs = new List<NamedOnnxValue>();
             vocoderInputs.Add(NamedOnnxValue.CreateFromTensor("mel", mel));
             vocoderInputs.Add(NamedOnnxValue.CreateFromTensor("f0",f0tensor));
             float[] samples;
-            var vocoderOutputs = vocoder.session.Run(vocoderInputs);
-            samples = vocoderOutputs.First().AsTensor<float>().ToArray();
+            lock(vocoder){
+                if(cancellation.IsCancellationRequested) {
+                    return null;
+                }
+                var vocoderOutputs = vocoder.session.Run(vocoderInputs);
+                samples = vocoderOutputs.First().AsTensor<float>().ToArray();
+            }
             return samples;
         }
 
         public RenderPitchResult LoadRenderedPitch(RenderPhrase phrase) {
-            return (phrase.singer as DiffSingerSinger).getPitchPredictor().Process(phrase);
+            var pitchPredictor = (phrase.singer as DiffSingerSinger).getPitchPredictor();
+            lock(pitchPredictor){
+                return pitchPredictor.Process(phrase);
+            }
         }
 
         public UExpressionDescriptor[] GetSuggestedExpressions(USinger singer, URenderSettings renderSettings) {

OpenUtau.Core/DiffSinger/DiffSingerSinger.cs

@@ -168,5 +168,33 @@ public DsVariance getVariancePredictor(){
             }
             return variancePredictor;
         }
+
+        public override void FreeMemory(){
+            Log.Information($"Freeing memory for singer {Id}");
+            if(acousticSession != null) {
+                lock(acousticSession) {
+                    acousticSession?.Dispose();
+                }
+                acousticSession = null;
+            }
+            if(vocoder != null) {
+                lock(vocoder) {
+                    vocoder?.Dispose();
+                }
+                vocoder = null;
+            }
+            if(pitchPredictor != null) {
+                lock(pitchPredictor) {
+                    pitchPredictor?.Dispose();
+                }
+                pitchPredictor = null;
+            }
+            if(variancePredictor != null){
+                lock(variancePredictor) {
+                    variancePredictor?.Dispose();
+                }
+                variancePredictor = null;
+            }
+        }
     }
 }

OpenUtau.Core/DiffSinger/DiffSingerVariance.cs

@@ -17,7 +17,7 @@ public struct VarianceResult{
         public float[] energy;
         public float[] breathiness;
     }
-    public class DsVariance{
+    public class DsVariance : IDisposable{
         string rootPath;
         DsConfig dsConfig;
         List<string> phonemes;
@@ -172,5 +172,22 @@ public VarianceResult Process(RenderPhrase phrase){
                 breathiness = breathiness_pred.ToArray()
             };
         }
+
+        private bool disposedValue;
+
+        protected virtual void Dispose(bool disposing) {
+            if (!disposedValue) {
+                if (disposing) {
+                    linguisticModel?.Dispose();
+                    varianceModel?.Dispose();
+                }
+                disposedValue = true;
+            }
+        }
+
+        public void Dispose() {
+            Dispose(disposing: true);
+            GC.SuppressFinalize(this);
+        }
     }
 }

OpenUtau.Core/DiffSinger/DiffSingerVocoder.cs

@@ -3,7 +3,7 @@
 using Microsoft.ML.OnnxRuntime;
 
 namespace OpenUtau.Core.DiffSinger {
-    public class DsVocoder {
+    public class DsVocoder : IDisposable {
         public string Location;
         public DsVocoderConfig config;
         public InferenceSession session;
@@ -27,6 +27,23 @@ public DsVocoder(string name) {
         public float frameMs() {
             return 1000f * config.hop_size / config.sample_rate;
         }
+
+        private bool disposedValue;
+
+        protected virtual void Dispose(bool disposing) {
+            if (!disposedValue) {
+                if (disposing) {
+                    session?.Dispose();
+                }
+                disposedValue = true;
+            }
+        }
+
+        public void Dispose() {
+            Dispose(disposing: true);
+            GC.SuppressFinalize(this);
+        }
+
     }
 
     [Serializable]

OpenUtau.Core/Render/RenderEngine.cs

@@ -169,6 +169,7 @@ public void PreRenderProject(ref CancellationTokenSource cancellation) {
 
         private RenderPartRequest[] PrepareRequests() {
             RenderPartRequest[] requests;
+            SingerManager.Inst.ReleaseSingersNotInUse(project);
             lock (project) {
                 requests = project.parts
                     .Where(part => part is UVoicePart)

OpenUtau.Core/SingerManager.cs

@@ -19,6 +19,8 @@ public class SingerManager : SingletonBase<SingerManager> {
 
         private readonly ConcurrentQueue<USinger> reloadQueue = new ConcurrentQueue<USinger>();
         private CancellationTokenSource reloadCancellation;
+        
+        private HashSet<USinger> singersUsed = new HashSet<USinger>();
 
         public void Initialize() {
             InitializationTask = Task.Run(() => {
@@ -109,5 +111,26 @@ private void Refresh() {
                 }).Start(DocManager.Inst.MainScheduler);
             }
         }
+
+        //Check which singers are in use and free memory for those that are not
+        public void ReleaseSingersNotInUse(UProject project) {
+            //Check which singers are in use
+            var singersInUse = new HashSet<USinger>();
+            foreach(var track in project.tracks){
+                var singer = track.Singer;
+                if(singer != null){
+                    singersInUse.Add(singer);
+                }
+            }
+            //Release singers that are no longer in use
+            foreach(var singer in singersUsed){
+                if(!singersInUse.Contains(singer)){
+                    singer.FreeMemory();
+                    singersUsed.Remove(singer);
+                }
+            }
+            //Update singers used
+            singersUsed.UnionWith(singersInUse);
+        }
     }
 }

OpenUtau.Core/Ustx/USinger.cs

@@ -280,5 +280,14 @@ public static USinger CreateMissing(string name) {
         private void NotifyPropertyChanged(string propertyName = "") {
             PropertyChanged?.Invoke(this, new PropertyChangedEventArgs(propertyName));
         }
+
+        /// <summary>
+        /// Some types of singers store their data in memory when rendering.
+        /// This method is called when the singer is no longer used.
+        /// Note:
+        /// - the voicebank may be used again even after this method is called.
+        /// - this method may be called even when the singer has not been used
+        /// </summary>
+        public virtual void FreeMemory(){ }
     }
 }