diff --git a/tts/tacotron2_v1_early_fakeyou/vocodes_common.py b/tts/tacotron2_v1_early_fakeyou/vocodes_common.py index 26741dc..b990156 100644 --- a/tts/tacotron2_v1_early_fakeyou/vocodes_common.py +++ b/tts/tacotron2_v1_early_fakeyou/vocodes_common.py @@ -6,7 +6,7 @@ import json import numpy as np import torch - +import time from tacotron2.hparams import create_hparams from tacotron2.model import Tacotron2 from tacotron2.layers import TacotronSTFT @@ -221,6 +221,9 @@ class FakeYouTt2Pipeline: def __init__(self): # Cache of "path" => loaded model self.tacotron_model_cache = {} + self.tacotron_model_cache_size_limit = None + self.tacotron_model_cache_usage_tracker = {} + self.tacotron_model_cache_eviction_percent = 0.1 self.waveglow = None self.waveglow_filename = None @@ -233,6 +236,11 @@ class FakeYouTt2Pipeline: self.hifigan_super_resolution_h = None self.hifigan_super_resolution_filename = None + + # method to set cache limit + def set_cache_limit(self, limit): + self.tacotron_model_cache_size_limit = limit + def maybe_load_hifigan(self, checkpoint_path): if not checkpoint_path or checkpoint_path == 'BOGUS_NOT_USED': # Requests no longer have to specify a hifigan to use (if using waveglow) @@ -278,10 +286,26 @@ class FakeYouTt2Pipeline: def maybe_load_tacotron_model_to_cache(self, checkpoint_path): if checkpoint_path in self.tacotron_model_cache: LOGGER.info('Tacotron model already loaded into memory: {}'.format(checkpoint_path)) + # update usage tracker with current timestamp + self.tacotron_model_cache_usage_tracker[checkpoint_path] = time.time() return + + if self.tacotron_model_cache_size_limit and len(self.tacotron_model_cache) >= self.tacotron_model_cache_size_limit: + # uncache at least a few least recently used models + num_to_uncache = int(self.tacotron_model_cache_size_limit * self.tacotron_model_cache_eviction_percent) + LOGGER.info('Cache limit reached, uncaching {} least recently used models'.format(num_to_uncache)) + # find the oldest models + oldest_models = sorted(self.tacotron_model_cache_usage_tracker.items(), key=lambda x: x[1])[:num_to_uncache] + for model_path, _ in oldest_models: + self.uncache_tacotron_model(model_path) + # update usage tracker to remove the uncached model + self.tacotron_model_cache_usage_tracker.pop(model_path, None) + LOGGER.info('Loading tacotron model into memory: {}'.format(checkpoint_path)) model = load_tacotron_model(checkpoint_path) self.tacotron_model_cache[checkpoint_path] = model + # update usage tracker with current timestamp + self.tacotron_model_cache_usage_tracker[checkpoint_path] = time.time() LOGGER.info('Now {} synthesizer models in memory'.format(len(self.tacotron_model_cache))) def uncache_tacotron_model(self, checkpoint_path): diff --git a/tts/tacotron2_v1_early_fakeyou/vocodes_server.py b/tts/tacotron2_v1_early_fakeyou/vocodes_server.py index 08e838b..e0fbde6 100755 --- a/tts/tacotron2_v1_early_fakeyou/vocodes_server.py +++ b/tts/tacotron2_v1_early_fakeyou/vocodes_server.py @@ -287,8 +287,11 @@ class ApiHandler(): def main(): parser = argparse.ArgumentParser() parser.add_argument('--port', type=int, default=8000) + parser.add_argument('--model-cache-limit', type=int, default=100) + args = parser.parse_args() + pipeline.set_cache_limit(args.model_cache_limit) api = falcon.API(middleware=[MultipartMiddleware()]) api.add_route('/', IndexHandler()) api.add_route('/_status', StatusHandler())