diff --git a/TTS/tts/models/vits.py b/TTS/tts/models/vits.py index 334e4526..c24fec68 100644 --- a/TTS/tts/models/vits.py +++ b/TTS/tts/models/vits.py @@ -222,6 +222,9 @@ class VitsArgs(Coqpit): speaker_encoder_config_path: str = "" speaker_encoder_model_path: str = "" fine_tuning_mode: int = 0 + freeze_encoder: bool = False + freeze_DP: bool = False + freeze_PE: bool = False @@ -781,6 +784,20 @@ class Vits(BaseTTS): self.waveform_decoder.train() self.disc.train() + if self.args.freeze_encoder: + for param in self.text_encoder.parameters(): + param.requires_grad = False + for param in self.emb_l.parameters(): + param.requires_grad = False + + if self.args.freeze_PE: + for param in self.posterior_encoder.parameters(): + param.requires_grad = False + + if self.args.freeze_DP: + for param in self.duration_predictor.parameters(): + param.requires_grad = False + if optimizer_idx == 0: text_input = batch["text_input"] text_lengths = batch["text_lengths"]