model push

VitsModelSplit/.gitignore

@@ -0,0 +1,163 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+models--facebook--mms-tts-ara/
+output/
+dataset/
+wandb/
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

VitsModelSplit/Arguments.py

@@ -0,0 +1,265 @@
+
+from transformers import TrainingArguments
+from typing import Any, Optional
+from dataclasses import dataclass, field
+
+
+#.............................................
+
+#### ARGUMENTS
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    feature_extractor_name: Optional[str] = field(
+        default=None, metadata={"help": "feature extractor name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    token: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
+                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
+            )
+        },
+    )
+    use_auth_token: bool = field(
+        default=None,
+        metadata={
+            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token`."
+        },
+    )
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option"
+                "should only be set to `True` for repositories you trust and in which you have read the code, as it will"
+                "execute code present on the Hub on your local machine."
+            )
+        },
+    )
+    override_speaker_embeddings: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "If `True` and if `speaker_id_column_name` is specified, it will replace current speaker embeddings with a new set of speaker embeddings."
+                "If the model from the checkpoint didn't have speaker embeddings, it will initialize speaker embeddings."
+            )
+        },
+    )
+
+    override_vocabulary_embeddings: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "If `True`, it will resize the token embeddings based on the vocabulary size of the tokenizer. In other words, use this when you use a different tokenizer than the one that was used during pretraining."
+            )
+        },
+    )
+
+#.............................................................................................
+
+
+@dataclass
+class VITSTrainingArguments(TrainingArguments):
+    do_step_schedule_per_epoch: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether or not to perform scheduler steps per epoch or per steps. If `True`, the scheduler will be `ExponentialLR` parametrized with `lr_decay`."
+            )
+        },
+    )
+
+    lr_decay: float = field(
+        default=0.999875,
+        metadata={"help": "Learning rate decay, used with `ExponentialLR` when `do_step_schedule_per_epoch`."},
+    )
+
+    weight_duration: float = field(default=1.0, metadata={"help": "Duration loss weight."})
+
+    weight_kl: float = field(default=1.5, metadata={"help": "KL loss weight."})
+
+    weight_mel: float = field(default=35.0, metadata={"help": "Mel-spectrogram loss weight"})
+
+    weight_disc: float = field(default=3.0, metadata={"help": "Discriminator loss weight"})
+
+    weight_gen: float = field(default=1.0, metadata={"help": "Generator loss weight"})
+
+    weight_fmaps: float = field(default=1.0, metadata={"help": "Feature map loss weight"})
+    d_learning_rate: float = field(default=2e-4, metadata={"help": "Feature map loss weight"})
+   
+    d_adam_beta1: float = field(default=0.8, metadata={"help": "Feature map loss weight"})
+    d_adam_beta2: float = field(default=0.99, metadata={"help": "Feature map loss weight"})
+
+
+#.............................................................................................
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    project_name: str = field(
+        default="vits_finetuning",
+        metadata={"help": "The project name associated to this run. Useful to track your experiment."},
+    )
+    dataset_name: str = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+                "value if set."
+            )
+        },
+    )
+    audio_column_name: str = field(
+        default="audio",
+        metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},
+    )
+    text_column_name: str = field(
+        default="text",
+        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},
+    )
+    speaker_id_column_name: str = field(
+        default=None,
+        metadata={
+            "help": """If set, corresponds to the name of the speaker id column containing the speaker ids.
+                If `override_speaker_embeddings=False`:
+                    it assumes that speakers are indexed from 0 to `num_speakers-1`.
+                    `num_speakers` and `speaker_embedding_size` have to be set in the model config.
+
+                If `override_speaker_embeddings=True`:
+                        It will use this column to compute how many speakers there are.
+
+                Defaults to None, i.e it is not used by default."""
+        },
+    )
+    filter_on_speaker_id: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If `speaker_id_column_name` and `filter_on_speaker_id` are set, will filter the dataset to keep a single speaker_id (`filter_on_speaker_id`)  "
+            )
+        },
+    )
+
+    max_tokens_length: float = field(
+        default=450,
+        metadata={
+            "help": ("Truncate audio files with a transcription that are longer than `max_tokens_length` tokens")
+        },
+    )
+    max_duration_in_seconds: float = field(
+        default=20.0,
+        metadata={
+            "help": (
+                "Truncate audio files that are longer than `max_duration_in_seconds` seconds to"
+                " 'max_duration_in_seconds`"
+            )
+        },
+    )
+    min_duration_in_seconds: float = field(
+        default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
+    )
+    preprocessing_only: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to only do data preprocessing and skip training. This is especially useful when data"
+                " preprocessing errors out in distributed training due to timeout. In this case, one should run the"
+                " preprocessing in a non-distributed setup with `preprocessing_only=True` so that the cached datasets"
+                " can consequently be loaded in distributed training"
+            )
+        },
+    )
+    train_split_name: str = field(
+        default="train",
+        metadata={
+            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
+        },
+    )
+    eval_split_name: str = field(
+        default="test",
+        metadata={
+            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
+        },
+    )
+    do_lower_case: bool = field(
+        default=False,
+        metadata={"help": "Whether the input text should be lower cased."},
+    )
+    do_normalize: bool = field(
+        default=False,
+        metadata={"help": "Whether the input waveform should be normalized."},
+    )
+    full_generation_sample_text: str = field(
+        default="This is a test, let's see what comes out of this.",
+        metadata={
+            "help": (
+                "Language for multilingual fine-tuning. This argument should be set for multilingual fine-tuning "
+                "only. For English speech recognition, it should be set to `None`."
+            )
+        },
+    )
+    uroman_path: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "Absolute path to the uroman package. To use if your model requires `uroman`."
+                "An easy way to check it is to go on your model card and manually check `is_uroman` in the `tokenizer_config.json,"
+                "e.g the French checkpoint doesn't need it: https://huggingface.co/facebook/mms-tts-fra/blob/main/tokenizer_config.json#L4"
+            )
+        },
+    )
+
+#.............................................................................................

VitsModelSplit/DATA

@@ -0,0 +1 @@
+

VitsModelSplit/PosteriorDecoderModel.py

@@ -0,0 +1,331 @@
+import os
+import sys
+from typing import Optional
+import numpy as np
+import torch
+from torch import nn
+from transformers import set_seed
+import wandb
+import logging
+import copy
+
+from .vits_config import VitsConfig, VitsPreTrainedModel
+from .feature_extraction import VitsFeatureExtractor
+from .vits_output import PosteriorDecoderModelOutput
+from  .dataset_features_collector import FeaturesCollectionDataset
+from .posterior_encoder import VitsPosteriorEncoder
+from .decoder import VitsHifiGan
+
+class PosteriorDecoderModel(torch.nn.Module):
+    
+    def __init__(self, config,posterior_encoder,decoder,device=None):
+        super().__init__()
+        
+        if device:
+            self.device = device
+        else:
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            
+        self.config = copy.deepcopy(config)
+        self.posterior_encoder = copy.deepcopy(posterior_encoder)
+        self.decoder = copy.deepcopy(decoder)
+        
+        if config.num_speakers > 1:
+            self.embed_speaker = nn.Embedding(config.num_speakers, 
+                                              config.speaker_embedding_size
+                                              )
+        self.sampling_rate = config.sampling_rate
+        self.speaking_rate = config.speaking_rate
+        self.noise_scale = config.noise_scale
+        self.noise_scale_duration = config.noise_scale_duration
+        self.segment_size = self.config.segment_size // self.config.hop_length
+        
+        self.to(self.device)
+        
+    
+    
+    #....................................
+    
+    def slice_segments(self,hidden_states, ids_str, segment_size=4):
+        
+        batch_size, channels, _ = hidden_states.shape
+        # 1d tensor containing the indices to keep
+        indices = torch.arange(segment_size).to(ids_str.device)
+        # extend the indices to match the shape of hidden_states
+        indices = indices.view(1, 1, -1).expand(batch_size, channels, -1)
+        # offset indices with ids_str
+        indices = indices + ids_str.view(-1, 1, 1)
+        # gather indices
+        output = torch.gather(hidden_states, dim=2, index=indices)
+
+        return output
+    
+    #....................................
+    
+    def rand_slice_segments(self,hidden_states, sample_lengths=None, segment_size=4):
+        batch_size, _, seq_len = hidden_states.size()
+        if sample_lengths is None:
+            sample_lengths = seq_len
+        ids_str_max = sample_lengths - segment_size + 1
+        ids_str = (torch.rand([batch_size]).to(device=hidden_states.device) * ids_str_max).to(dtype=torch.long)
+        ret = self.slice_segments(hidden_states, ids_str, segment_size)
+        
+        return ret, ids_str
+
+    #....................................
+    
+    def forward(
+        self,
+        labels: Optional[torch.FloatTensor] = None,
+        labels_attention_mask: Optional[torch.Tensor] = None,
+        speaker_id: Optional[int] = None,
+        return_dict: Optional[bool] = True,
+        ) :
+        
+        if self.config.num_speakers > 1 and speaker_id is not None:
+            if isinstance(speaker_id, int):
+                speaker_id = torch.full(size=(1,), fill_value=speaker_id, device=self.device)
+            elif isinstance(speaker_id, (list, tuple, np.ndarray)):
+                speaker_id = torch.tensor(speaker_id, device=self.device)
+
+            if not ((0 <= speaker_id).all() and (speaker_id < self.config.num_speakers).all()).item():
+                raise ValueError(f"Set `speaker_id` in the range 0-{self.config.num_speakers - 1}.")
+            
+            if not (len(speaker_id) == 1 or len(speaker_id == len(labels))):
+                raise ValueError(
+                    f"You passed {len(speaker_id)} `speaker_id` but you should either pass one speaker id or `batch_size` `speaker_id`."
+                )
+
+            speaker_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1)
+        else:
+            speaker_embeddings = None
+        
+        
+        if labels_attention_mask is not None:
+            labels_padding_mask = labels_attention_mask.unsqueeze(1).float()
+        else:
+            labels_attention_mask = torch.ones((labels.shape[0], labels.shape[2])).float().to(self.device)
+            labels_padding_mask = labels_attention_mask.unsqueeze(1)
+        
+        
+        posterior_latents, posterior_means, posterior_log_variances = self.posterior_encoder(
+            labels, labels_padding_mask, speaker_embeddings
+        )
+        
+        label_lengths = labels_attention_mask.sum(dim=1)
+        latents_slice, ids_slice = self.rand_slice_segments(posterior_latents, 
+                                                            label_lengths, 
+                                                            segment_size=self.segment_size
+                                                            )
+        
+        waveform = self.decoder(latents_slice, speaker_embeddings)
+        
+        if not return_dict:
+            outputs = (
+                labels_padding_mask,
+                posterior_latents,
+                posterior_means,
+                posterior_log_variances,
+                latents_slice,
+                ids_slice,
+                waveform,
+            )
+            return outputs
+        
+        return PosteriorDecoderModelOutput(
+                labels_padding_mask = labels_padding_mask,
+                posterior_latents = posterior_latents,
+                posterior_means = posterior_means,
+                posterior_log_variances = posterior_log_variances,
+                latents_slice = latents_slice,
+                ids_slice = ids_slice,
+                waveform = waveform,
+        )
+    
+    
+    
+    #....................................
+    
+    def trainer(self,
+              train_dataset_dir = None,
+              eval_dataset_dir = None,
+              full_generation_dir = None,
+              feature_extractor = VitsFeatureExtractor(),
+              training_args = None,
+              full_generation_sample_index= 0,
+              project_name = "Posterior_Decoder_Finetuning",
+              wandbKey = "782b6a6e82bbb5a5348de0d3c7d40d1e76351e79",
+              
+              
+              ):
+        
+        os.makedirs(training_args.output_dir,exist_ok=True)
+        logger = logging.getLogger(f"{__name__} Training")
+        log_level = training_args.get_process_log_level()
+        logger.setLevel(log_level)
+        
+        wandb.login(key= wandbKey)
+        wandb.init(project= project_name,config = training_args.to_dict())
+        
+        
+        set_seed(training_args.seed)
+        # Apply Weight Norm Decoder
+        self.decoder.apply_weight_norm()
+        # Save Config
+        self.config.save_pretrained(training_args.output_dir)
+        
+        train_dataset = FeaturesCollectionDataset(dataset_dir = train_dataset_dir,
+                                                  device = self.device
+                                                  )
+        
+        eval_dataset = None
+        if training_args.do_eval:
+            eval_dataset = FeaturesCollectionDataset(dataset_dir = eval_dataset_dir,
+                                                     device = self.device
+                                                     )
+        
+        full_generation_dataset = FeaturesCollectionDataset(dataset_dir = full_generation_dir,
+                                                            device = self.device
+                                                            )
+        self.full_generation_sample = full_generation_dataset[full_generation_sample_index]
+      
+        # init optimizer, lr_scheduler 
+        
+        optimizer = torch.optim.AdamW(
+            self.parameters(),
+            training_args.learning_rate,
+            betas=[training_args.adam_beta1, training_args.adam_beta2],
+            eps=training_args.adam_epsilon,
+        )
+        
+        lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
+            optimizer, gamma=training_args.lr_decay, last_epoch=-1
+            )
+        
+        
+        logger.info("***** Running training *****")
+        logger.info(f"  Num Epochs = {training_args.num_train_epochs}")
+
+        
+        #.......................loop training............................
+       
+        global_step = 0
+     
+        for epoch in range(training_args.num_train_epochs):
+            train_losses_sum = 0
+            lr_scheduler.step()
+            
+            for step, batch in enumerate(train_dataset):
+                
+                # forward through model
+                outputs = self.forward(
+                    labels=batch["labels"],
+                    labels_attention_mask=batch["labels_attention_mask"],
+                    speaker_id=batch["speaker_id"]
+                    )
+                
+                mel_scaled_labels = batch["mel_scaled_input_features"]
+                mel_scaled_target = self.slice_segments(mel_scaled_labels, outputs.ids_slice,self.segment_size)
+                mel_scaled_generation = feature_extractor._torch_extract_fbank_features(outputs.waveform.squeeze(1))[1]
+
+                target_waveform = batch["waveform"].transpose(1, 2)
+                target_waveform = self.slice_segments(
+                                    target_waveform, 
+                                    outputs.ids_slice * feature_extractor.hop_length, 
+                                    self.config.segment_size
+                                )
+                
+                
+                # backpropagate
+                
+                loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
+                loss = loss_mel.detach().item()
+                train_losses_sum = train_losses_sum + loss
+                loss_mel.backward()
+                optimizer.step()
+                optimizer.zero_grad()
+                
+                print(f"TRAINIG - batch {step}, waveform {(batch['waveform'].shape)}, step_loss_mel {loss}, lr {lr_scheduler.get_last_lr()[0]}... ")
+                global_step +=1
+                
+                # validation  
+                
+                do_eval = training_args.do_eval and (global_step % training_args.eval_steps == 0)
+                if do_eval:
+                    logger.info("Running validation... ")
+                    eval_losses_sum = 0
+                    for step, batch in enumerate(eval_dataset):
+                       
+                        with torch.no_grad():
+                            outputs = self.forward(
+                                labels=batch["labels"],
+                                labels_attention_mask=batch["labels_attention_mask"],
+                                speaker_id=batch["speaker_id"]
+                                )
+                        
+                        mel_scaled_labels = batch["mel_scaled_input_features"]
+                        mel_scaled_target = self.slice_segments(mel_scaled_labels, outputs.ids_slice,self.segment_size)
+                        mel_scaled_generation = feature_extractor._torch_extract_fbank_features(outputs.waveform.squeeze(1))[1]
+                        loss = loss_mel.detach().item()
+                        eval_losses_sum +=loss
+                        loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
+                        print(f"VALIDATION - batch {step}, waveform {(batch['waveform'].shape)}, step_loss_mel {loss} ... ")
+                        
+                    
+                    
+                    with torch.no_grad():
+                        full_generation_sample = self.full_generation_sample
+                        full_generation =self.forward(
+                                labels=full_generation_sample["labels"],
+                                labels_attention_mask=full_generation_sample["labels_attention_mask"],
+                                speaker_id=full_generation_sample["speaker_id"]
+                                )
+                    
+                        full_generation_waveform = full_generation.waveform.cpu().numpy()
+                    
+                    wandb.log({
+                    "eval_losses": eval_losses_sum,
+                    "full generations samples": [
+                        wandb.Audio(w.reshape(-1), caption=f"Full generation sample {epoch}", sample_rate=self.sampling_rate)
+                        for w in full_generation_waveform],})
+                
+            wandb.log({"train_losses":train_losses_sum})
+            
+        # add weight norms
+        self.decoder.remove_weight_norm()
+        
+        
+        torch.save(self.posterior_encoder.state_dict(), os.path.join(training_args.output_dir,"posterior_encoder.pt"))
+        torch.save(self.decoder.state_dict(), os.path.join(training_args.output_dir,"decoder.pt"))
+
+      
+        
+        logger.info("Running final full generations samples... ")
+        
+        
+        with torch.no_grad():
+            
+            full_generation_sample = self.full_generation_sample
+            full_generation = self.forward(
+                    labels=full_generation_sample["labels"],
+                    labels_attention_mask=full_generation_sample["labels_attention_mask"],
+                    speaker_id=full_generation_sample["speaker_id"]
+                    )
+       
+            full_generation_waveform = full_generation.waveform.cpu().numpy()
+            
+        wandb.log({"eval_losses": eval_losses_sum,
+                "full generations samples": [
+                    wandb.Audio(w.reshape(-1), caption=f"Full generation sample {epoch}",
+                                sample_rate=self.sampling_rate) for w in full_generation_waveform],
+                })
+        
+        
+        logger.info("***** Training / Inference Done *****")
+        
+    #....................................
+    
+    
+    
+    
+    #....................................

VitsModelSplit/Trainer.py

@@ -0,0 +1,848 @@
+import os
+import shutil
+import tempfile
+import numpy as np
+import wandb
+from transformers import VitsModel
+import math
+import torch
+from accelerate.utils import ProjectConfiguration, is_wandb_available, set_seed
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from transformers.utils import send_example_telemetry
+import logging
+import sys
+from transformers.trainer_utils import get_last_checkpoint, is_main_process
+from transformers.trainer_pt_utils import LengthGroupedSampler
+from transformers.optimization import get_scheduler
+
+
+from .data_collator import DataCollatorTTSWithPadding
+from .discriminator import VitsDiscriminator
+from .feature_extraction import VitsFeatureExtractor
+from .plot import plot_alignment_to_numpy, plot_spectrogram_to_numpy
+
+#.............................................
+
+if is_wandb_available():
+    import wandb
+
+ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
+logger = logging.getLogger(__name__)
+#.............................................
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    real_losses = 0
+    generated_losses = 0
+    for disc_real, disc_generated in zip(disc_real_outputs, disc_generated_outputs):
+        real_loss = torch.mean((1 - disc_real) ** 2)
+        generated_loss = torch.mean(disc_generated**2)
+        loss += real_loss + generated_loss
+        real_losses += real_loss
+        generated_losses += generated_loss
+
+    return loss, real_losses, generated_losses
+
+
+def feature_loss(feature_maps_real, feature_maps_generated):
+    loss = 0
+    for feature_map_real, feature_map_generated in zip(feature_maps_real, feature_maps_generated):
+        for real, generated in zip(feature_map_real, feature_map_generated):
+            real = real.detach()
+            loss += torch.mean(torch.abs(real - generated))
+
+    return loss * 2
+
+
+def generator_loss(disc_outputs):
+    total_loss = 0
+    gen_losses = []
+    for disc_output in disc_outputs:
+        disc_output = disc_output
+        loss = torch.mean((1 - disc_output) ** 2)
+        gen_losses.append(loss)
+        total_loss += loss
+
+    return total_loss, gen_losses
+
+
+def kl_loss(prior_latents, posterior_log_variance, prior_means, prior_log_variance, labels_mask):
+    """
+    z_p, logs_q: [b, h, t_t]
+    prior_means, prior_log_variance: [b, h, t_t]
+    """
+
+    kl = prior_log_variance - posterior_log_variance - 0.5
+    kl += 0.5 * ((prior_latents - prior_means) ** 2) * torch.exp(-2.0 * prior_log_variance)
+    kl = torch.sum(kl * labels_mask)
+    loss = kl / torch.sum(labels_mask)
+    return loss
+
+
+def log_on_trackers(
+    trackers,
+    generated_audio,
+    generated_attn,
+    generated_spec,
+    target_spec,
+    full_generation_waveform,
+    epoch,
+    sampling_rate,
+):
+    max_num_samples = min(len(generated_audio), 50)
+    generated_audio = generated_audio[:max_num_samples]
+    generated_attn = generated_attn[:max_num_samples]
+    generated_spec = generated_spec[:max_num_samples]
+    target_spec = target_spec[:max_num_samples]
+
+    for tracker in trackers:
+        if tracker.name == "tensorboard":
+            for cpt, audio in enumerate(generated_audio):
+                tracker.writer.add_audio(f"train_step_audio_{cpt}", audio[None, :], epoch, sample_rate=sampling_rate)
+
+            for cpt, audio in enumerate(full_generation_waveform):
+                tracker.writer.add_audio(
+                    f"full_generation_sample{cpt}", audio[None, :], epoch, sample_rate=sampling_rate
+                )
+
+            tracker.writer.add_images("alignements", np.stack(generated_attn), dataformats="NHWC")
+            tracker.writer.add_images("spectrogram", np.stack(generated_spec), dataformats="NHWC")
+            tracker.writer.add_images("target spectrogram", np.stack(target_spec), dataformats="NHWC")
+        elif tracker.name == "wandb":
+            # wandb can only loads 100 audios per step
+            tracker.log(
+                {
+                    "alignments": [wandb.Image(attn, caption=f"Audio epoch {epoch}") for attn in generated_attn],
+                    "spectrogram": [wandb.Image(spec, caption=f"Audio epoch {epoch}") for spec in generated_spec],
+                    "target spectrogram": [wandb.Image(spec, caption=f"Audio epoch {epoch}") for spec in target_spec],
+                    "train generated audio": [
+                        wandb.Audio(
+                            audio[0],
+                            caption=f"Audio during train step epoch {epoch}",
+                            sample_rate=sampling_rate,
+                        )
+                        for audio in generated_audio
+                    ],
+                    "full generations samples": [
+                        wandb.Audio(w, caption=f"Full generation sample {epoch}", sample_rate=sampling_rate)
+                        for w in full_generation_waveform
+                    ],
+                }
+            )
+        else:
+            logger.warn(f"audio logging not implemented for {tracker.name}")
+
+
+def compute_val_metrics_and_losses(
+    val_losses,
+    accelerator,
+    model_outputs,
+    mel_scaled_generation,
+    mel_scaled_target,
+    batch_size,
+    compute_clap_similarity=False,
+):
+    loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
+    loss_kl = kl_loss(
+        model_outputs.prior_latents,
+        model_outputs.posterior_log_variances,
+        model_outputs.prior_means,
+        model_outputs.prior_log_variances,
+        model_outputs.labels_padding_mask,
+    )
+
+    losses_mel_kl = loss_mel + loss_kl
+
+    losses = torch.stack([loss_mel, loss_kl, losses_mel_kl])
+    losses = accelerator.gather(losses.repeat(batch_size, 1)).mean(0)
+
+    for key, loss in zip(["val_loss_mel", "val_loss_kl", "val_loss_mel_kl"], losses):
+        val_losses[key] = val_losses.get(key, 0) + loss.item()
+
+    return val_losses
+
+
+#.............................................
+
+def vits_trainin(
+                model,
+                tokenizer,
+                model_args,
+                data_args,
+                training_args,
+                train_dataset,
+                eval_dataset,
+                
+                ):
+    
+    
+
+
+    send_example_telemetry("run_vits_finetuning", model_args, data_args)
+    
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(log_level)
+    # datasets.utils.logging.set_verbosity(log_level)
+    # transformers.utils.logging.set_verbosity(log_level)
+    # transformers.utils.logging.enable_default_handler()
+    # transformers.utils.logging.enable_explicit_format()
+    # # logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)
+    # if is_main_process(training_args.local_rank):
+    #     transformers.utils.logging.set_verbosity_info()    
+
+    
+    
+     
+    set_seed(training_args.seed)
+    
+ 
+    
+    config = model.config
+    feature_extractor = VitsFeatureExtractor()
+    
+    forward_attention_mask = True
+    
+    
+    with training_args.main_process_first(desc="apply_weight_norm"):
+        # apply weight norms
+        model.decoder.apply_weight_norm()
+        for flow in model.flow.flows:
+            torch.nn.utils.weight_norm(flow.conv_pre)
+            torch.nn.utils.weight_norm(flow.conv_post)
+    
+    
+    
+    with training_args.main_process_first():
+        # only the main process saves them
+        if is_main_process(training_args.local_rank):
+            # save feature extractor, tokenizer and config
+            feature_extractor.save_pretrained(training_args.output_dir)
+            tokenizer.save_pretrained(training_args.output_dir)
+            config.save_pretrained(training_args.output_dir)
+    
+    
+    data_collator = DataCollatorTTSWithPadding(
+        tokenizer=tokenizer,
+        feature_extractor=feature_extractor,
+        forward_attention_mask=forward_attention_mask,
+    )
+
+    with training_args.main_process_first():
+        input_str = data_args.full_generation_sample_text
+        full_generation_sample = tokenizer(input_str, return_tensors="pt")
+    
+  
+    project_name = data_args.project_name
+    logging_dir = os.path.join(training_args.output_dir, training_args.logging_dir)
+    accelerator_project_config = ProjectConfiguration(project_dir=training_args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+        log_with=training_args.report_to,
+        project_config=accelerator_project_config,
+        kwargs_handlers=[ddp_kwargs],
+    )
+
+    per_device_train_batch_size = (
+        training_args.per_device_train_batch_size if training_args.per_device_train_batch_size else 1
+    )
+    total_batch_size = (
+        per_device_train_batch_size * accelerator.num_processes * training_args.gradient_accumulation_steps
+    )
+
+    num_speakers = model.config.num_speakers
+    if training_args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+    
+        
+    
+    train_dataloader = None
+    if training_args.do_train:
+        sampler = (
+            LengthGroupedSampler(
+                batch_size=per_device_train_batch_size,
+                dataset=train_dataset,
+                lengths=train_dataset["tokens_input_length"],
+            )
+            if training_args.group_by_length
+            else None
+        )
+        train_dataloader = torch.utils.data.DataLoader(
+            train_dataset,
+            shuffle=False,#not training_args.group_by_length,
+            collate_fn=data_collator,
+            batch_size=training_args.per_device_train_batch_size,
+            num_workers=training_args.dataloader_num_workers,
+            sampler=sampler,
+        )
+
+    eval_dataloader = None
+    if training_args.do_eval:
+        eval_sampler = (
+            LengthGroupedSampler(
+                batch_size=training_args.per_device_eval_batch_size,
+                dataset=eval_dataset,
+                lengths=eval_dataset["tokens_input_length"],
+            )
+            if training_args.group_by_length
+            else None
+        )
+
+        eval_dataloader = torch.utils.data.DataLoader(
+            eval_dataset,
+            shuffle=False,
+            collate_fn=data_collator,
+            batch_size=training_args.per_device_eval_batch_size,
+            num_workers=training_args.dataloader_num_workers,
+            sampler=eval_sampler,
+        )
+
+    model_segment_size = model.segment_size
+    config_segment_size = model.config.segment_size
+    sampling_rate = model.config.sampling_rate
+    
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / training_args.gradient_accumulation_steps)
+    if training_args.max_steps == -1:
+        training_args.max_steps = training_args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / training_args.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        training_args.max_steps = int(training_args.num_train_epochs * num_update_steps_per_epoch)
+    # Afterwards we recalculate our number of training epochs
+    training_args.num_train_epochs = math.ceil(training_args.max_steps / num_update_steps_per_epoch)
+
+    # hack to be able to train on multiple device
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        model.discriminator.save_pretrained(tmpdirname)
+        discriminator = VitsDiscriminator.from_pretrained(tmpdirname)
+        for disc in discriminator.discriminators:
+            disc.apply_weight_norm()
+    del model.discriminator
+    
+    # init gen_optimizer, gen_lr_scheduler, disc_optimizer, dics_lr_scheduler
+    gen_optimizer = torch.optim.AdamW(
+        model.parameters(),
+        training_args.learning_rate,
+        betas=[training_args.adam_beta1, training_args.adam_beta2],
+        eps=training_args.adam_epsilon,
+    )
+
+    disc_optimizer = torch.optim.AdamW(
+        discriminator.parameters(),
+        training_args.learning_rate,
+        betas=[training_args.adam_beta1, training_args.adam_beta2],
+        eps=training_args.adam_epsilon,
+    )
+
+    num_warmups_steps = training_args.get_warmup_steps(training_args.num_train_epochs * accelerator.num_processes)
+    num_training_steps = training_args.num_train_epochs * accelerator.num_processes
+    
+    gen_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
+        gen_optimizer, gamma=training_args.lr_decay, last_epoch=-1
+    )
+    disc_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
+        disc_optimizer, gamma=training_args.lr_decay, last_epoch=-1
+    )
+   
+
+    # Prepare everything with our `accelerator`.
+    (
+        model,
+        discriminator,
+        gen_optimizer,
+        gen_lr_scheduler,
+        disc_optimizer,
+        disc_lr_scheduler,
+        train_dataloader,
+        eval_dataloader,
+    ) = accelerator.prepare(
+        model,
+        discriminator,
+        gen_optimizer,
+        gen_lr_scheduler,
+        disc_optimizer,
+        disc_lr_scheduler,
+        train_dataloader,
+        eval_dataloader,
+    )
+
+    
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if accelerator.is_main_process:
+        tracker_config = training_args.to_sanitized_dict()
+        accelerator.init_trackers(project_name, tracker_config)
+    
+    
+    
+    # Train!
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {training_args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {training_args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {training_args.max_steps}")
+    global_step = 0
+    first_epoch = 0
+    
+    
+    
+    # Potentially load in the weights and states from a previous save
+    if training_args.resume_from_checkpoint:
+        if training_args.resume_from_checkpoint != "latest":
+            path = os.path.basename(training_args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(training_args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+
+        if path is None:
+            accelerator.print(
+                f"Checkpoint '{training_args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            training_args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            accelerator.print(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(training_args.output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+
+    else:
+        initial_global_step = 0
+
+
+    
+    #.......................loop training............................
+
+    for epoch in range(first_epoch, training_args.num_train_epochs):
+        # keep track of train losses
+        train_losses = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+        
+        disc_lr_scheduler.step()
+        gen_lr_scheduler.step()
+        
+        for step, batch in enumerate(train_dataloader):
+            print(f"TRAINIG - batch {step}, process{accelerator.process_index}, waveform {(batch['waveform'].shape)}, tokens {(batch['input_ids'].shape)}... ")
+            with accelerator.accumulate(model, discriminator):
+                # forward through model
+                model_outputs = model(
+                    input_ids=batch["input_ids"],
+                    attention_mask=batch["attention_mask"],
+                    labels=batch["labels"],
+                    labels_attention_mask=batch["labels_attention_mask"],
+                    speaker_id=batch["speaker_id"],
+                    encoder_output = batch['text_encoder_output'],
+                    
+                    return_dict=True,
+                    monotonic_alignment_function=None,
+                )
+                
+                mel_scaled_labels = batch["mel_scaled_input_features"]
+                mel_scaled_target = model.slice_segments(mel_scaled_labels, model_outputs.ids_slice, model_segment_size)
+                mel_scaled_generation = feature_extractor._torch_extract_fbank_features(
+                    model_outputs.waveform.squeeze(1)
+                )[1]
+
+                target_waveform = batch["waveform"].transpose(1, 2)
+                target_waveform = model.slice_segments(
+                    target_waveform, model_outputs.ids_slice * feature_extractor.hop_length, config_segment_size
+                )
+
+                # -----------------------
+                #  Train Discriminator
+                # -----------------------
+            
+                discriminator_target, _ = discriminator(target_waveform)
+                discriminator_candidate, _ = discriminator(model_outputs.waveform.detach())
+
+                loss_disc, loss_real_disc, loss_fake_disc = discriminator_loss(
+                    discriminator_target, discriminator_candidate
+                )
+
+                # backpropagate
+                accelerator.backward(loss_disc * training_args.weight_disc)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(discriminator.parameters(), training_args.max_grad_norm)
+                disc_optimizer.step()
+                if not training_args.do_step_schedule_per_epoch:
+                    disc_lr_scheduler.step()
+                disc_optimizer.zero_grad()
+
+                # -----------------------
+                #  Train Generator
+                # -----------------------
+                
+                _, fmaps_target = discriminator(target_waveform)
+                discriminator_candidate, fmaps_candidate = discriminator(model_outputs.waveform)
+
+                loss_duration = torch.sum(model_outputs.log_duration)
+                loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
+                loss_kl = kl_loss(
+                    model_outputs.prior_latents,
+                    model_outputs.posterior_log_variances,
+                    model_outputs.prior_means,
+                    model_outputs.prior_log_variances,
+                    model_outputs.labels_padding_mask,
+                )
+                loss_fmaps = feature_loss(fmaps_target, fmaps_candidate)
+                loss_gen, losses_gen = generator_loss(discriminator_candidate)
+
+                total_generator_loss = (
+                    loss_duration * training_args.weight_duration
+                    + loss_mel * training_args.weight_mel
+                    + loss_kl * training_args.weight_kl
+                    + loss_fmaps * training_args.weight_fmaps
+                    + loss_gen * training_args.weight_gen
+                )
+
+                # backpropagate
+                accelerator.backward(total_generator_loss)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), training_args.max_grad_norm)
+                gen_optimizer.step()
+                if not training_args.do_step_schedule_per_epoch:
+                    gen_lr_scheduler.step()
+                gen_optimizer.zero_grad()
+
+                # update and gather losses
+                losses = torch.stack(
+                    [
+                        # for fair comparison, don't use weighted loss
+                        loss_duration + loss_mel + loss_kl + loss_fmaps + loss_gen,
+                        loss_duration,
+                        loss_mel,
+                        loss_kl,
+                        loss_fmaps,
+                        loss_gen,
+                        loss_disc,
+                        loss_real_disc,
+                        loss_fake_disc,
+                    ]
+                )
+                losses = accelerator.gather(losses.repeat(per_device_train_batch_size, 1)).mean(0)
+
+                train_losses = [
+                    l + losses[i].item() / training_args.gradient_accumulation_steps
+                    for (i, l) in enumerate(train_losses)
+                ]
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                (
+                    train_summed_losses,
+                    train_loss_duration,
+                    train_loss_mel,
+                    train_loss_kl,
+                    train_loss_fmaps,
+                    train_loss_gen,
+                    train_loss_disc,
+                    train_loss_real_disc,
+                    train_loss_fake_disc,
+                ) = train_losses
+                
+                global_step += 1
+                accelerator.log(
+                    {
+                        "train_summed_losses": train_summed_losses,
+                        "train_loss_disc": train_loss_disc,
+                        "train_loss_real_disc": train_loss_real_disc,
+                        "train_loss_fake_disc": train_loss_fake_disc,
+                        "train_loss_duration": train_loss_duration,
+                        "train_loss_mel": train_loss_mel,
+                        "train_loss_kl": train_loss_kl,
+                        "train_loss_fmaps": train_loss_fmaps,
+                        "train_loss_gen": train_loss_gen,
+                        "lr": disc_lr_scheduler.get_last_lr()[0],
+                    },
+                    step=global_step,
+                )
+                train_losses = [0.0 for _ in train_losses]
+                
+                if global_step % training_args.save_steps == 0:
+                    if accelerator.is_main_process:
+                        # _before_ saving state, check if this save would set us over the `save_total_limit`
+                        if training_args.save_total_limit is not None:
+                            checkpoints = os.listdir(training_args.output_dir)
+                            checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+                            checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+                            # before we save the new checkpoint, we need to have at _most_ `save_total_limit - 1` checkpoints
+                            if len(checkpoints) >= training_args.save_total_limit:
+                                num_to_remove = len(checkpoints) - training_args.save_total_limit + 1
+                                removing_checkpoints = checkpoints[0:num_to_remove]
+
+                                logger.info(
+                                    f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+                                )
+                                logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+                                for removing_checkpoint in removing_checkpoints:
+                                    removing_checkpoint = os.path.join(training_args.output_dir, removing_checkpoint)
+                                    shutil.rmtree(removing_checkpoint)
+
+                        save_path = os.path.join(training_args.output_dir, f"checkpoint-{global_step}")
+                        accelerator.save_state(save_path)
+                        logger.info(f"Saved state to {save_path}")
+
+            logs = {
+                "step_loss": total_generator_loss.detach().item(),
+                "lr": disc_lr_scheduler.get_last_lr()[0],
+                "step_loss_duration": loss_duration.detach().item(),
+                "step_loss_mel": loss_mel.detach().item(),
+                "step_loss_kl": loss_kl.detach().item(),
+                "step_loss_fmaps": loss_fmaps.detach().item(),
+                "step_loss_gen": loss_gen.detach().item(),
+                "step_loss_disc": loss_disc.detach().item(),
+                "step_loss_real_disc": loss_real_disc.detach().item(),
+                "step_loss_fake_disc": loss_fake_disc.detach().item(),
+            }
+           
+
+            if global_step >= training_args.max_steps:
+                break
+
+            eval_steps = training_args.eval_steps if training_args.eval_steps else 1
+            do_eval = training_args.do_eval and (global_step % eval_steps == 0) and accelerator.sync_gradients
+
+            if do_eval:
+                logger.info("Running validation... ")
+                generated_audio = []
+                generated_attn = []
+                generated_spec = []
+                target_spec = []
+                val_losses = {}
+                for step, batch in enumerate(eval_dataloader):
+                    print(
+                        f"VALIDATION - batch {step}, process{accelerator.process_index}, waveform {(batch['waveform'].shape)}, tokens {(batch['input_ids'].shape)}... "
+                    )
+                    with torch.no_grad():
+                        model_outputs_train = model(
+                            input_ids=batch["input_ids"],
+                            attention_mask=batch["attention_mask"],
+                            labels=batch["labels"],
+                            labels_attention_mask=batch["labels_attention_mask"],
+                            speaker_id=batch["speaker_id"],
+                            encoder_output = batch['text_encoder_output'],
+                            
+                            return_dict=True,
+                            monotonic_alignment_function=None,
+                        )
+
+                        mel_scaled_labels = batch["mel_scaled_input_features"]
+                        mel_scaled_target = model.slice_segments(
+                            mel_scaled_labels, model_outputs_train.ids_slice, model_segment_size
+                        )
+                        mel_scaled_generation = feature_extractor._torch_extract_fbank_features(
+                            model_outputs_train.waveform.squeeze(1)
+                        )[1]
+
+                        val_losses = compute_val_metrics_and_losses(
+                            val_losses,
+                            accelerator,
+                            model_outputs_train,
+                            mel_scaled_generation,
+                            mel_scaled_target,
+                            per_device_train_batch_size,
+                            compute_clap_similarity=False,
+                        )
+
+                    print(f"VALIDATION - batch {step}, process{accelerator.process_index}, PADDING AND GATHER... ")
+                    specs = feature_extractor._torch_extract_fbank_features(model_outputs_train.waveform.squeeze(1))[0]
+                    padded_attn, specs, target_specs = accelerator.pad_across_processes(
+                        [model_outputs_train.attn.squeeze(1), specs, batch["labels"]], dim=1
+                    )
+                    padded_attn, specs, target_specs = accelerator.pad_across_processes(
+                        [padded_attn, specs, target_specs], dim=2
+                    )
+
+                    generated_train_waveform, padded_attn, specs, target_specs = accelerator.gather_for_metrics(
+                        [model_outputs_train.waveform, padded_attn, specs, target_specs]
+                    )
+
+
+                    if accelerator.is_main_process:
+                        with torch.no_grad():
+                            speaker_id = None if num_speakers < 2 else list(range(min(5, num_speakers)))
+                            full_generation = model(**full_generation_sample.to(model.device), speaker_id=speaker_id)
+
+                        generated_audio.append(generated_train_waveform.cpu())
+                        generated_attn.append(padded_attn.cpu())
+                        generated_spec.append(specs.cpu())
+                        target_spec.append(target_specs.cpu())
+
+                logger.info("Validation inference done, now evaluating... ")
+                if accelerator.is_main_process:
+                    generated_audio = [audio.numpy() for audio_batch in generated_audio for audio in audio_batch]
+                    generated_attn = [
+                        plot_alignment_to_numpy(attn.numpy()) for attn_batch in generated_attn for attn in attn_batch
+                    ]
+                    generated_spec = [
+                        plot_spectrogram_to_numpy(attn.numpy()) for attn_batch in generated_spec for attn in attn_batch
+                    ]
+                    target_spec = [
+                        plot_spectrogram_to_numpy(attn.numpy()) for attn_batch in target_spec for attn in attn_batch
+                    ]
+                    full_generation_waveform = full_generation.waveform.cpu().numpy()
+
+                    accelerator.log(val_losses, step=global_step)
+
+                    log_on_trackers(
+                        accelerator.trackers,
+                        generated_audio,
+                        generated_attn,
+                        generated_spec,
+                        target_spec,
+                        full_generation_waveform,
+                        epoch,
+                        sampling_rate,
+                    )
+
+                    logger.info("Validation finished... ")
+
+                accelerator.wait_for_everyone()
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        epoch = training_args.num_train_epochs if training_args.num_train_epochs else 1
+        eval_steps = training_args.eval_steps if training_args.eval_steps else 1
+
+        # Run a final round of inference.
+        do_eval = training_args.do_eval
+
+        if do_eval:
+            logger.info("Running final validation... ")
+            generated_audio = []
+            generated_attn = []
+            generated_spec = []
+            target_spec = []
+            val_losses = {}
+            for step, batch in enumerate(eval_dataloader):
+                print(
+                    f"VALIDATION - batch {step}, process{accelerator.process_index}, waveform {(batch['waveform'].shape)}, tokens {(batch['input_ids'].shape)}... "
+                )
+                with torch.no_grad():
+                    model_outputs_train = model(
+                        input_ids=batch["input_ids"],
+                        attention_mask=batch["attention_mask"],
+                        labels=batch["labels"],
+                        labels_attention_mask=batch["labels_attention_mask"],
+                        speaker_id=batch["speaker_id"],
+                        encoder_output = batch['text_encoder_output'],
+                        
+                        return_dict=True,
+                        monotonic_alignment_function=None,
+                    )
+
+                    mel_scaled_labels = batch["mel_scaled_input_features"]
+                    mel_scaled_target = model.slice_segments(
+                        mel_scaled_labels, model_outputs_train.ids_slice, model_segment_size
+                    )
+                    mel_scaled_generation = feature_extractor._torch_extract_fbank_features(
+                        model_outputs_train.waveform.squeeze(1)
+                    )[1]
+
+                    val_losses = compute_val_metrics_and_losses(
+                        val_losses,
+                        accelerator,
+                        model_outputs_train,
+                        mel_scaled_generation,
+                        mel_scaled_target,
+                        per_device_train_batch_size,
+                        compute_clap_similarity=False,
+                    )
+                specs = feature_extractor._torch_extract_fbank_features(model_outputs_train.waveform.squeeze(1))[0]
+                padded_attn, specs, target_specs = accelerator.pad_across_processes(
+                    [model_outputs_train.attn.squeeze(1), specs, batch["labels"]], dim=1
+                )
+                padded_attn, specs, target_specs = accelerator.pad_across_processes(
+                    [padded_attn, specs, target_specs], dim=2
+                )
+
+                generated_train_waveform, padded_attn, specs, target_specs = accelerator.gather_for_metrics(
+                    [model_outputs_train.waveform, padded_attn, specs, target_specs]
+                )
+
+                if accelerator.is_main_process:
+                    with torch.no_grad():
+                        speaker_id = None if num_speakers < 2 else list(range(min(5, num_speakers)))
+                        full_generation = model(**full_generation_sample.to(model.device), speaker_id=speaker_id)
+
+                    generated_audio.append(generated_train_waveform.cpu())
+                    generated_attn.append(padded_attn.cpu())
+                    generated_spec.append(specs.cpu())
+                    target_spec.append(target_specs.cpu())
+
+            logger.info("Validation inference done, now evaluating... ")
+            if accelerator.is_main_process:
+                generated_audio = [audio.numpy() for audio_batch in generated_audio for audio in audio_batch]
+                generated_attn = [
+                    plot_alignment_to_numpy(attn.numpy()) for attn_batch in generated_attn for attn in attn_batch
+                ]
+                generated_spec = [
+                    plot_spectrogram_to_numpy(attn.numpy()) for attn_batch in generated_spec for attn in attn_batch
+                ]
+                target_spec = [
+                    plot_spectrogram_to_numpy(attn.numpy()) for attn_batch in target_spec for attn in attn_batch
+                ]
+                full_generation_waveform = full_generation.waveform.cpu().numpy()
+
+                log_on_trackers(
+                    accelerator.trackers,
+                    generated_audio,
+                    generated_attn,
+                    generated_spec,
+                    target_spec,
+                    full_generation_waveform,
+                    epoch,
+                    sampling_rate,
+                )
+
+                accelerator.log(val_losses, step=global_step)
+                logger.info("Validation finished... ")
+
+            accelerator.wait_for_everyone()
+
+        # unwrap, save and push final model
+        model = accelerator.unwrap_model(model)
+        discriminator = accelerator.unwrap_model(discriminator)
+
+        model.discriminator = discriminator
+
+        # add weight norms
+        for disc in model.discriminator.discriminators:
+            disc.remove_weight_norm()
+        model.decoder.remove_weight_norm()
+        for flow in model.flow.flows:
+            torch.nn.utils.remove_weight_norm(flow.conv_pre)
+            torch.nn.utils.remove_weight_norm(flow.conv_post)
+
+        model.save_pretrained(training_args.output_dir)
+
+        if training_args.push_to_hub:
+            VitsModel.from_pretrained(training_args.output_dir).push_to_hub(training_args.hub_model_id)
+
+    accelerator.end_training()
+
+
+
+    logger.info("***** Training / Inference Done *****")
+        
+    
+    
+    
+#............................................................................... 

VitsModelSplit/__init__.py

@@ -0,0 +1,4 @@
+
+# from . import vits_config,vits_model,vits_output,arguments,decoder,encoder,discriminator,duration_predictor,flow,posterior_encoder
+# from . import PosteriorDecoderModel,plot,trainer
+# from . import dataset_features_collector,feature_extraction

VitsModelSplit/data_collator.py

@@ -0,0 +1,119 @@
+from typing import Any, Union,List,Dict
+import numpy as np
+import torch
+from dataclasses import dataclass
+from transformers.feature_extraction_utils import BatchFeature
+
+from .vits_output import VitsTextEncoderOutput
+#.............................................
+
+
+@dataclass
+class DataCollatorTTSWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received.
+    Args:
+        tokenizer ([`VitsTokenizer`])
+            The tokenizer used for processing the data.
+        feature_extractor ([`VitsFeatureExtractor`])
+            The tokenizer used for processing the data.
+        forward_attention_mask (`bool`)
+            Whether to return attention_mask.
+    """
+
+    tokenizer: Any
+    feature_extractor: Any
+    forward_attention_mask: bool
+
+    def pad_waveform(self, raw_speech):
+        is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
+        if is_batched_numpy and len(raw_speech.shape) > 2:
+            raise ValueError(f"Only mono-channel audio is supported for input to {self}")
+        is_batched = is_batched_numpy or (
+            isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
+        )
+
+        if is_batched:
+            raw_speech = [np.asarray([speech], dtype=np.float32).T for speech in raw_speech]
+        elif not is_batched and not isinstance(raw_speech, np.ndarray):
+            raw_speech = np.asarray(raw_speech, dtype=np.float32)
+        elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):
+            raw_speech = raw_speech.astype(np.float32)
+
+        # always return batch
+        if not is_batched:
+            raw_speech = [np.asarray([raw_speech]).T]
+
+        batched_speech = BatchFeature({"input_features": raw_speech})
+
+        # convert into correct format for padding
+
+        padded_inputs = self.feature_extractor.pad(
+            batched_speech,
+            padding=True,
+            return_attention_mask=False,
+            return_tensors="pt",
+        )["input_features"]
+
+        return padded_inputs
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+        
+        model_input_name = "input_ids"
+        
+        input_ids = [{model_input_name: feature[model_input_name][0]} for feature in features]
+        
+        # pad input tokens
+        batch = self.tokenizer.pad(input_ids, return_tensors="pt", return_attention_mask=self.forward_attention_mask)
+   
+        # pad waveform
+        waveforms = [np.array(feature["waveform"]) for feature in features]
+        batch["waveform"] = self.pad_waveform(waveforms)
+
+        # pad spectrogram
+        label_features = [np.array(feature["labels"]) for feature in features]
+        labels_batch = self.feature_extractor.pad(
+            {"input_features": [i.T for i in label_features]}, return_tensors="pt", return_attention_mask=True
+        )
+
+        labels = labels_batch["input_features"].transpose(1, 2)
+        batch["labels"] = labels
+        batch["labels_attention_mask"] = labels_batch["attention_mask"]
+
+        # pad mel spectrogram
+        mel_scaled_input_features = {
+            "input_features": [np.array(feature["mel_scaled_input_features"]).squeeze().T for feature in features]
+        }
+        mel_scaled_input_features = self.feature_extractor.pad(
+            mel_scaled_input_features, return_tensors="pt", return_attention_mask=True
+        )["input_features"].transpose(1, 2)
+
+        batch["mel_scaled_input_features"] = mel_scaled_input_features
+        batch["speaker_id"] = (
+            torch.tensor([feature["speaker_id"] for feature in features]) if "speaker_id" in features[0] else None
+        )
+        
+   
+            
+
+        
+        # text_encoder_output = [{
+        #     'last_hidden_state':torch.tensor(features["text_encoder_output"]['last_hidden_state']),
+        #     'prior_log_variances':torch.tensor(feature["text_encoder_output"]['prior_log_variances']),
+        #     'prior_means':torch.tensor(feature["text_encoder_output"]['prior_means']),
+        #     } for feature in features]
+        
+        batch['text_encoder_output'] = VitsTextEncoderOutput(
+                last_hidden_state=torch.tensor(features[0]["text_encoder_output"]['last_hidden_state']),
+                prior_means=torch.tensor(features[0]["text_encoder_output"]['prior_means']),
+                prior_log_variances=torch.tensor(features[0]["text_encoder_output"]['prior_log_variances']),   
+            )
+        
+        # print("DataColl   ",batch.keys())
+        
+        return batch
+
+
+#.............................................................................................

VitsModelSplit/dataset_features_collector.py

@@ -0,0 +1,402 @@
+
+import numpy as np
+import os
+from datasets import Dataset,DatasetDict
+from typing import Union,List,Dict
+import torch
+from dataclasses import dataclass
+from transformers.feature_extraction_utils import BatchFeature
+from VitsModelSplit.feature_extraction import VitsFeatureExtractor
+from VitsModelSplit.vits_model import VitsModel
+from transformers import AutoTokenizer
+
+#.............................................
+
+
+@dataclass
+class DataSetFeaturesCollector:
+
+    def __init__(self,tokenizer,model,feature_extractor,forward_attention_mask=True) -> None:
+        self.tokenizer=tokenizer
+        self.feature_extractor = feature_extractor
+        self.model=model
+        self.forward_attention_mask = forward_attention_mask
+
+    #.............................................
+
+    def pad_waveform(self, raw_speech):
+        
+        is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
+        if is_batched_numpy and len(raw_speech.shape) > 2:
+            raise ValueError(f"Only mono-channel audio is supported for input to {self}")
+        is_batched = is_batched_numpy or (
+            isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
+        )
+
+        if is_batched:
+            raw_speech = [np.asarray([speech], dtype=np.float32).T for speech in raw_speech]
+        elif not is_batched and not isinstance(raw_speech, np.ndarray):
+            raw_speech = np.asarray(raw_speech, dtype=np.float32)
+        elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):
+            raw_speech = raw_speech.astype(np.float32)
+
+        # always return batch
+        if not is_batched:
+            raw_speech = [np.asarray([raw_speech]).T]
+
+        batched_speech = BatchFeature({"input_features": raw_speech})
+
+        # convert into correct format for padding
+
+        padded_inputs = self.feature_extractor.pad(
+            batched_speech,
+            padding=True,
+            return_attention_mask=False,
+            return_tensors="pt",
+        )["input_features"]
+
+        return padded_inputs
+    
+    #.............................................
+    
+    def prepare_dataset(self,batch):
+        
+        sample = batch['audio']
+        audio_inputs = self.feature_extractor(
+                sample,
+                sampling_rate=16000,
+                return_attention_mask=False,
+                do_normalize=False,
+            )
+
+        batch["labels"] = audio_inputs.get("input_features")[0]
+        batch["waveform_input_length"] = len(sample)
+        batch["waveform"] = batch['audio']
+        batch["mel_scaled_input_features"] = audio_inputs.get("mel_scaled_input_features")[0]
+        textsample = batch['text']
+        inputs = self.tokenizer(textsample, return_tensors="pt")
+        inputs = self.tokenizer.pad({'input_ids':inputs.input_ids})
+        batch['input_ids'] = inputs.input_ids
+        batch['attention_mask'] = inputs.attention_mask
+       # batch['speaker_id']=batch['speaker_id']
+
+
+        return batch
+    
+    
+    #.............................................
+
+
+    def __call__(self, dataset: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+
+        dataset = Dataset.from_list(dataset)
+        features = dataset.map(
+                self.prepare_dataset,
+                remove_columns=dataset.column_names,
+                desc="preprocess",
+                )
+
+        features = list(features)
+     
+        model_input_name = "input_ids"
+        
+        input_ids = [{model_input_name: feature[model_input_name][0]} for feature in features]
+        
+        # pad input tokens
+        batch = self.tokenizer.pad(input_ids, return_tensors="pt", return_attention_mask=self.forward_attention_mask)
+   
+        # pad waveform
+        waveforms = [np.array(feature["waveform"]) for feature in features]
+        batch["waveform"] = self.pad_waveform(waveforms)
+
+        # pad spectrogram
+        label_features = [np.array(feature["labels"]) for feature in features]
+        labels_batch = self.feature_extractor.pad(
+            {"input_features": [i.T for i in label_features]}, return_tensors="pt", return_attention_mask=True
+        )
+
+        labels = labels_batch["input_features"].transpose(1, 2)
+        batch["labels"] = labels
+        batch["labels_attention_mask"] = labels_batch["attention_mask"]
+
+        # pad mel spectrogram
+        mel_scaled_input_features = {
+            "input_features": [np.array(feature["mel_scaled_input_features"]).squeeze().T for feature in features]
+        }
+        mel_scaled_input_features = self.feature_extractor.pad(
+            mel_scaled_input_features, return_tensors="pt", return_attention_mask=True
+        )["input_features"].transpose(1, 2)
+
+        batch["mel_scaled_input_features"] = mel_scaled_input_features
+        batch["speaker_id"] = (
+            torch.tensor([feature["speaker_id"] for feature in dataset]) if "speaker_id" in dataset[0] else None
+        )
+        
+        # with torch.no_grad():
+        #     padding_mask  =torch.ones_like(batch['input_ids']).unsqueeze(-1).float()
+        #     text_encoder_output = self.model.text_encoder(batch['input_ids'],
+        #                                                 padding_mask=padding_mask,
+        #                                                 attention_mask = batch['attention_mask']
+        #                                                 )
+        #     batch['text_encoder_output'] = text_encoder_output 
+        #     posterior_latents, posterior_means, posterior_log_variances = self.model.posterior_encoder(
+        #             batch['labels'], batch['labels_attention_mask'].unsqueeze(1).float()
+        #            )
+        #     posterior_encode_output={
+        #       'posterior_latents':posterior_latents,
+        #       'posterior_means':posterior_means,
+        #       'posterior_log_variances':posterior_log_variances
+        #     }
+        #     batch['posterior_encode_output']=posterior_encode_output
+
+
+        
+        return batch
+
+
+#..............................................................
+
+
+
+#.............................................
+
+def run_dataset_features_collection(
+                         dataset_dir,
+                         train_split_name ="train",
+                         eval_split_name="eval",
+                         full_generation_name = 'full_generation',
+                         tokenizer = None,
+                         model = None,
+                         feature_extractor = None,
+                         train_batch_size = 1,
+                         eval_batch_size = 1,
+                         output_dir = "dataset_features"
+                         
+                         ):
+    
+    dataset = DatasetDict.load_from_disk(dataset_dir)
+    
+    data_collator = DataSetFeaturesCollector(
+        tokenizer = tokenizer,
+        model = model,
+        feature_extractor = feature_extractor,
+        forward_attention_mask = True
+        )
+    
+    if train_split_name:
+        train_dataloader = torch.utils.data.DataLoader(
+            dataset[train_split_name],
+            shuffle=False,
+            collate_fn=data_collator,
+            batch_size=train_batch_size,
+            sampler=None,
+        )
+        
+        train_dir = os.path.join(output_dir,"train")
+        os.makedirs(train_dir,exist_ok=True)
+        
+        for step, batch in enumerate(train_dataloader):
+            print(f"Train Dataset - batch {step}, waveform {(batch['waveform'].shape)},tokens {(batch['input_ids'].shape)}... ")
+            fname = os.path.join(train_dir,f"train-batch-{step}.bin")
+            with open(fname, "wb") as f:
+                torch.save(batch, f)
+            
+    if eval_split_name:
+        
+        eval_dataloader = torch.utils.data.DataLoader(
+            dataset[eval_split_name],
+            shuffle=False,
+            collate_fn=data_collator,
+            batch_size=eval_batch_size,
+            sampler=None,
+        )
+        
+        eval_dir = os.path.join(output_dir,"eval")
+        os.makedirs(eval_dir,exist_ok=True)
+        
+        for step, batch in enumerate(eval_dataloader):
+            print(f"Eval Dataset - batch {step}, waveform {(batch['waveform'].shape)},tokens {(batch['input_ids'].shape)}... ")
+            fname = os.path.join(eval_dir,f"eval-batch-{step}.bin")
+            with open(fname, "wb") as f:
+                torch.save(batch, f)  
+    
+    if full_generation_name:
+        
+        full_generation_dataloader = torch.utils.data.DataLoader(
+            dataset[full_generation_name],
+            shuffle=False,
+            collate_fn=data_collator,
+            batch_size=1,
+            sampler=None,
+        )
+        
+        full_generation_dir = os.path.join(output_dir,"full_generation")
+        os.makedirs(full_generation_dir,exist_ok=True)
+        
+        for step, batch in enumerate(full_generation_dataloader):
+            print(f"Full Generation Dataset - batch {step}, waveform {(batch['waveform'].shape)},tokens {(batch['input_ids'].shape)}... ")
+            fname = os.path.join(full_generation_dir,f"full-generation-batch-{step}.bin")
+            with open(fname, "wb") as f:
+                torch.save(batch, f)  
+
+#...........................................................................
+
+import torch.utils.data 
+
+class FeaturesCollectionDataset(torch.utils.data.Dataset):
+    
+    def __init__(self,dataset_dir,device='cpu') -> None:
+        self.dataset_dir = dataset_dir
+        self.batchs_path = sorted([os.path.join(self.dataset_dir,file) for file in os.listdir(dataset_dir) if file.endswith('.bin')])
+        self.device = device
+        
+    def __len__(self):
+        return len(self.batchs_path)
+    
+    def __getitem__(self, idx):
+        batch_name = self.batchs_path[idx]
+        with open(batch_name, "rb") as f:
+            batch = torch.load(f,map_location=torch.device(self.device))
+        return batch
+        
+        
+class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
+    """
+    Maintain similar input lengths in a batch.
+    Length groups are specified by boundaries.
+    Ex) boundaries = [b1, b2, b3] -> any batch is included either {x | b1 < length(x) <=b2} or {x | b2 < length(x) <= b3}.
+  
+    It removes samples which are not included in the boundaries.
+    Ex) boundaries = [b1, b2, b3] -> any x s.t. length(x) <= b1 or length(x) > b3 are discarded.
+    """
+    def __init__(self, dataset, batch_size, boundaries, num_replicas=None, rank=None, shuffle=True):
+        super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)
+        self.lengths =dataset.lengths
+        self.batch_size = batch_size
+        self.boundaries = boundaries
+  
+        self.buckets, self.num_samples_per_bucket = self._create_buckets()
+        self.total_size = sum(self.num_samples_per_bucket)
+        self.num_samples = self.total_size // self.num_replicas
+  
+    def _create_buckets(self):
+        buckets = [[] for _ in range(len(self.boundaries) - 1)]
+        for i in range(len(self.lengths)):
+            length = self.lengths[i]
+            idx_bucket = self._bisect(length)
+            if idx_bucket != -1:
+                buckets[idx_bucket].append(i)
+  
+        for i in range(len(buckets) - 1, 0, -1):
+            if len(buckets[i]) == 0:
+                buckets.pop(i)
+                self.boundaries.pop(i+1)
+  
+        num_samples_per_bucket = []
+        for i in range(len(buckets)):
+            len_bucket = len(buckets[i])
+            total_batch_size = self.num_replicas * self.batch_size
+            rem = (total_batch_size - (len_bucket % total_batch_size)) % total_batch_size
+            num_samples_per_bucket.append(len_bucket + rem)
+        return buckets, num_samples_per_bucket
+  
+    def __iter__(self):
+      # deterministically shuffle based on epoch
+      g = torch.Generator()
+      g.manual_seed(self.epoch)
+  
+      indices = []
+      if self.shuffle:
+          for bucket in self.buckets:
+              indices.append(torch.randperm(len(bucket), generator=g).tolist())
+      else:
+          for bucket in self.buckets:
+              indices.append(list(range(len(bucket))))
+  
+      batches = []
+      for i in range(len(self.buckets)):
+          bucket = self.buckets[i]
+          len_bucket = len(bucket)
+          ids_bucket = indices[i]
+          num_samples_bucket = self.num_samples_per_bucket[i]
+  
+          # add extra samples to make it evenly divisible
+          rem = num_samples_bucket - len_bucket
+          ids_bucket = ids_bucket + ids_bucket * (rem // len_bucket) + ids_bucket[:(rem % len_bucket)]
+  
+          # subsample
+          ids_bucket = ids_bucket[self.rank::self.num_replicas]
+  
+          # batching
+          for j in range(len(ids_bucket) // self.batch_size):
+              batch = [bucket[idx] for idx in ids_bucket[j*self.batch_size:(j+1)*self.batch_size]]
+              batches.append(batch)
+  
+      if self.shuffle:
+          batch_ids = torch.randperm(len(batches), generator=g).tolist()
+          batches = [batches[i] for i in batch_ids]
+      self.batches = batches
+  
+      assert len(self.batches) * self.batch_size == self.num_samples
+      return iter(self.batches)
+  
+    def _bisect(self, x, lo=0, hi=None):
+      if hi is None:
+          hi = len(self.boundaries) - 1
+  
+      if hi > lo:
+          mid = (hi + lo) // 2
+          if self.boundaries[mid] < x and x <= self.boundaries[mid+1]:
+              return mid
+          elif x <= self.boundaries[mid]:
+              return self._bisect(x, lo, mid)
+          else:
+              return self._bisect(x, mid + 1, hi)
+      else:
+          return -1
+
+    def __len__(self):
+        return self.num_samples // self.batch_size
+class VitsCollectionDataset(torch.utils.data.Dataset):
+    
+    def __init__(self,dataset,hop_length=256,rate=16_000,device='cpu') -> None:
+        self.dataset = dataset
+        self.lengths =(torch.tensor(dataset['secs'])*rate//(2*hop_length)).tolist()
+        self.device = device
+
+
+        
+    def __len__(self):
+        return self.dataset.num_rows
+     
+    
+    def __getitem__(self, idx):
+        return self.dataset[idx]
+
+def  get_dataloader(dir_db_train,feature_extractor,name_db='train',batch_size=8,num_workers=0):
+    dataset = DatasetDict.load_from_disk(dir_db_train)
+    db_train=VitsCollectionDataset(dataset[name_db])
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model=VitsModel.from_pretrained("facebook/mms-tts-ara").to(device)
+    tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-ara",cache_dir="./")#.to("cuda")
+    train_sampler = DistributedBucketSampler(
+          db_train,
+          batch_size,
+          [32,300,400,500,600,700,800,900,1000],
+          num_replicas=1,
+          rank=0,
+          shuffle=True)
+    data_collator = DataSetFeaturesCollector(
+        tokenizer = tokenizer,
+        model = model,
+        feature_extractor = feature_extractor,
+        forward_attention_mask = True
+        )
+    train_dataloader = torch.utils.data.DataLoader(
+              db_train,
+              num_workers=num_workers, shuffle=False, pin_memory=True,
+          collate_fn=data_collator, batch_sampler=train_sampler
+            )
+    return train_dataloader   

VitsModelSplit/decoder.py

@@ -0,0 +1,168 @@
+import math
+from typing import Optional
+import numpy as np
+import torch
+from torch import nn
+from .vits_config import VitsConfig
+
+#.............................................
+
+
+
+
+
+# Copied from transformers.models.speecht5.modeling_speecht5.HifiGanResidualBlock
+class HifiGanResidualBlock(nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), leaky_relu_slope=0.1):
+        super().__init__()
+        self.leaky_relu_slope = leaky_relu_slope
+
+        self.convs1 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    stride=1,
+                    dilation=dilation[i],
+                    padding=self.get_padding(kernel_size, dilation[i]),
+                )
+                for i in range(len(dilation))
+            ]
+        )
+        self.convs2 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    stride=1,
+                    dilation=1,
+                    padding=self.get_padding(kernel_size, 1),
+                )
+                for _ in range(len(dilation))
+            ]
+        )
+
+    def get_padding(self, kernel_size, dilation=1):
+        return (kernel_size * dilation - dilation) // 2
+
+    def apply_weight_norm(self):
+        for layer in self.convs1:
+            nn.utils.weight_norm(layer)
+        for layer in self.convs2:
+            nn.utils.weight_norm(layer)
+
+    def remove_weight_norm(self):
+        for layer in self.convs1:
+            nn.utils.remove_weight_norm(layer)
+        for layer in self.convs2:
+            nn.utils.remove_weight_norm(layer)
+
+    def forward(self, hidden_states):
+        for conv1, conv2 in zip(self.convs1, self.convs2):
+            residual = hidden_states
+            hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
+            hidden_states = conv1(hidden_states)
+            hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
+            hidden_states = conv2(hidden_states)
+            hidden_states = hidden_states + residual
+        return hidden_states
+
+
+#.............................................................................................
+
+
+class VitsHifiGan(nn.Module):
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.config = config
+        self.num_kernels = len(config.resblock_kernel_sizes)
+        self.num_upsamples = len(config.upsample_rates)
+        self.conv_pre = nn.Conv1d(
+            config.flow_size,
+            config.upsample_initial_channel,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+        )
+
+        self.upsampler = nn.ModuleList()
+        for i, (upsample_rate, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes)):
+            self.upsampler.append(
+                nn.ConvTranspose1d(
+                    config.upsample_initial_channel // (2**i),
+                    config.upsample_initial_channel // (2 ** (i + 1)),
+                    kernel_size=kernel_size,
+                    stride=upsample_rate,
+                    padding=(kernel_size - upsample_rate) // 2,
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.upsampler)):
+            channels = config.upsample_initial_channel // (2 ** (i + 1))
+            for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes):
+                self.resblocks.append(HifiGanResidualBlock(channels, kernel_size, dilation, config.leaky_relu_slope))
+
+        self.conv_post = nn.Conv1d(channels, 1, kernel_size=7, stride=1, padding=3, bias=False)
+
+        if config.speaker_embedding_size != 0:
+            self.cond = nn.Conv1d(config.speaker_embedding_size, config.upsample_initial_channel, 1)
+
+    def resize_speaker_embedding(self, speaker_embedding_size):
+        self.config.speaker_embedding_size = speaker_embedding_size
+        self.cond = nn.Conv1d(speaker_embedding_size, self.config.upsample_initial_channel, 1)
+        nn.init.kaiming_normal_(self.cond.weight)
+        if self.cond.bias is not None:
+            k = math.sqrt(self.cond.groups / (self.cond.in_channels * self.cond.kernel_size[0]))
+            nn.init.uniform_(self.cond.bias, a=-k, b=k)
+
+    def apply_weight_norm(self):
+        for layer in self.upsampler:
+            nn.utils.weight_norm(layer)
+        for layer in self.resblocks:
+            layer.apply_weight_norm()
+
+    def remove_weight_norm(self):
+        for layer in self.upsampler:
+            nn.utils.remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+
+    def forward(
+        self, spectrogram: torch.FloatTensor, global_conditioning: Optional[torch.FloatTensor] = None
+    ) -> torch.FloatTensor:
+        r"""
+        Converts a spectrogram into a speech waveform.
+
+        Args:
+            spectrogram (`torch.FloatTensor` of shape `(batch_size, config.spectrogram_bins, sequence_length)`):
+                Tensor containing the spectrograms.
+            global_conditioning (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_size, 1)`, *optional*):
+                Tensor containing speaker embeddings, for multispeaker models.
+
+        Returns:
+            `torch.FloatTensor`: Tensor of shape shape `(batch_size, 1, num_frames)` containing the speech waveform.
+        """
+        hidden_states = self.conv_pre(spectrogram)
+
+        if global_conditioning is not None:
+            hidden_states = hidden_states + self.cond(global_conditioning)
+
+        for i in range(self.num_upsamples):
+            hidden_states = nn.functional.leaky_relu(hidden_states, self.config.leaky_relu_slope)
+            hidden_states = self.upsampler[i](hidden_states)
+
+            res_state = self.resblocks[i * self.num_kernels](hidden_states)
+            for j in range(1, self.num_kernels):
+                res_state += self.resblocks[i * self.num_kernels + j](hidden_states)
+            hidden_states = res_state / self.num_kernels
+
+        hidden_states = nn.functional.leaky_relu(hidden_states)
+        hidden_states = self.conv_post(hidden_states)
+        waveform = torch.tanh(hidden_states)
+        return waveform
+
+
+#.............................................................................................

VitsModelSplit/discriminator.py

@@ -0,0 +1,162 @@
+from torch import nn
+import torch
+
+from .vits_config import VitsPreTrainedModel
+
+
+#.............................................
+
+
+class VitsHifiGanDiscriminatorScaleResidualBlock(nn.Module):
+    def __init__(self, discriminator_scale_channels, leaky_relu_slope=0.1):
+        super().__init__()
+        self.leaky_relu_slope = leaky_relu_slope
+
+        in_channels, out_channels = discriminator_scale_channels[:2]
+        self.convs = nn.ModuleList([nn.Conv1d(in_channels, out_channels, 15, 1, padding=7)])
+
+        groups = 4
+        for in_channels, out_channels in zip(discriminator_scale_channels[1:-1], discriminator_scale_channels[2:]):
+            self.convs.append(nn.Conv1d(in_channels, out_channels, 41, 4, groups=groups, padding=20))
+            groups = groups * 4
+
+        channel_size = discriminator_scale_channels[-1]
+        self.convs.append(nn.Conv1d(channel_size, channel_size, 41, 4, groups=groups, padding=20))
+        self.convs.append(nn.Conv1d(channel_size, channel_size, 5, 1, padding=2))
+        self.final_conv = nn.Conv1d(channel_size, 1, 3, 1, padding=1)
+
+    def apply_weight_norm(self):
+        for layer in self.convs:
+            nn.utils.weight_norm(layer)
+        nn.utils.weight_norm(self.final_conv)
+
+    def remove_weight_norm(self):
+        for layer in self.convs:
+            nn.utils.remove_weight_norm(layer)
+        nn.utils.remove_weight_norm(self.final_conv)
+
+    def forward(self, hidden_states):
+        fmap = []
+
+        for conv in self.convs:
+            hidden_states = conv(hidden_states)
+            hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
+            fmap.append(hidden_states)
+
+        hidden_states = self.final_conv(hidden_states)
+        fmap.append(hidden_states)
+        hidden_states = torch.flatten(hidden_states, 1, -1)
+
+        return hidden_states, fmap
+
+
+#.............................................................................................
+
+class VitsHifiGanDiscriminatorPeriodResidualBlock(nn.Module):
+    def __init__(self, discriminator_period_channels, period, kernel_size=5, stride=3, leaky_relu_slope=0.1):
+        super().__init__()
+        self.leaky_relu_slope = leaky_relu_slope
+        self.period = period
+
+        self.convs = nn.ModuleList()
+        for in_channels, out_channels in zip(discriminator_period_channels[:-1], discriminator_period_channels[1:]):
+            self.convs.append(
+                nn.Conv2d(
+                    in_channels,
+                    out_channels,
+                    (kernel_size, 1),
+                    (stride, 1),
+                    padding=(self.get_padding(kernel_size, 1), 0),
+                )
+            )
+
+        channel_size = discriminator_period_channels[-1]
+        self.convs.append(
+            nn.Conv2d(channel_size, channel_size, (kernel_size, 1), 1, padding=(self.get_padding(kernel_size, 1), 0))
+        )
+        self.final_conv = nn.Conv2d(channel_size, 1, (3, 1), 1, padding=(1, 0))
+
+    def get_padding(self, kernel_size, dilation=1):
+        return (kernel_size * dilation - dilation) // 2
+
+    def apply_weight_norm(self):
+        for layer in self.convs:
+            nn.utils.weight_norm(layer)
+        nn.utils.weight_norm(self.final_conv)
+
+    def remove_weight_norm(self):
+        for layer in self.convs:
+            nn.utils.remove_weight_norm(layer)
+        nn.utils.remove_weight_norm(self.final_conv)
+
+    def forward(self, hidden_states):
+        fmap = []
+
+        # from 1D to 2D
+        batch_size, channels, length = hidden_states.shape
+        if length % self.period != 0:
+            # pad first
+            n_pad = self.period - (length % self.period)
+            hidden_states = nn.functional.pad(hidden_states, (0, n_pad), "reflect")
+            length = length + n_pad
+        hidden_states = hidden_states.view(batch_size, channels, length // self.period, self.period)
+
+        for conv in self.convs:
+            hidden_states = conv(hidden_states)
+            hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
+            fmap.append(hidden_states)
+
+        hidden_states = self.final_conv(hidden_states)
+        fmap.append(hidden_states)
+        hidden_states = torch.flatten(hidden_states, 1, -1)
+
+        return hidden_states, fmap
+
+
+#.............................................................................................
+
+class VitsDiscriminator(VitsPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        if config.discriminator_scale_channels is not None:
+            self.discriminators = nn.ModuleList(
+                [VitsHifiGanDiscriminatorScaleResidualBlock(config.discriminator_scale_channels, config.leaky_relu_slope)]
+            )
+        else:
+            self.discriminators = nn.ModuleList([])
+        
+        self.discriminators.extend(
+            [
+                VitsHifiGanDiscriminatorPeriodResidualBlock(
+                    config.discriminator_period_channels,
+                    period,
+                    config.discriminator_kernel_size,
+                    config.discriminator_stride,
+                    config.leaky_relu_slope,
+                )
+                for period in config.discriminator_periods
+            ]
+        )
+
+    def forward(self, hidden_states):
+        fmaps = []
+        discriminated_hidden_states_list = []
+
+        for discriminator in self.discriminators:
+            discriminated_hidden_states, fmap = discriminator(hidden_states)
+            fmaps.append(fmap)
+            discriminated_hidden_states_list.append(discriminated_hidden_states)
+
+        return discriminated_hidden_states_list, fmaps
+
+    def apply_weight_norm(self):
+        for disc in self.discriminators:
+            disc.apply_weight_norm()
+
+    def remove_weight_norm(self):
+        for disc in self.discriminators:
+            disc.remove_weight_norm()
+
+
+#.............................................................................................

VitsModelSplit/duration_predictor.py

@@ -0,0 +1,489 @@
+import math
+import numpy as np
+import torch
+from torch import nn
+
+from .vits_config import VitsConfig
+
+#.............................................
+
+
+def _rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    reverse,
+    tail_bound,
+    min_bin_width,
+    min_bin_height,
+    min_derivative,
+):
+    """
+    This transformation represents a monotonically increasing piecewise rational quadratic function. Unlike the
+    function `_unconstrained_rational_quadratic_spline`, the function behaves the same across the `tail_bound`.
+
+    Args:
+        inputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Second half of the hidden-states input to the Vits convolutional flow module.
+        unnormalized_widths (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            First `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        unnormalized_heights (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            Second `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        unnormalized_derivatives (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            Third `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        reverse (`bool`):
+            Whether the model is being run in reverse mode.
+        tail_bound (`float`):
+            Upper and lower limit bound for the rational quadratic function. Outside of this `tail_bound`, the
+            transform behaves as an identity function.
+        min_bin_width (`float`):
+            Minimum bin value across the width dimension for the piecewise rational quadratic function.
+        min_bin_height (`float`):
+            Minimum bin value across the height dimension for the piecewise rational quadratic function.
+        min_derivative (`float`):
+            Minimum bin value across the derivatives for the piecewise rational quadratic function.
+    Returns:
+        outputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Hidden-states as transformed by the piecewise rational quadratic function.
+        log_abs_det (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Logarithm of the absolute value of the determinants corresponding to the `outputs`.
+    """
+    upper_bound = tail_bound
+    lower_bound = -tail_bound
+    if torch.min(inputs) < lower_bound or torch.max(inputs) > upper_bound:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError(f"Minimal bin width {min_bin_width} too large for the number of bins {num_bins}")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError(f"Minimal bin height {min_bin_height} too large for the number of bins {num_bins}")
+
+    widths = nn.functional.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = nn.functional.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (upper_bound - lower_bound) * cumwidths + lower_bound
+    cumwidths[..., 0] = lower_bound
+    cumwidths[..., -1] = upper_bound
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + nn.functional.softplus(unnormalized_derivatives)
+
+    heights = nn.functional.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = nn.functional.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (upper_bound - lower_bound) * cumheights + lower_bound
+    cumheights[..., 0] = lower_bound
+    cumheights[..., -1] = upper_bound
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    bin_locations = cumheights if reverse else cumwidths
+    bin_locations[..., -1] += 1e-6
+    bin_idx = torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+    bin_idx = bin_idx[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    intermediate1 = input_derivatives + input_derivatives_plus_one - 2 * input_delta
+    if not reverse:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)
+        denominator = input_delta + intermediate1 * theta_one_minus_theta
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+        return outputs, log_abs_det
+    else:
+        # find the roots of a quadratic equation
+        intermediate2 = inputs - input_cumheights
+        intermediate3 = intermediate2 * intermediate1
+        a = input_heights * (input_delta - input_derivatives) + intermediate3
+        b = input_heights * input_derivatives - intermediate3
+        c = -input_delta * intermediate2
+
+        discriminant = b.pow(2) - 4 * a * c
+        if not (discriminant >= 0).all():
+            raise RuntimeError(f"invalid discriminant {discriminant}")
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + intermediate1 * theta_one_minus_theta
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+        return outputs, -log_abs_det
+
+#.............................................
+
+def _unconstrained_rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    reverse=False,
+    tail_bound=5.0,
+    min_bin_width=1e-3,
+    min_bin_height=1e-3,
+    min_derivative=1e-3,
+):
+    """
+    This transformation represents a monotonically increasing piecewise rational quadratic function. Outside of the
+    `tail_bound`, the transform behaves as an identity function.
+
+    Args:
+        inputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Second half of the hidden-states input to the Vits convolutional flow module.
+        unnormalized_widths (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            First `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        unnormalized_heights (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            Second `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        unnormalized_derivatives (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
+            Third `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
+            layer in the convolutional flow module
+        reverse (`bool`, *optional*, defaults to `False`):
+            Whether the model is being run in reverse mode.
+        tail_bound (`float`, *optional* defaults to 5):
+            Upper and lower limit bound for the rational quadratic function. Outside of this `tail_bound`, the
+            transform behaves as an identity function.
+        min_bin_width (`float`, *optional*, defaults to 1e-3):
+            Minimum bin value across the width dimension for the piecewise rational quadratic function.
+        min_bin_height (`float`, *optional*, defaults to 1e-3):
+            Minimum bin value across the height dimension for the piecewise rational quadratic function.
+        min_derivative (`float`, *optional*, defaults to 1e-3):
+            Minimum bin value across the derivatives for the piecewise rational quadratic function.
+    Returns:
+        outputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Hidden-states as transformed by the piecewise rational quadratic function with the `tail_bound` limits
+            applied.
+        log_abs_det (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
+            Logarithm of the absolute value of the determinants corresponding to the `outputs` with the `tail_bound`
+            limits applied.
+    """
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    log_abs_det = torch.zeros_like(inputs)
+    constant = np.log(np.exp(1 - min_derivative) - 1)
+
+    unnormalized_derivatives = nn.functional.pad(unnormalized_derivatives, pad=(1, 1))
+    unnormalized_derivatives[..., 0] = constant
+    unnormalized_derivatives[..., -1] = constant
+
+    outputs[outside_interval_mask] = inputs[outside_interval_mask]
+    log_abs_det[outside_interval_mask] = 0.0
+
+    outputs[inside_interval_mask], log_abs_det[inside_interval_mask] = _rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        reverse=reverse,
+        tail_bound=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+    return outputs, log_abs_det
+
+
+#.............................................................................................
+
+class VitsConvFlow(nn.Module):
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.filter_channels = config.hidden_size
+        self.half_channels = config.depth_separable_channels // 2
+        self.num_bins = config.duration_predictor_flow_bins
+        self.tail_bound = config.duration_predictor_tail_bound
+
+        self.conv_pre = nn.Conv1d(self.half_channels, self.filter_channels, 1)
+        self.conv_dds = VitsDilatedDepthSeparableConv(config)
+        self.conv_proj = nn.Conv1d(self.filter_channels, self.half_channels * (self.num_bins * 3 - 1), 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
+        first_half, second_half = torch.split(inputs, [self.half_channels] * 2, dim=1)
+
+        hidden_states = self.conv_pre(first_half)
+        hidden_states = self.conv_dds(hidden_states, padding_mask, global_conditioning)
+        hidden_states = self.conv_proj(hidden_states) * padding_mask
+
+        batch_size, channels, length = first_half.shape
+        hidden_states = hidden_states.reshape(batch_size, channels, -1, length).permute(0, 1, 3, 2)
+
+        unnormalized_widths = hidden_states[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = hidden_states[..., self.num_bins : 2 * self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_derivatives = hidden_states[..., 2 * self.num_bins :]
+
+        second_half, log_abs_det = _unconstrained_rational_quadratic_spline(
+            second_half,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            reverse=reverse,
+            tail_bound=self.tail_bound,
+        )
+
+        outputs = torch.cat([first_half, second_half], dim=1) * padding_mask
+        if not reverse:
+            log_determinant = torch.sum(log_abs_det * padding_mask, [1, 2])
+            return outputs, log_determinant
+        else:
+            return outputs, None
+
+
+#.............................................................................................
+
+class VitsElementwiseAffine(nn.Module):
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.channels = config.depth_separable_channels
+        self.translate = nn.Parameter(torch.zeros(self.channels, 1))
+        self.log_scale = nn.Parameter(torch.zeros(self.channels, 1))
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
+        if not reverse:
+            outputs = self.translate + torch.exp(self.log_scale) * inputs
+            outputs = outputs * padding_mask
+            log_determinant = torch.sum(self.log_scale * padding_mask, [1, 2])
+            return outputs, log_determinant
+        else:
+            outputs = (inputs - self.translate) * torch.exp(-self.log_scale) * padding_mask
+            return outputs, None
+
+#.............................................................................................
+
+class VitsDilatedDepthSeparableConv(nn.Module):
+    def __init__(self, config: VitsConfig, dropout_rate=0.0):
+        super().__init__()
+        kernel_size = config.duration_predictor_kernel_size
+        channels = config.hidden_size
+        self.num_layers = config.depth_separable_num_layers
+
+        self.dropout = nn.Dropout(dropout_rate)
+        self.convs_dilated = nn.ModuleList()
+        self.convs_pointwise = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(self.num_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_dilated.append(
+                nn.Conv1d(
+                    in_channels=channels,
+                    out_channels=channels,
+                    kernel_size=kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_pointwise.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(nn.LayerNorm(channels))
+            self.norms_2.append(nn.LayerNorm(channels))
+
+    def forward(self, inputs, padding_mask, global_conditioning=None):
+        if global_conditioning is not None:
+            inputs = inputs + global_conditioning
+
+        for i in range(self.num_layers):
+            hidden_states = self.convs_dilated[i](inputs * padding_mask)
+            hidden_states = self.norms_1[i](hidden_states.transpose(1, -1)).transpose(1, -1)
+            hidden_states = nn.functional.gelu(hidden_states)
+            hidden_states = self.convs_pointwise[i](hidden_states)
+            hidden_states = self.norms_2[i](hidden_states.transpose(1, -1)).transpose(1, -1)
+            hidden_states = nn.functional.gelu(hidden_states)
+            hidden_states = self.dropout(hidden_states)
+            inputs = inputs + hidden_states
+
+        return inputs * padding_mask
+
+#.............................................................................................
+
+class VitsStochasticDurationPredictor(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        embed_dim = config.speaker_embedding_size
+        filter_channels = config.hidden_size
+
+        self.conv_pre = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.conv_dds = VitsDilatedDepthSeparableConv(
+            config,
+            dropout_rate=config.duration_predictor_dropout,
+        )
+
+        if embed_dim != 0:
+            self.cond = nn.Conv1d(embed_dim, filter_channels, 1)
+
+        self.flows = nn.ModuleList()
+        self.flows.append(VitsElementwiseAffine(config))
+        for _ in range(config.duration_predictor_num_flows):
+            self.flows.append(VitsConvFlow(config))
+
+        self.post_conv_pre = nn.Conv1d(1, filter_channels, 1)
+        self.post_conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.post_conv_dds = VitsDilatedDepthSeparableConv(
+            config,
+            dropout_rate=config.duration_predictor_dropout,
+        )
+
+        self.post_flows = nn.ModuleList()
+        self.post_flows.append(VitsElementwiseAffine(config))
+        for _ in range(config.duration_predictor_num_flows):
+            self.post_flows.append(VitsConvFlow(config))
+
+        self.filter_channels = filter_channels
+
+    def resize_speaker_embeddings(self, speaker_embedding_size):
+        self.cond = nn.Conv1d(speaker_embedding_size, self.filter_channels, 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, durations=None, reverse=False, noise_scale=1.0):
+        inputs = torch.detach(inputs)
+        inputs = self.conv_pre(inputs)
+
+        if global_conditioning is not None:
+            global_conditioning = torch.detach(global_conditioning)
+            inputs = inputs + self.cond(global_conditioning)
+
+        inputs = self.conv_dds(inputs, padding_mask)
+        inputs = self.conv_proj(inputs) * padding_mask
+
+        if not reverse:
+            hidden_states = self.post_conv_pre(durations)
+            hidden_states = self.post_conv_dds(hidden_states, padding_mask)
+            hidden_states = self.post_conv_proj(hidden_states) * padding_mask
+
+            random_posterior = (
+                torch.randn(durations.size(0), 2, durations.size(2)).to(device=inputs.device, dtype=inputs.dtype)
+                * padding_mask
+            )
+            latents_posterior = random_posterior
+
+            latents_posterior, log_determinant = self.post_flows[0](
+                latents_posterior, padding_mask, global_conditioning=inputs + hidden_states
+            )
+            log_determinant_posterior_sum = log_determinant
+
+            for flow in self.post_flows[1:]:
+                latents_posterior, log_determinant = flow(
+                    latents_posterior, padding_mask, global_conditioning=inputs + hidden_states
+                )
+                latents_posterior = torch.flip(latents_posterior, [1])
+                log_determinant_posterior_sum += log_determinant
+
+            first_half, second_half = torch.split(latents_posterior, [1, 1], dim=1)
+
+            log_determinant_posterior_sum += torch.sum(
+                (nn.functional.logsigmoid(first_half) + nn.functional.logsigmoid(-first_half)) * padding_mask, [1, 2]
+            )
+            logq = (
+                torch.sum(-0.5 * (math.log(2 * math.pi) + (random_posterior**2)) * padding_mask, [1, 2])
+                - log_determinant_posterior_sum
+            )
+
+            first_half = (durations - torch.sigmoid(first_half)) * padding_mask
+            first_half = torch.log(torch.clamp_min(first_half, 1e-5)) * padding_mask
+            log_determinant_sum = torch.sum(-first_half, [1, 2])
+
+            latents = torch.cat([first_half, second_half], dim=1)
+            latents, log_determinant = self.flows[0](latents, padding_mask, global_conditioning=inputs)
+
+            log_determinant_sum += log_determinant
+            for flow in self.flows[1:]:
+                latents, log_determinant = flow(latents, padding_mask, global_conditioning=inputs)
+                latents = torch.flip(latents, [1])
+                log_determinant_sum += log_determinant
+
+            nll = torch.sum(0.5 * (math.log(2 * math.pi) + (latents**2)) * padding_mask, [1, 2]) - log_determinant_sum
+            return nll + logq
+        else:
+            flows = list(reversed(self.flows))
+            flows = flows[:-2] + [flows[-1]]  # remove a useless vflow
+
+            latents = (
+                torch.randn(inputs.size(0), 2, inputs.size(2)).to(device=inputs.device, dtype=inputs.dtype)
+                * noise_scale
+            )
+            for flow in flows:
+                latents = torch.flip(latents, [1])
+                latents, _ = flow(latents, padding_mask, global_conditioning=inputs, reverse=True)
+
+            log_duration, _ = torch.split(latents, [1, 1], dim=1)
+            return log_duration
+
+#.............................................................................................
+
+class VitsDurationPredictor(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        kernel_size = config.duration_predictor_kernel_size
+        filter_channels = config.duration_predictor_filter_channels
+
+        self.dropout = nn.Dropout(config.duration_predictor_dropout)
+        self.conv_1 = nn.Conv1d(config.hidden_size, filter_channels, kernel_size, padding=kernel_size // 2)
+        self.norm_1 = nn.LayerNorm(filter_channels, eps=config.layer_norm_eps)
+        self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)
+        self.norm_2 = nn.LayerNorm(filter_channels, eps=config.layer_norm_eps)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if config.speaker_embedding_size != 0:
+            self.cond = nn.Conv1d(config.speaker_embedding_size, config.hidden_size, 1)
+
+        self.hidden_size = config.hidden_size
+
+    def resize_speaker_embeddings(self, speaker_embedding_size):
+        self.cond = nn.Conv1d(speaker_embedding_size, self.hidden_size, 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None):
+        inputs = torch.detach(inputs)
+
+        if global_conditioning is not None:
+            global_conditioning = torch.detach(global_conditioning)
+            inputs = inputs + self.cond(global_conditioning)
+
+        inputs = self.conv_1(inputs * padding_mask)
+        inputs = torch.relu(inputs)
+        inputs = self.norm_1(inputs.transpose(1, -1)).transpose(1, -1)
+        inputs = self.dropout(inputs)
+
+        inputs = self.conv_2(inputs * padding_mask)
+        inputs = torch.relu(inputs)
+        inputs = self.norm_2(inputs.transpose(1, -1)).transpose(1, -1)
+        inputs = self.dropout(inputs)
+
+        inputs = self.proj(inputs * padding_mask)
+        return inputs * padding_mask
+
+#.............................................................................................

VitsModelSplit/encoder.py

@@ -0,0 +1,407 @@
+import math
+from typing import Optional, Tuple, Union
+import numpy as np
+import torch
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+from transformers.modeling_outputs import BaseModelOutput  
+
+from .vits_config import VitsConfig
+from .vits_output import VitsTextEncoderOutput
+
+
+#....................................................
+
+
+
+
+
+class VitsFeedForward(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.conv_1 = nn.Conv1d(config.hidden_size, config.ffn_dim, config.ffn_kernel_size)
+        self.conv_2 = nn.Conv1d(config.ffn_dim, config.hidden_size, config.ffn_kernel_size)
+        self.dropout = nn.Dropout(config.activation_dropout)
+
+        if isinstance(config.hidden_act, str):
+            self.act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.act_fn = config.hidden_act
+
+        if config.ffn_kernel_size > 1:
+            pad_left = (config.ffn_kernel_size - 1) // 2
+            pad_right = config.ffn_kernel_size // 2
+            self.padding = [pad_left, pad_right, 0, 0, 0, 0]
+        else:
+            self.padding = None
+
+    def forward(self, hidden_states, padding_mask):
+        hidden_states = hidden_states.permute(0, 2, 1)
+        padding_mask = padding_mask.permute(0, 2, 1)
+
+        hidden_states = hidden_states * padding_mask
+        if self.padding is not None:
+            hidden_states = nn.functional.pad(hidden_states, self.padding)
+
+        hidden_states = self.conv_1(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        hidden_states = hidden_states * padding_mask
+        if self.padding is not None:
+            hidden_states = nn.functional.pad(hidden_states, self.padding)
+
+        hidden_states = self.conv_2(hidden_states)
+        hidden_states = hidden_states * padding_mask
+
+        hidden_states = hidden_states.permute(0, 2, 1)
+        return hidden_states
+
+
+#.............................................................................................
+
+class VitsAttention(nn.Module):
+    """Multi-headed attention with relative positional representation."""
+
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.dropout = config.attention_dropout
+        self.window_size = config.window_size
+
+        self.head_dim = self.embed_dim // self.num_heads
+        self.scaling = self.head_dim**-0.5
+
+        if (self.head_dim * self.num_heads) != self.embed_dim:
+            raise ValueError(
+                f"hidden_size must be divisible by num_attention_heads (got `hidden_size`: {self.embed_dim}"
+                f" and `num_attention_heads`: {self.num_heads})."
+            )
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
+
+        if self.window_size:
+            self.emb_rel_k = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.head_dim) * self.scaling)
+            self.emb_rel_v = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.head_dim) * self.scaling)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Input shape: Batch x Time x Channel"""
+
+        # if key_value_states are provided this layer is used as a cross-attention layer
+        # for the decoder
+
+        bsz, tgt_len, _ = hidden_states.size()
+
+        # get query proj
+        query_states = self.q_proj(hidden_states) * self.scaling
+
+        # self_attention
+        key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+        value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+
+        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
+        key_states = key_states.view(*proj_shape)
+        value_states = value_states.view(*proj_shape)
+
+        src_len = key_states.size(1)
+        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
+
+        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if self.window_size is not None:
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, src_len)
+            relative_logits = torch.matmul(query_states, key_relative_embeddings.transpose(-2, -1))
+            rel_pos_bias = self._relative_position_to_absolute_position(relative_logits)
+            attn_weights += rel_pos_bias
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+        if layer_head_mask is not None:
+            if layer_head_mask.size() != (self.num_heads,):
+                raise ValueError(
+                    f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
+                    f" {layer_head_mask.size()}"
+                )
+            attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if output_attentions:
+            # this operation is a bit awkward, but it's required to
+            # make sure that attn_weights keeps its gradient.
+            # In order to do so, attn_weights have to be reshaped
+            # twice and have to be reused in the following
+            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
+        else:
+            attn_weights_reshaped = None
+
+        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+
+        attn_output = torch.bmm(attn_probs, value_states)
+
+        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        if self.window_size is not None:
+            value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, src_len)
+            relative_weights = self._absolute_position_to_relative_position(attn_probs)
+            rel_pos_bias = torch.matmul(relative_weights, value_relative_embeddings)
+            attn_output += rel_pos_bias
+
+        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        attn_output = attn_output.transpose(1, 2)
+
+        # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
+        # partitioned aross GPUs when using tensor-parallelism.
+        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights_reshaped
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        pad_length = max(length - (self.window_size + 1), 0)
+        if pad_length > 0:
+            relative_embeddings = nn.functional.pad(relative_embeddings, [0, 0, pad_length, pad_length, 0, 0])
+
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        return relative_embeddings[:, slice_start_position:slice_end_position]
+
+    def _relative_position_to_absolute_position(self, x):
+        batch_heads, length, _ = x.size()
+
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = nn.functional.pad(x, [0, 1, 0, 0, 0, 0])
+
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch_heads, length * 2 * length])
+        x_flat = nn.functional.pad(x_flat, [0, length - 1, 0, 0])
+
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch_heads, length + 1, 2 * length - 1])
+        x_final = x_final[:, :length, length - 1 :]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        batch_heads, length, _ = x.size()
+
+        # Pad along column
+        x = nn.functional.pad(x, [0, length - 1, 0, 0, 0, 0])
+        x_flat = x.view([batch_heads, length**2 + length * (length - 1)])
+
+        # Add 0's in the beginning that will skew the elements after reshape
+        x_flat = nn.functional.pad(x_flat, [length, 0, 0, 0])
+        x_final = x_flat.view([batch_heads, length, 2 * length])[:, :, 1:]
+        return x_final
+
+
+#.............................................................................................
+
+class VitsEncoderLayer(nn.Module):
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.attention = VitsAttention(config)
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = VitsFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        padding_mask: torch.FloatTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ):
+        residual = hidden_states
+        hidden_states, attn_weights = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.layer_norm(residual + hidden_states)
+
+        residual = hidden_states
+        hidden_states = self.feed_forward(hidden_states, padding_mask)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.final_layer_norm(residual + hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+#.............................................................................................
+
+class VitsEncoder(nn.Module):
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([VitsEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+        self.layerdrop = config.layerdrop
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        padding_mask: torch.FloatTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+
+        # expand attention_mask
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+        hidden_states = hidden_states * padding_mask
+
+        deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled()
+
+        for encoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
+            dropout_probability = np.random.uniform(0, 1)
+
+            skip_the_layer = self.training and (dropout_probability < self.layerdrop)
+            if not skip_the_layer or deepspeed_zero3_is_enabled:
+                # under deepspeed zero3 all gpus must run in sync
+                if self.gradient_checkpointing and self.training:
+                    layer_outputs = self._gradient_checkpointing_func(
+                        encoder_layer.__call__,
+                        hidden_states,
+                        padding_mask,
+                        attention_mask,
+                        output_attentions,
+                    )
+                else:
+                    layer_outputs = encoder_layer(
+                        hidden_states,
+                        attention_mask=attention_mask,
+                        padding_mask=padding_mask,
+                        output_attentions=output_attentions,
+                    )
+                hidden_states = layer_outputs[0]
+
+            if skip_the_layer:
+                layer_outputs = (None, None)
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        hidden_states = hidden_states * padding_mask
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+        
+#.............................................................................................
+
+class VitsTextEncoder(nn.Module):
+    """
+    Transformer encoder that uses relative positional representation instead of absolute positional encoding.
+    """
+
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        
+        self.encoder = VitsEncoder(config)
+        self.project = nn.Conv1d(config.hidden_size, config.flow_size * 2, kernel_size=1)
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        padding_mask: torch.FloatTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = True,
+    ) -> Union[Tuple[torch.Tensor], VitsTextEncoderOutput]:
+        hidden_states = self.embed_tokens(input_ids) * math.sqrt(self.config.hidden_size)
+
+        encoder_outputs = self.encoder(
+            hidden_states=hidden_states,
+            padding_mask=padding_mask,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0] if not return_dict else encoder_outputs.last_hidden_state
+
+        stats = self.project(last_hidden_state.transpose(1, 2)).transpose(1, 2) * padding_mask
+        prior_means, prior_log_variances = torch.split(stats, self.config.flow_size, dim=2)
+
+        if not return_dict:
+            outputs = (last_hidden_state, prior_means, prior_log_variances) + encoder_outputs[1:]
+            return outputs
+
+        return VitsTextEncoderOutput(
+            last_hidden_state=last_hidden_state,
+            prior_means=prior_means,
+            prior_log_variances=prior_log_variances,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+        
+#.............................................................................................

VitsModelSplit/feature_extraction.py

@@ -0,0 +1,280 @@
+"""
+Feature extractor class for Vits
+"""
+import copy
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
+from transformers import is_torch_available
+from transformers.audio_utils import mel_filter_bank
+from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils import TensorType, logging
+
+
+MAX_WAV_VALUE = 32768.0
+
+if is_torch_available():
+    import torch
+
+logger = logging.get_logger(__name__)
+
+
+class VitsFeatureExtractor(SequenceFeatureExtractor):
+    r"""
+    Constructs a Vits feature extractor.
+
+    This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains
+    most of the main methods. Users should refer to this superclass for more information regarding those methods.
+
+    This class extracts `Short Time Fourier Transform` from raw speech using a custom numpy implementation which should
+    match pytorch's `torch.stft`.
+
+    Args:
+        feature_size (`int`, defaults to 80):
+            The feature dimension of the extracted features.
+        sampling_rate (`int`, defaults to 22050):
+            The sampling rate at which the audio files should be digitalized expressed in hertz (Hz).
+        hop_length (`int`, defaults to 256):
+            Length of the overlaping windows for the STFT used to obtain the Mel Frequency coefficients.
+        n_fft (`int`, defaults to 1024):
+            Size of the Fourier transform.
+        padding_value (`float`, *optional*, defaults to 0.0):
+            Padding value used to pad the audio. Should correspond to silences.
+        return_attention_mask (`bool`, *optional*, defaults to `False`):
+            Whether to return the attention mask.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            <Tip>
+
+            For Vits finetuning, `attention_mask` should always be passed for batched inference, to avoid subtle bugs.
+
+            </Tip>
+
+        max_wav_value (`float`, defaults to 32768.0):
+            Maximum wav value. Used to normalize the input waveforms if `do_normalize=True` in the forward pass of this
+            feature extractor.
+    """
+
+    model_input_names = ["input_features"]
+
+    def __init__(
+        self,
+        feature_size=80,
+        sampling_rate=16000,
+        hop_length=256,
+        n_fft=1024,
+        padding_value=0.0,
+        return_attention_mask=False,  # pad inputs to max length with silence token (zero) and no attention mask,
+        max_wav_value=32768.0,
+        **kwargs,
+    ):
+        super().__init__(
+            feature_size=feature_size,
+            sampling_rate=sampling_rate,
+            padding_value=padding_value,
+            return_attention_mask=return_attention_mask,
+            **kwargs,
+        )
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.sampling_rate = sampling_rate
+        self.mel_filters = mel_filter_bank(
+            num_frequency_bins=1 + n_fft // 2,
+            num_mel_filters=feature_size,
+            min_frequency=0.0,
+            max_frequency=sampling_rate // 2,
+            sampling_rate=sampling_rate,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+        self.max_wav_value = max_wav_value
+
+    def _torch_extract_fbank_features(self, waveform: np.array) -> Tuple[torch.Tensor]:
+        """
+        Compute the log-mel spectrogram of the provided audio using the PyTorch STFT implementation.
+        """
+        if len(waveform.shape) == 1:
+            waveform = waveform.unsqueeze(0)
+
+        waveform = torch.nn.functional.pad(
+            waveform,
+            (int((self.n_fft - self.hop_length) / 2), int((self.n_fft - self.hop_length) / 2)),
+            mode="reflect",
+        )
+
+        window = torch.hann_window(self.n_fft).to(waveform.device)
+        stft = torch.stft(
+            waveform,
+            self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.n_fft,
+            window=window,
+            center=False,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=False,
+        )
+        magnitudes = torch.sqrt(stft.pow(2).sum(-1) + 1e-6)
+
+        mel_filters = torch.from_numpy(self.mel_filters).type(torch.float32).to(waveform.device)
+        mel_spec = mel_filters.T @ magnitudes
+
+        log_spec = torch.clamp(mel_spec, min=1e-5).log()
+        return magnitudes, log_spec
+
+    def __call__(
+        self,
+        raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],
+        truncation: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: Optional[bool] = True,
+        padding: Optional[str] = True,
+        max_length: Optional[int] = None,
+        sampling_rate: Optional[int] = None,
+        do_normalize: Optional[bool] = None,
+        **kwargs,
+    ) -> BatchFeature:
+        """
+        Main method to featurize and prepare for the model one or several sequence(s).
+
+        Args:
+            raw_speech (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`):
+                The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float
+                values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not
+                stereo, i.e. single float per timestep.
+            truncation (`bool`, *optional*, default to `False`):
+                Activates truncation to cut input sequences longer than *max_length* to *max_length*.
+            pad_to_multiple_of (`int`, *optional*, defaults to None):
+                If set will pad the sequence to a multiple of the provided value.
+
+                This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
+                `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors instead of list of python integers. Acceptable values are:
+
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return Numpy `np.ndarray` objects.
+            return_attention_mask (`bool`, *optional*, defaults to `True`):
+                Whether to return the attention mask. If left to the default, will return the attention mask according
+                to the specific feature_extractor's default.
+
+                [What are attention masks?](../glossary#attention-mask)
+
+                <Tip>
+
+                For Vits finetuning, `attention_mask` should always be passed for batched inference, to avoid subtle
+                bugs.
+
+                </Tip>
+
+            padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
+                Select a strategy to pad the returned sequences (according to the model's padding side and padding
+                index) among:
+
+                - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+                  sequence if provided).
+                - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
+                  acceptable input length for the model if that argument is not provided.
+                - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
+                  lengths).
+            max_length (`int`, *optional*):
+                Maximum length of the returned list and optionally padding length (see above).
+            sampling_rate (`int`, *optional*):
+                The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass
+                `sampling_rate` at the forward call to prevent silent errors and allow automatic speech recognition
+                pipeline.
+            do_normalize (`bool`, *optional*):
+                Whether or not to divide the input waveform by `self.max_wav_value`.
+        """
+
+        if sampling_rate is not None:
+            if sampling_rate != self.sampling_rate:
+                raise ValueError(
+                    f"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"
+                    f" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"
+                    f" was sampled with {self.sampling_rate} and not {sampling_rate}."
+                )
+        else:
+            logger.warning(
+                "It is strongly recommended to pass the `sampling_rate` argument to this function. "
+                "Failing to do so can result in silent errors that might be hard to debug."
+            )
+
+        is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
+        if is_batched_numpy and len(raw_speech.shape) > 2:
+            raise ValueError(f"Only mono-channel audio is supported for input to {self}")
+        is_batched = is_batched_numpy or (
+            isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
+        )
+
+        if is_batched:
+            raw_speech = [np.asarray([speech], dtype=np.float32).T for speech in raw_speech]
+        elif not is_batched and not isinstance(raw_speech, np.ndarray):
+            raw_speech = np.asarray(raw_speech, dtype=np.float32)
+        elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):
+            raw_speech = raw_speech.astype(np.float32)
+
+        # always return batch
+        if not is_batched:
+            raw_speech = [np.asarray([raw_speech]).T]
+
+        if self.max_wav_value is not None and do_normalize:
+            raw_speech = [
+                speech if self.max_wav_value is None else speech / self.max_wav_value for speech in raw_speech
+            ]
+
+        batched_speech = BatchFeature({"input_features": raw_speech})
+
+        # convert into correct format for padding
+        padded_inputs = self.pad(
+            batched_speech,
+            padding=padding,
+            max_length=max_length,
+            truncation=truncation,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_attention_mask=return_attention_mask or do_normalize,
+            return_tensors="pt",
+        )
+
+        # make sure list is in array format
+        if isinstance(padded_inputs.get("input_features"),list):
+            input_features = torch.tensor(padded_inputs.get("input_features")).transpose(1, 2).transpose(0, 1)
+        else:
+            input_features = padded_inputs.get("input_features").clone().detach().transpose(1, 2).transpose(0, 1)
+
+
+        input_features = self._torch_extract_fbank_features(input_features[0])
+
+        mel_scaled_input_features = input_features[1]
+        input_features = input_features[0]
+
+        padded_inputs["input_features"] = input_features
+        padded_inputs["mel_scaled_input_features"] = mel_scaled_input_features
+
+        if return_attention_mask:
+            # rescale from sample (48000) to feature (3000)
+            padded_inputs["attention_mask"] = padded_inputs["attention_mask"][:, :: self.hop_length]
+
+        if return_tensors is not None:
+            padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
+
+        return padded_inputs
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Serializes this instance to a Python dictionary.
+
+        Returns:
+            `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+        """
+        output = copy.deepcopy(self.__dict__)
+        output["feature_extractor_type"] = self.__class__.__name__
+        if "mel_filters" in output:
+            del output["mel_filters"]
+        return output

VitsModelSplit/finetune_config_ara.json

@@ -0,0 +1,55 @@
+{
+    "project_name": "vits_ara",
+    "push_to_hub": false,
+    "hub_model_id": "ara_tts_finetuning/ara_tts_finetuning",
+    "overwrite_output_dir": true,
+    "output_dir": "./output",
+
+    "dataset_name": "./dataset/",      
+    "dataset_config_name": "welsh_female", 
+    "audio_column_name": "audio", 
+    "text_column_name":"text",
+    "train_split_name": "train",
+    "eval_split_name": "eval",
+  
+    "override_speaker_embeddings": false,
+    "filter_on_speaker_id": 5223,
+
+    
+    "max_duration_in_seconds": 20,
+    "min_duration_in_seconds": 1.0,
+    "max_tokens_length": 500,
+
+    "model_name_or_path": "facebook/mms-tts-ara",
+
+    "full_generation_sample_text": "اوريه و اخليه يعرف هو حاط نفسه في مواجهة مع مين",
+    "preprocessing_num_workers": 4,
+
+    "do_train": true,
+    "num_train_epochs": 300,
+    "gradient_accumulation_steps": 1,
+    "per_device_train_batch_size": 10,
+    "learning_rate": 2e-5,
+    "adam_beta1": 0.8,
+    "adam_beta2": 0.99,
+    "warmup_ratio": 0.01,
+    "d_learning_rate": 2e-5,
+    "d_adam_beta1": 0.7,
+    "d_adam_beta2": 0.99,
+
+    "do_eval": true, 
+    "eval_steps": 10,
+    "per_device_eval_batch_size": 10,
+    "max_eval_samples": 2, 
+    "do_step_schedule_per_epoch": true,
+
+    "weight_disc": 3,
+    "weight_fmaps": 1,
+    "weight_gen": 1,
+    "weight_kl": 1.5,
+    "weight_duration": 1,
+    "weight_mel": 35,
+
+    "fp16": false,
+    "seed": 456
+}

VitsModelSplit/flow.py

@@ -0,0 +1,190 @@
+
+import torch
+from torch import nn
+from typing import  Optional
+from .vits_config import VitsConfig
+#.............................................
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, num_channels):
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :num_channels, :])
+    s_act = torch.sigmoid(in_act[:, num_channels:, :])
+    acts = t_act * s_act
+    return acts
+
+
+
+#.............................................
+
+class VitsWaveNet(torch.nn.Module):
+    def __init__(self, config: VitsConfig, num_layers: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_layers = num_layers
+        self.speaker_embedding_size = config.speaker_embedding_size
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.dropout = nn.Dropout(config.wavenet_dropout)
+
+        if hasattr(nn.utils.parametrizations, "weight_norm"):
+            weight_norm = nn.utils.parametrizations.weight_norm
+        else:
+            weight_norm = nn.utils.weight_norm
+
+        if config.speaker_embedding_size != 0:
+            cond_layer = torch.nn.Conv1d(config.speaker_embedding_size, 2 * config.hidden_size * num_layers, 1)
+            self.cond_layer = weight_norm(cond_layer, name="weight")
+
+        for i in range(num_layers):
+            dilation = config.wavenet_dilation_rate**i
+            padding = (config.wavenet_kernel_size * dilation - dilation) // 2
+            in_layer = torch.nn.Conv1d(
+                in_channels=config.hidden_size,
+                out_channels=2 * config.hidden_size,
+                kernel_size=config.wavenet_kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < num_layers - 1:
+                res_skip_channels = 2 * config.hidden_size
+            else:
+                res_skip_channels = config.hidden_size
+
+            res_skip_layer = torch.nn.Conv1d(config.hidden_size, res_skip_channels, 1)
+            res_skip_layer = weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None):
+        outputs = torch.zeros_like(inputs)
+        num_channels_tensor = torch.IntTensor([self.hidden_size])
+
+        if global_conditioning is not None:
+            global_conditioning = self.cond_layer(global_conditioning)
+
+        for i in range(self.num_layers):
+            hidden_states = self.in_layers[i](inputs)
+
+            if global_conditioning is not None:
+                cond_offset = i * 2 * self.hidden_size
+                global_states = global_conditioning[:, cond_offset : cond_offset + 2 * self.hidden_size, :]
+            else:
+                global_states = torch.zeros_like(hidden_states)
+
+            acts = fused_add_tanh_sigmoid_multiply(hidden_states, global_states, num_channels_tensor[0])
+            acts = self.dropout(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.num_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_size, :]
+                inputs = (inputs + res_acts) * padding_mask
+                outputs = outputs + res_skip_acts[:, self.hidden_size :, :]
+            else:
+                outputs = outputs + res_skip_acts
+
+        return outputs * padding_mask
+
+    def remove_weight_norm(self):
+        if self.speaker_embedding_size != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for layer in self.in_layers:
+            torch.nn.utils.remove_weight_norm(layer)
+        for layer in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(layer)
+
+    def apply_weight_norm(self):
+        if hasattr(nn.utils.parametrizations, "weight_norm"):
+            weight_norm = nn.utils.parametrizations.weight_norm
+        else:
+            weight_norm = nn.utils.weight_norm
+
+        if self.speaker_embedding_size != 0:
+            weight_norm(self.cond_layer)
+        for layer in self.in_layers:
+            weight_norm(layer)
+        for layer in self.res_skip_layers:
+            weight_norm(layer)
+
+
+#.............................................................................................
+
+class VitsResidualCouplingLayer(nn.Module):
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.half_channels = config.flow_size // 2
+
+        self.conv_pre = nn.Conv1d(self.half_channels, config.hidden_size, 1)
+        self.wavenet = VitsWaveNet(config, num_layers=config.prior_encoder_num_wavenet_layers)
+        self.conv_post = nn.Conv1d(config.hidden_size, self.half_channels, 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
+        first_half, second_half = torch.split(inputs, [self.half_channels] * 2, dim=1)
+        hidden_states = self.conv_pre(first_half) * padding_mask
+        hidden_states = self.wavenet(hidden_states, padding_mask, global_conditioning)
+        mean = self.conv_post(hidden_states) * padding_mask
+        log_stddev = torch.zeros_like(mean)
+
+        if not reverse:
+            second_half = mean + second_half * torch.exp(log_stddev) * padding_mask
+            outputs = torch.cat([first_half, second_half], dim=1)
+            log_determinant = torch.sum(log_stddev, [1, 2])
+            return outputs, log_determinant
+        else:
+            second_half = (second_half - mean) * torch.exp(-log_stddev) * padding_mask
+            outputs = torch.cat([first_half, second_half], dim=1)
+            return outputs, None
+
+    def apply_weight_norm(self):
+        nn.utils.weight_norm(self.conv_pre)
+        self.wavenet.apply_weight_norm()
+        nn.utils.weight_norm(self.conv_post)
+
+    def remove_weight_norm(self):
+        nn.utils.remove_weight_norm(self.conv_pre)
+        self.wavenet.remove_weight_norm()
+        nn.utils.remove_weight_norm(self.conv_post)
+
+
+
+#.............................................................................................
+
+class VitsResidualCouplingBlock(nn.Module):
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.flows = nn.ModuleList()
+        for _ in range(config.prior_encoder_num_flows):
+            self.flows.append(VitsResidualCouplingLayer(config))
+
+    def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                inputs, _ = flow(inputs, padding_mask, global_conditioning)
+                inputs = torch.flip(inputs, [1])
+        else:
+            for flow in reversed(self.flows):
+                inputs = torch.flip(inputs, [1])
+                inputs, _ = flow(inputs, padding_mask, global_conditioning, reverse=True)
+        return inputs
+
+    def apply_weight_norm(self):
+        for flow in self.flows:
+            flow.apply_weight_norm()
+
+    def remove_weight_norm(self):
+        for flow in self.flows:
+            flow.remove_weight_norm()
+
+    def resize_speaker_embeddings(self, speaker_embedding_size: Optional[int] = None):
+        for flow in self.flows:
+            flow.wavenet.speaker_embedding_size = speaker_embedding_size
+            hidden_size = flow.wavenet.hidden_size
+            num_layers = flow.wavenet.num_layers
+
+            cond_layer = torch.nn.Conv1d(speaker_embedding_size, 2 * hidden_size * num_layers, 1)
+            flow.wavenet.cond_layer = nn.utils.weight_norm(cond_layer, name="weight")
+

VitsModelSplit/mk

@@ -0,0 +1 @@
+

VitsModelSplit/monotonic_align/__init__.py

@@ -0,0 +1,19 @@
+import numpy as np
+import torch
+from .monotonic_align.core import maximum_path_c
+
+
+def maximum_path(neg_cent, mask):
+  """ Cython optimized version.
+  neg_cent: [b, t_t, t_s]
+  mask: [b, t_t, t_s]
+  """
+  device = neg_cent.device
+  dtype = neg_cent.dtype
+  neg_cent = neg_cent.data.cpu().numpy().astype(np.float32)
+  path = np.zeros(neg_cent.shape, dtype=np.int32)
+
+  t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(np.int32)
+  t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(np.int32)
+  maximum_path_c(path, neg_cent, t_t_max, t_s_max)
+  return torch.from_numpy(path).to(device=device, dtype=dtype)

VitsModelSplit/monotonic_align/core.pyx

@@ -0,0 +1,42 @@
+cimport cython
+from cython.parallel import prange
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cdef void maximum_path_each(int[:,::1] path, float[:,::1] value, int t_y, int t_x, float max_neg_val=-1e9) nogil:
+  cdef int x
+  cdef int y
+  cdef float v_prev
+  cdef float v_cur
+  cdef float tmp
+  cdef int index = t_x - 1
+
+  for y in range(t_y):
+    for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+      if x == y:
+        v_cur = max_neg_val
+      else:
+        v_cur = value[y-1, x]
+      if x == 0:
+        if y == 0:
+          v_prev = 0.
+        else:
+          v_prev = max_neg_val
+      else:
+        v_prev = value[y-1, x-1]
+      value[y, x] += max(v_prev, v_cur)
+
+  for y in range(t_y - 1, -1, -1):
+    path[y, index] = 1
+    if index != 0 and (index == y or value[y-1, index] < value[y-1, index-1]):
+      index = index - 1
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cpdef void maximum_path_c(int[:,:,::1] paths, float[:,:,::1] values, int[::1] t_ys, int[::1] t_xs) nogil:
+  cdef int b = paths.shape[0]
+  cdef int i
+  for i in prange(b, nogil=True):
+    maximum_path_each(paths[i], values[i], t_ys[i], t_xs[i])

VitsModelSplit/monotonic_align/data

@@ -0,0 +1 @@
+

VitsModelSplit/monotonic_align/setup.py

@@ -0,0 +1,9 @@
+from distutils.core import setup
+from Cython.Build import cythonize
+import numpy
+
+setup(
+  name = 'monotonic_align',
+  ext_modules = cythonize("core.pyx"),
+  include_dirs=[numpy.get_include()]
+)

VitsModelSplit/plot.py

@@ -0,0 +1,63 @@
+import logging
+import matplotlib
+
+
+matplotlib.use("Agg")
+
+MATPLOTLIB_FLAG = False
+
+
+def plot_spectrogram_to_numpy(spectrogram):
+    global MATPLOTLIB_FLAG
+    if not MATPLOTLIB_FLAG:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        MATPLOTLIB_FLAG = True
+        mpl_logger = logging.getLogger("matplotlib")
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+    import numpy as np
+
+    fig, ax = plt.subplots(figsize=(10, 2))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+    
+    plt.colorbar(im, ax=ax)
+    plt.xlabel("Frames")
+    plt.ylabel("Channels")
+    plt.tight_layout()
+    fig.canvas.draw()
+    
+    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def plot_alignment_to_numpy(alignment, info=None):
+    global MATPLOTLIB_FLAG
+    if not MATPLOTLIB_FLAG:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        MATPLOTLIB_FLAG = True
+        mpl_logger = logging.getLogger("matplotlib")
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+    import numpy as np
+
+    fig, ax = plt.subplots(figsize=(6, 4))
+    im = ax.imshow(alignment.transpose(), aspect="auto", origin="lower", interpolation="none")
+    fig.colorbar(im, ax=ax)
+    xlabel = "Decoder timestep"
+    if info is not None:
+        xlabel += "\n\n" + info
+    plt.xlabel(xlabel)
+    plt.ylabel("Encoder timestep")
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data

VitsModelSplit/posterior_encoder.py

@@ -0,0 +1,50 @@
+import math
+from typing import Optional
+import torch
+from torch import nn
+from .vits_config import VitsConfig
+from .flow import VitsWaveNet
+
+#.............................................
+
+
+
+class VitsPosteriorEncoder(nn.Module):
+    def __init__(self, config: VitsConfig):
+        super().__init__()
+        self.out_channels = config.flow_size
+
+        self.conv_pre = nn.Conv1d(config.spectrogram_bins, config.hidden_size, 1)
+        self.wavenet = VitsWaveNet(config, num_layers=config.posterior_encoder_num_wavenet_layers)
+        self.conv_proj = nn.Conv1d(config.hidden_size, self.out_channels * 2, 1)
+
+    def forward(self, inputs, padding_mask, global_conditioning=None):
+        inputs = self.conv_pre(inputs) * padding_mask
+        inputs = self.wavenet(inputs, padding_mask, global_conditioning)
+        stats = self.conv_proj(inputs) * padding_mask
+        mean, log_stddev = torch.split(stats, self.out_channels, dim=1)
+        sampled = (mean + torch.randn_like(mean) * torch.exp(log_stddev)) * padding_mask
+        return sampled, mean, log_stddev
+
+    def apply_weight_norm(self):
+        self.wavenet.apply_weight_norm()
+
+    def remove_weight_norm(self):
+        self.wavenet.remove_weight_norm()
+
+    def resize_speaker_embeddings(self, speaker_embedding_size: Optional[int] = None):
+        self.wavenet.speaker_embedding_size = speaker_embedding_size
+        hidden_size = self.wavenet.hidden_size
+        num_layers = self.wavenet.num_layers
+
+        cond_layer = torch.nn.Conv1d(speaker_embedding_size, 2 * hidden_size * num_layers, 1)
+        self.wavenet.cond_layer = nn.utils.weight_norm(cond_layer, name="weight")
+        nn.init.kaiming_normal_(self.wavenet.cond_layer.weight)
+        if self.wavenet.cond_layer.bias is not None:
+            k = math.sqrt(
+                self.wavenet.cond_layer.groups
+                / (self.wavenet.cond_layer.in_channels * self.wavenet.cond_layer.kernel_size[0])
+            )
+            nn.init.uniform_(self.wavenet.cond_layer.bias, a=-k, b=k)
+
+#.............................................................................................

VitsModelSplit/requirements.txt

@@ -0,0 +1,10 @@
+Cython==0.29.21
+librosa==0.8.0
+matplotlib==3.3.1
+numpy==1.18.5
+phonemizer==2.2.1
+scipy==1.5.2
+tensorboard==2.3.0
+torch==1.6.0
+torchvision==0.7.0
+Unidecode==1.1.1

VitsModelSplit/vits_config.py

@@ -0,0 +1,162 @@
+import math
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_utils import PreTrainedModel
+from torch import nn
+
+#.............................................
+
+
+
+class VitsConfig(PretrainedConfig):
+    model_type = "vits"
+    
+    def __init__(
+        self,
+        vocab_size=38,
+        hidden_size=192,
+        num_hidden_layers=6,
+        num_attention_heads=2,
+        window_size=4,
+        use_bias=True,
+        ffn_dim=768,
+        layerdrop=0.1,
+        ffn_kernel_size=3,
+        flow_size=192,
+        spectrogram_bins=513,
+        hidden_act="relu",
+        hidden_dropout=0.1,
+        attention_dropout=0.1,
+        activation_dropout=0.1,
+        initializer_range=0.02,
+        layer_norm_eps=1e-5,
+        use_stochastic_duration_prediction=True,
+        num_speakers=1,
+        speaker_embedding_size=0,
+        upsample_initial_channel=512,
+        upsample_rates=[8, 8, 2, 2],
+        upsample_kernel_sizes=[16, 16, 4, 4],
+        resblock_kernel_sizes=[3, 7, 11],
+        resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        leaky_relu_slope=0.1,
+        depth_separable_channels=2,
+        depth_separable_num_layers=3,
+        duration_predictor_flow_bins=10,
+        duration_predictor_tail_bound=5.0,
+        duration_predictor_kernel_size=3,
+        duration_predictor_dropout=0.5,
+        duration_predictor_num_flows=4,
+        duration_predictor_filter_channels=256,
+        prior_encoder_num_flows=4,
+        prior_encoder_num_wavenet_layers=4,
+        posterior_encoder_num_wavenet_layers=16,
+        wavenet_kernel_size=5,
+        wavenet_dilation_rate=1,
+        wavenet_dropout=0.0,
+        speaking_rate=1.0,
+        noise_scale=0.667,
+        noise_scale_duration=0.8,
+        sampling_rate=16_000,
+        discriminator_kernel_size=5,
+        discriminator_stride=3,
+        discriminator_periods=[2, 3, 5, 7, 11],
+        discriminator_period_channels=[1, 32, 128, 512, 1024],
+        discriminator_scale_channels=[1, 16, 64, 256, 1024],
+        segment_size=8192,
+        hop_length=256,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.window_size = window_size
+        self.use_bias = use_bias
+        self.ffn_dim = ffn_dim
+        self.layerdrop = layerdrop
+        self.ffn_kernel_size = ffn_kernel_size
+        self.flow_size = flow_size
+        self.spectrogram_bins = spectrogram_bins
+        self.hidden_act = hidden_act
+        self.hidden_dropout = hidden_dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.use_stochastic_duration_prediction = use_stochastic_duration_prediction
+        self.num_speakers = num_speakers
+        self.speaker_embedding_size = speaker_embedding_size
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_rates = upsample_rates
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.leaky_relu_slope = leaky_relu_slope
+        self.depth_separable_channels = depth_separable_channels
+        self.depth_separable_num_layers = depth_separable_num_layers
+        self.duration_predictor_flow_bins = duration_predictor_flow_bins
+        self.duration_predictor_tail_bound = duration_predictor_tail_bound
+        self.duration_predictor_kernel_size = duration_predictor_kernel_size
+        self.duration_predictor_dropout = duration_predictor_dropout
+        self.duration_predictor_num_flows = duration_predictor_num_flows
+        self.duration_predictor_filter_channels = duration_predictor_filter_channels
+        self.prior_encoder_num_flows = prior_encoder_num_flows
+        self.prior_encoder_num_wavenet_layers = prior_encoder_num_wavenet_layers
+        self.posterior_encoder_num_wavenet_layers = posterior_encoder_num_wavenet_layers
+        self.wavenet_kernel_size = wavenet_kernel_size
+        self.wavenet_dilation_rate = wavenet_dilation_rate
+        self.wavenet_dropout = wavenet_dropout
+        self.speaking_rate = speaking_rate
+        self.noise_scale = noise_scale
+        self.noise_scale_duration = noise_scale_duration
+        self.sampling_rate = sampling_rate
+
+        # used for training
+        self.discriminator_kernel_size = discriminator_kernel_size
+        self.discriminator_stride = discriminator_stride
+        self.discriminator_periods = discriminator_periods
+        self.discriminator_period_channels = discriminator_period_channels
+        self.discriminator_scale_channels = discriminator_scale_channels
+        self.segment_size = segment_size
+        self.hop_length = hop_length
+
+        if len(upsample_kernel_sizes) != len(upsample_rates):
+            raise ValueError(
+                f"The length of `upsample_kernel_sizes` ({len(upsample_kernel_sizes)}) must match the length of "
+                f"`upsample_rates` ({len(upsample_rates)})"
+            )
+
+        super().__init__(**kwargs)
+
+#.............................................................................................
+
+class VitsPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+    config_class = VitsConfig
+    base_model_prefix = "vits"
+    main_input_name = "input_ids"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.Conv1d):
+            nn.init.kaiming_normal_(module.weight)
+            if module.bias is not None:
+                k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))
+                nn.init.uniform_(module.bias, a=-k, b=k)
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+#.............................................................................................

VitsModelSplit/vits_model.py

@@ -0,0 +1,447 @@
+
+import numpy as np
+import torch
+from torch import nn
+import math
+from typing import Any, Callable, Optional, Tuple, Union
+from torch.cuda.amp import autocast, GradScaler
+
+from .vits_config import VitsConfig,VitsPreTrainedModel
+from .flow import VitsResidualCouplingBlock
+from .duration_predictor import VitsDurationPredictor, VitsStochasticDurationPredictor
+from .encoder import VitsTextEncoder
+from .decoder import VitsHifiGan
+from .posterior_encoder import VitsPosteriorEncoder
+from .discriminator import VitsDiscriminator
+from .vits_output import VitsModelOutput, VitsTrainingOutput
+
+
+class VitsModel(VitsPreTrainedModel):
+
+    def __init__(self, config: VitsConfig):
+        super().__init__(config)
+
+        self.config = config
+        self.text_encoder = VitsTextEncoder(config)
+        self.flow = VitsResidualCouplingBlock(config)
+        self.decoder = VitsHifiGan(config)
+
+
+
+        if config.use_stochastic_duration_prediction:
+            self.duration_predictor = VitsStochasticDurationPredictor(config)
+        else:
+            self.duration_predictor = VitsDurationPredictor(config)
+
+        if config.num_speakers > 1:
+            self.embed_speaker = nn.Embedding(config.num_speakers, config.speaker_embedding_size)
+
+        # This is used only for training.
+        self.posterior_encoder = VitsPosteriorEncoder(config)
+        self.discriminator = VitsDiscriminator(config)
+
+        # These parameters control the synthesised speech properties
+        self.speaking_rate = config.speaking_rate
+        self.noise_scale = config.noise_scale
+        self.noise_scale_duration = config.noise_scale_duration
+        self.segment_size = self.config.segment_size // self.config.hop_length
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+
+    #....................................
+
+    def monotonic_align_max_path(self,log_likelihoods, mask):
+        # used for training - awfully slow
+        # an alternative is proposed in examples/pytorch/text-to-speech/run_vits_finetuning.py
+        path = torch.zeros_like(log_likelihoods)
+
+        text_length_maxs = mask.sum(1)[:, 0]
+        latent_length_maxs = mask.sum(2)[:, 0]
+
+        indexes = latent_length_maxs - 1
+
+        max_neg_val = -1e9
+
+        for batch_id in range(len(path)):
+            index = int(indexes[batch_id].item())
+            text_length_max = int(text_length_maxs[batch_id].item())
+            latent_length_max = int(latent_length_maxs[batch_id].item())
+
+            for y in range(text_length_max):
+                for x in range(max(0, latent_length_max + y - text_length_max), min(latent_length_max, y + 1)):
+                    if x == y:
+                        v_cur = max_neg_val
+                    else:
+                        v_cur = log_likelihoods[batch_id, y - 1, x]
+                    if x == 0:
+                        if y == 0:
+                            v_prev = 0.0
+                        else:
+                            v_prev = max_neg_val
+                    else:
+                        v_prev = log_likelihoods[batch_id, y - 1, x - 1]
+                    log_likelihoods[batch_id, y, x] += max(v_prev, v_cur)
+
+            for y in range(text_length_max - 1, -1, -1):
+                path[batch_id, y, index] = 1
+                if index != 0 and (
+                    index == y or log_likelihoods[batch_id, y - 1, index] < log_likelihoods[batch_id, y - 1, index - 1]
+                ):
+                    index = index - 1
+        return path
+
+    #....................................
+
+    def slice_segments(self,hidden_states, ids_str, segment_size=4):
+
+        batch_size, channels, _ = hidden_states.shape
+        # 1d tensor containing the indices to keep
+        indices = torch.arange(segment_size).to(ids_str.device)
+        # extend the indices to match the shape of hidden_states
+        indices = indices.view(1, 1, -1).expand(batch_size, channels, -1)
+        # offset indices with ids_str
+        indices = indices + ids_str.view(-1, 1, 1)
+        # gather indices
+        output = torch.gather(hidden_states, dim=2, index=indices)
+
+        return output
+
+
+    #....................................
+
+
+    def rand_slice_segments(self,hidden_states, sample_lengths=None, segment_size=4):
+
+        batch_size, _, seq_len = hidden_states.size()
+        if sample_lengths is None:
+            sample_lengths = seq_len
+        ids_str_max = sample_lengths - segment_size + 1
+        ids_str = (torch.rand([batch_size]).to(device=hidden_states.device) * ids_str_max).to(dtype=torch.long)
+        ret = self.slice_segments(hidden_states, ids_str, segment_size)
+
+        return ret, ids_str
+
+    #....................................
+
+    def resize_speaker_embeddings(
+        self,
+        new_num_speakers: int,
+        speaker_embedding_size: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = 2,
+    ):
+        if pad_to_multiple_of is not None:
+            new_num_speakers = ((new_num_speakers + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of
+
+        # first, take care of embed_speaker
+        if self.config.num_speakers <= 1:
+            if speaker_embedding_size is None:
+                raise ValueError(
+                    "The current model had no previous speaker embedding, but `speaker_embedding_size` is not specified. Pass `speaker_embedding_size` to this method."
+                )
+            # create new embedding layer
+            new_embeddings = nn.Embedding(
+                new_num_speakers,
+                speaker_embedding_size,
+                device=self.device,
+            )
+            # initialize all new embeddings
+            self._init_weights(new_embeddings)
+        else:
+            new_embeddings = self._get_resized_embeddings(self.embed_speaker, new_num_speakers)
+
+        self.embed_speaker = new_embeddings
+
+        # then take care of sub-models
+        self.flow.resize_speaker_embeddings(speaker_embedding_size)
+        for flow in self.flow.flows:
+            self._init_weights(flow.wavenet.cond_layer)
+
+        self.decoder.resize_speaker_embedding(speaker_embedding_size)
+        self._init_weights(self.decoder.cond)
+
+        self.duration_predictor.resize_speaker_embeddings(speaker_embedding_size)
+        self._init_weights(self.duration_predictor.cond)
+
+        self.posterior_encoder.resize_speaker_embeddings(speaker_embedding_size)
+        self._init_weights(self.posterior_encoder.wavenet.cond_layer)
+
+        self.config.num_speakers = new_num_speakers
+        self.config.speaker_embedding_size = speaker_embedding_size
+
+    #....................................
+
+    def get_input_embeddings(self):
+        return self.text_encoder.get_input_embeddings()
+
+    #....................................
+
+    def set_input_embeddings(self, value):
+        self.text_encoder.set_input_embeddings(value)
+
+    #....................................
+
+    def apply_weight_norm(self):
+        self.decoder.apply_weight_norm()
+        self.flow.apply_weight_norm()
+        self.posterior_encoder.apply_weight_norm()
+
+    #....................................
+
+    def remove_weight_norm(self):
+        self.decoder.remove_weight_norm()
+        self.flow.remove_weight_norm()
+        self.posterior_encoder.remove_weight_norm()
+
+    #....................................
+
+    def discriminate(self, hidden_states):
+        return self.discriminator(hidden_states)
+
+    #....................................
+
+    def get_encoder(self):
+        return self.text_encoder
+
+    #....................................
+
+    def _inference_forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        speaker_embeddings: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+    ):
+        text_encoder_output = self.text_encoder(
+            input_ids=input_ids,
+            padding_mask=padding_mask,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
+        hidden_states = hidden_states.transpose(1, 2)
+        input_padding_mask = padding_mask.transpose(1, 2)
+
+        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
+        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances
+
+        if self.config.use_stochastic_duration_prediction:
+            log_duration = self.duration_predictor(
+                hidden_states,
+                input_padding_mask,
+                speaker_embeddings,
+                reverse=True,
+                noise_scale=self.noise_scale_duration,
+            )
+        else:
+            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
+
+        length_scale = 1.0 / self.speaking_rate
+        duration = torch.ceil(torch.exp(log_duration) * input_padding_mask * length_scale)
+        predicted_lengths = torch.clamp_min(torch.sum(duration, [1, 2]), 1).long()
+
+
+        # Create a padding mask for the output lengths of shape (batch, 1, max_output_length)
+        indices = torch.arange(predicted_lengths.max(), dtype=predicted_lengths.dtype, device=predicted_lengths.device)
+        output_padding_mask = indices.unsqueeze(0) < predicted_lengths.unsqueeze(1)
+        output_padding_mask = output_padding_mask.unsqueeze(1).to(input_padding_mask.dtype)
+
+        # Reconstruct an attention tensor of shape (batch, 1, out_length, in_length)
+        attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(output_padding_mask, -1)
+        batch_size, _, output_length, input_length = attn_mask.shape
+        cum_duration = torch.cumsum(duration, -1).view(batch_size * input_length, 1)
+        indices = torch.arange(output_length, dtype=duration.dtype, device=duration.device)
+        valid_indices = indices.unsqueeze(0) < cum_duration
+        valid_indices = valid_indices.to(attn_mask.dtype).view(batch_size, input_length, output_length)
+        padded_indices = valid_indices - nn.functional.pad(valid_indices, [0, 0, 1, 0, 0, 0])[:, :-1]
+        attn = padded_indices.unsqueeze(1).transpose(2, 3) * attn_mask
+
+        # Expand prior distribution
+        prior_means = torch.matmul(attn.squeeze(1), prior_means).transpose(1, 2)
+        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances).transpose(1, 2)
+
+        prior_latents = prior_means + torch.randn_like(prior_means) * torch.exp(prior_log_variances) * self.noise_scale
+        latents = self.flow(prior_latents, output_padding_mask, speaker_embeddings, reverse=True)
+
+        spectrogram = latents * output_padding_mask
+        waveform = self.decoder(spectrogram, speaker_embeddings)
+        waveform = waveform.squeeze(1)
+        sequence_lengths = predicted_lengths * np.prod(self.config.upsample_rates)
+
+        if not return_dict:
+            outputs = (waveform, sequence_lengths, spectrogram) + text_encoder_output[3:]
+            return outputs
+
+        return VitsModelOutput(
+            waveform=waveform,
+            sequence_lengths=sequence_lengths,
+            spectrogram=spectrogram,
+            hidden_states=text_encoder_output.hidden_states,
+            attentions=text_encoder_output.attentions,
+        )
+
+    #....................................
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        speaker_id: Optional[int] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        labels: Optional[torch.FloatTensor] = None,
+        labels_attention_mask: Optional[torch.Tensor] = None,
+        monotonic_alignment_function: Optional[Callable] = None,
+    ) -> Union[Tuple[Any], VitsModelOutput]:
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        monotonic_alignment_function = (
+            self.monotonic_align_max_path if monotonic_alignment_function is None else monotonic_alignment_function
+        )
+
+        if attention_mask is not None:
+            input_padding_mask = attention_mask.unsqueeze(-1).float()
+        else:
+            input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).float()
+
+        if self.config.num_speakers > 1 and speaker_id is not None:
+            if isinstance(speaker_id, int):
+                speaker_id = torch.full(size=(1,), fill_value=speaker_id, device=self.device)
+            elif isinstance(speaker_id, (list, tuple, np.ndarray)):
+                speaker_id = torch.tensor(speaker_id, device=self.device)
+
+            if not ((0 <= speaker_id).all() and (speaker_id < self.config.num_speakers).all()).item():
+                raise ValueError(f"Set `speaker_id` in the range 0-{self.config.num_speakers - 1}.")
+            if not (len(speaker_id) == 1 or len(speaker_id == len(input_ids))):
+                raise ValueError(
+                    f"You passed {len(speaker_id)} `speaker_id` but you should either pass one speaker id or `batch_size` `speaker_id`."
+                )
+
+            speaker_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1)
+        else:
+            speaker_embeddings = None
+
+        # if inference, return inference forward of VitsModel
+        if labels is None:
+            return self._inference_forward(
+                input_ids,
+                attention_mask,
+                speaker_embeddings,
+                output_attentions,
+                output_hidden_states,
+                return_dict,
+                input_padding_mask,
+            )
+
+        if labels_attention_mask is not None:
+            labels_padding_mask = labels_attention_mask.unsqueeze(1).float()
+        else:
+            labels_attention_mask = torch.ones((labels.shape[0], labels.shape[2])).float().to(self.device)
+            labels_padding_mask = labels_attention_mask.unsqueeze(1)
+
+        text_encoder_output = self.text_encoder(
+            input_ids=input_ids,
+            padding_mask=input_padding_mask,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
+        hidden_states = hidden_states.transpose(1, 2)
+        input_padding_mask = input_padding_mask.transpose(1, 2)
+        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
+        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances
+
+        latents, posterior_means, posterior_log_variances = self.posterior_encoder(
+            labels, labels_padding_mask, speaker_embeddings
+        )
+        prior_latents = self.flow(latents, labels_padding_mask, speaker_embeddings, reverse=False)
+
+        prior_means, prior_log_variances = prior_means.transpose(1, 2), prior_log_variances.transpose(1, 2)
+        with torch.no_grad():
+            # negative cross-entropy
+
+            # [batch_size, d, latent_length]
+            prior_variances = torch.exp(-2 * prior_log_variances)
+            # [batch_size, 1, latent_length]
+            neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - prior_log_variances, [1], keepdim=True)
+            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
+            neg_cent2 = torch.matmul(-0.5 * (prior_latents**2).transpose(1, 2), prior_variances)
+            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
+            neg_cent3 = torch.matmul(prior_latents.transpose(1, 2), (prior_means * prior_variances))
+            # [batch_size, 1, latent_length]
+            neg_cent4 = torch.sum(-0.5 * (prior_means**2) * prior_variances, [1], keepdim=True)
+
+            # [batch_size, text_length, latent_length]
+            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+
+            attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(labels_padding_mask, -1)
+
+            attn = monotonic_alignment_function(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()
+
+        durations = attn.sum(2)
+
+        if self.config.use_stochastic_duration_prediction:
+            log_duration = self.duration_predictor(
+                hidden_states, input_padding_mask, speaker_embeddings, durations=durations, reverse=False
+            )
+            log_duration = log_duration / torch.sum(input_padding_mask)
+        else:
+            log_duration_padded = torch.log(durations + 1e-6) * input_padding_mask
+            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
+            log_duration = torch.sum((log_duration - log_duration_padded) ** 2, [1, 2]) / torch.sum(input_padding_mask)
+
+        # expand priors
+        prior_means = torch.matmul(attn.squeeze(1), prior_means.transpose(1, 2)).transpose(1, 2)
+        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances.transpose(1, 2)).transpose(1, 2)
+
+        label_lengths = labels_attention_mask.sum(dim=1)
+        latents_slice, ids_slice = self.rand_slice_segments(latents, label_lengths, segment_size=self.segment_size)
+
+        waveform = self.decoder(latents_slice, speaker_embeddings)
+
+        if not return_dict:
+            outputs = (
+                waveform,
+                log_duration,
+                attn,
+                ids_slice,
+                input_padding_mask,
+                labels_padding_mask,
+                latents,
+                prior_latents,
+                prior_means,
+                prior_log_variances,
+                posterior_means,
+                posterior_log_variances,
+            )
+            return outputs
+
+        return VitsTrainingOutput(
+            waveform=waveform,
+            log_duration=log_duration,
+            attn=attn,
+            ids_slice=ids_slice,
+            input_padding_mask=input_padding_mask,
+            labels_padding_mask=labels_padding_mask,
+            latents=latents,
+            prior_latents=prior_latents,
+            prior_means=prior_means,
+            prior_log_variances=prior_log_variances,
+            posterior_means=posterior_means,
+            posterior_log_variances=posterior_log_variances,
+        )
+ 

VitsModelSplit/vits_model3.py

@@ -0,0 +1,670 @@
+
+import numpy as np
+import torch
+from torch import nn
+import math
+from typing import Any, Callable, Optional, Tuple, Union
+from torch.cuda.amp import autocast, GradScaler
+
+from .vits_config import VitsConfig,VitsPreTrainedModel
+from .flow import VitsResidualCouplingBlock
+from .duration_predictor import VitsDurationPredictor, VitsStochasticDurationPredictor
+from .encoder import VitsTextEncoder
+from .decoder import VitsHifiGan
+from .posterior_encoder import VitsPosteriorEncoder
+from .discriminator import VitsDiscriminator
+from .vits_output import VitsModelOutput, VitsTrainingOutput
+from  .dataset_features_collector import FeaturesCollectionDataset
+from .feature_extraction import VitsFeatureExtractor
+
+import os
+import sys
+from typing import Optional
+import tempfile
+from torch.cuda.amp import autocast, GradScaler
+
+from IPython.display import clear_output
+from transformers import set_seed
+import wandb
+import logging
+import copy
+Lst=['input_ids',
+ 'attention_mask',
+ 'waveform',
+ 'labels',
+ 'labels_attention_mask',
+ 'mel_scaled_input_features']
+
+def covert_cuda_batch(d):
+  #return d
+  for key in Lst:
+      d[key]=d[key].cuda(non_blocking=True)
+  # for key in d['text_encoder_output']:
+  #   d['text_encoder_output'][key]=d['text_encoder_output'][key].cuda(non_blocking=True)
+  for key in d['posterior_encode_output']:
+    d['posterior_encode_output'][key]=d['posterior_encode_output'][key].cuda(non_blocking=True)
+
+  return d
+def generator_loss(disc_outputs):
+    total_loss = 0
+    gen_losses = []
+    for disc_output in disc_outputs:
+        disc_output = disc_output
+        loss = torch.mean((1 - disc_output) ** 2)
+        gen_losses.append(loss)
+        total_loss += loss
+
+    return total_loss, gen_losses
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    real_losses = 0
+    generated_losses = 0
+    for disc_real, disc_generated in zip(disc_real_outputs, disc_generated_outputs):
+        real_loss = torch.mean((1 - disc_real) ** 2)
+        generated_loss = torch.mean(disc_generated**2)
+        loss += real_loss + generated_loss
+        real_losses += real_loss
+        generated_losses += generated_loss
+
+    return loss, real_losses, generated_losses
+
+def feature_loss(feature_maps_real, feature_maps_generated):
+    loss = 0
+    for feature_map_real, feature_map_generated in zip(feature_maps_real, feature_maps_generated):
+        for real, generated in zip(feature_map_real, feature_map_generated):
+            real = real.detach()
+            loss += torch.mean(torch.abs(real - generated))
+
+    return loss * 2
+def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
+  """
+  z_p, logs_q: [b, h, t_t]
+  m_p, logs_p: [b, h, t_t]
+  """
+  z_p = z_p.float()
+  logs_q = logs_q.float()
+  m_p = m_p.float()
+  logs_p = logs_p.float()
+  z_mask = z_mask.float()
+
+  kl = logs_p - logs_q - 0.5
+  kl += 0.5 * ((z_p - m_p)**2) * torch.exp(-2. * logs_p)
+  kl = torch.sum(kl * z_mask)
+  l = kl / torch.sum(z_mask)
+  return l
+#.............................................
+# def kl_loss(prior_latents, posterior_log_variance, prior_means, prior_log_variance, labels_mask):
+
+
+#   kl = prior_log_variance - posterior_log_variance - 0.5
+#   kl += 0.5 * ((prior_latents - prior_means) ** 2) * torch.exp(-2.0 * prior_log_variance)
+#   kl = torch.sum(kl * labels_mask)
+#   loss = kl / torch.sum(labels_mask)
+#   return loss
+
+def get_state_grad_loss(k1=True,
+                             mel=True,
+                             duration=True,
+                             generator=True,
+                             discriminator=True):
+                             return {'k1':k1,'mel':mel,'duration':duration,'generator':generator,'discriminator':discriminator}
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+  if isinstance(parameters, torch.Tensor):
+    parameters = [parameters]
+  parameters = list(filter(lambda p: p.grad is not None, parameters))
+  norm_type = float(norm_type)
+  if clip_value is not None:
+    clip_value = float(clip_value)
+
+  total_norm = 0
+  for p in parameters:
+    param_norm = p.grad.data.norm(norm_type)
+    total_norm += param_norm.item() ** norm_type
+    if clip_value is not None:
+      p.grad.data.clamp_(min=-clip_value, max=clip_value)
+  total_norm = total_norm ** (1. / norm_type)
+  return total_norm
+
+
+class VitsModel(VitsPreTrainedModel):
+
+    def __init__(self, config: VitsConfig):
+        super().__init__(config)
+
+        self.config = config
+        self.text_encoder = VitsTextEncoder(config)
+        self.flow = VitsResidualCouplingBlock(config)
+        self.decoder = VitsHifiGan(config)
+
+
+
+        if config.use_stochastic_duration_prediction:
+            self.duration_predictor = VitsStochasticDurationPredictor(config)
+        else:
+            self.duration_predictor = VitsDurationPredictor(config)
+
+        if config.num_speakers > 1:
+            self.embed_speaker = nn.Embedding(config.num_speakers, config.speaker_embedding_size)
+
+        # This is used only for training.
+        self.posterior_encoder = VitsPosteriorEncoder(config)
+        self.discriminator = VitsDiscriminator(config)
+
+        # These parameters control the synthesised speech properties
+        self.speaking_rate = config.speaking_rate
+        self.noise_scale = config.noise_scale
+        self.noise_scale_duration = config.noise_scale_duration
+        self.segment_size = self.config.segment_size // self.config.hop_length
+
+        # Initialize weights and apply final processing
+        self.post_init()
+        self.monotonic_alignment_function=self.monotonic_align_max_path
+
+
+
+    #....................................
+    def setMfA(self,fn):
+      self.monotonic_alignment_function=fn
+
+
+
+    def monotonic_align_max_path(self,log_likelihoods, mask):
+        # used for training - awfully slow
+        # an alternative is proposed in examples/pytorch/text-to-speech/run_vits_finetuning.py
+        path = torch.zeros_like(log_likelihoods)
+
+        text_length_maxs = mask.sum(1)[:, 0]
+        latent_length_maxs = mask.sum(2)[:, 0]
+
+        indexes = latent_length_maxs - 1
+
+        max_neg_val = -1e9
+
+        for batch_id in range(len(path)):
+            index = int(indexes[batch_id].item())
+            text_length_max = int(text_length_maxs[batch_id].item())
+            latent_length_max = int(latent_length_maxs[batch_id].item())
+
+            for y in range(text_length_max):
+                for x in range(max(0, latent_length_max + y - text_length_max), min(latent_length_max, y + 1)):
+                    if x == y:
+                        v_cur = max_neg_val
+                    else:
+                        v_cur = log_likelihoods[batch_id, y - 1, x]
+                    if x == 0:
+                        if y == 0:
+                            v_prev = 0.0
+                        else:
+                            v_prev = max_neg_val
+                    else:
+                        v_prev = log_likelihoods[batch_id, y - 1, x - 1]
+                    log_likelihoods[batch_id, y, x] += max(v_prev, v_cur)
+
+            for y in range(text_length_max - 1, -1, -1):
+                path[batch_id, y, index] = 1
+                if index != 0 and (
+                    index == y or log_likelihoods[batch_id, y - 1, index] < log_likelihoods[batch_id, y - 1, index - 1]
+                ):
+                    index = index - 1
+        return path
+
+    #....................................
+
+    def slice_segments(self,hidden_states, ids_str, segment_size=4):
+
+        batch_size, channels, _ = hidden_states.shape
+        # 1d tensor containing the indices to keep
+        indices = torch.arange(segment_size).to(ids_str.device)
+        # extend the indices to match the shape of hidden_states
+        indices = indices.view(1, 1, -1).expand(batch_size, channels, -1)
+        # offset indices with ids_str
+        indices = indices + ids_str.view(-1, 1, 1)
+        # gather indices
+        output = torch.gather(hidden_states, dim=2, index=indices)
+
+        return output
+
+
+    #....................................
+
+
+    def rand_slice_segments(self,hidden_states, sample_lengths=None, segment_size=4):
+
+        batch_size, _, seq_len = hidden_states.size()
+        if sample_lengths is None:
+            sample_lengths = seq_len
+        ids_str_max = sample_lengths - segment_size + 1
+        ids_str = (torch.rand([batch_size]).to(device=hidden_states.device) * ids_str_max).to(dtype=torch.long)
+        ret = self.slice_segments(hidden_states, ids_str, segment_size)
+
+        return ret, ids_str
+
+    #....................................
+
+    def resize_speaker_embeddings(
+        self,
+        new_num_speakers: int,
+        speaker_embedding_size: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = 2,
+    ):
+        if pad_to_multiple_of is not None:
+            new_num_speakers = ((new_num_speakers + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of
+
+        # first, take care of embed_speaker
+        if self.config.num_speakers <= 1:
+            if speaker_embedding_size is None:
+                raise ValueError(
+                    "The current model had no previous speaker embedding, but `speaker_embedding_size` is not specified. Pass `speaker_embedding_size` to this method."
+                )
+            # create new embedding layer
+            new_embeddings = nn.Embedding(
+                new_num_speakers,
+                speaker_embedding_size,
+                device=self.device,
+            )
+            # initialize all new embeddings
+            self._init_weights(new_embeddings)
+        else:
+            new_embeddings = self._get_resized_embeddings(self.embed_speaker, new_num_speakers)
+
+        self.embed_speaker = new_embeddings
+
+        # then take care of sub-models
+        self.flow.resize_speaker_embeddings(speaker_embedding_size)
+        for flow in self.flow.flows:
+            self._init_weights(flow.wavenet.cond_layer)
+
+        self.decoder.resize_speaker_embedding(speaker_embedding_size)
+        self._init_weights(self.decoder.cond)
+
+        self.duration_predictor.resize_speaker_embeddings(speaker_embedding_size)
+        self._init_weights(self.duration_predictor.cond)
+
+        self.posterior_encoder.resize_speaker_embeddings(speaker_embedding_size)
+        self._init_weights(self.posterior_encoder.wavenet.cond_layer)
+
+        self.config.num_speakers = new_num_speakers
+        self.config.speaker_embedding_size = speaker_embedding_size
+
+    #....................................
+
+    def get_input_embeddings(self):
+        return self.text_encoder.get_input_embeddings()
+
+    #....................................
+
+    def set_input_embeddings(self, value):
+        self.text_encoder.set_input_embeddings(value)
+
+    #....................................
+
+    def apply_weight_norm(self):
+        self.decoder.apply_weight_norm()
+        self.flow.apply_weight_norm()
+        self.posterior_encoder.apply_weight_norm()
+
+    #....................................
+
+    def remove_weight_norm(self):
+        self.decoder.remove_weight_norm()
+        self.flow.remove_weight_norm()
+        self.posterior_encoder.remove_weight_norm()
+
+    #....................................
+
+    def discriminate(self, hidden_states):
+        return self.discriminator(hidden_states)
+
+    #....................................
+
+    def get_encoder(self):
+        return self.text_encoder
+
+    #....................................
+
+    def _inference_forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        speaker_embeddings: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+    ):
+        text_encoder_output = self.text_encoder(
+            input_ids=input_ids,
+            padding_mask=padding_mask,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
+        hidden_states = hidden_states.transpose(1, 2)
+        input_padding_mask = padding_mask.transpose(1, 2)
+
+        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
+        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances
+
+        if self.config.use_stochastic_duration_prediction:
+            log_duration = self.duration_predictor(
+                hidden_states,
+                input_padding_mask,
+                speaker_embeddings,
+                reverse=True,
+                noise_scale=self.noise_scale_duration,
+            )
+        else:
+            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
+
+        length_scale = 1.0 / self.speaking_rate
+        duration = torch.ceil(torch.exp(log_duration) * input_padding_mask * length_scale)
+        predicted_lengths = torch.clamp_min(torch.sum(duration, [1, 2]), 1).long()
+
+
+        # Create a padding mask for the output lengths of shape (batch, 1, max_output_length)
+        indices = torch.arange(predicted_lengths.max(), dtype=predicted_lengths.dtype, device=predicted_lengths.device)
+        output_padding_mask = indices.unsqueeze(0) < predicted_lengths.unsqueeze(1)
+        output_padding_mask = output_padding_mask.unsqueeze(1).to(input_padding_mask.dtype)
+
+        # Reconstruct an attention tensor of shape (batch, 1, out_length, in_length)
+        attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(output_padding_mask, -1)
+        batch_size, _, output_length, input_length = attn_mask.shape
+        cum_duration = torch.cumsum(duration, -1).view(batch_size * input_length, 1)
+        indices = torch.arange(output_length, dtype=duration.dtype, device=duration.device)
+        valid_indices = indices.unsqueeze(0) < cum_duration
+        valid_indices = valid_indices.to(attn_mask.dtype).view(batch_size, input_length, output_length)
+        padded_indices = valid_indices - nn.functional.pad(valid_indices, [0, 0, 1, 0, 0, 0])[:, :-1]
+        attn = padded_indices.unsqueeze(1).transpose(2, 3) * attn_mask
+
+        # Expand prior distribution
+        prior_means = torch.matmul(attn.squeeze(1), prior_means).transpose(1, 2)
+        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances).transpose(1, 2)
+
+        prior_latents = prior_means + torch.randn_like(prior_means) * torch.exp(prior_log_variances) * self.noise_scale
+        latents = self.flow(prior_latents, output_padding_mask, speaker_embeddings, reverse=True)
+
+        spectrogram = latents * output_padding_mask
+        waveform = self.decoder(spectrogram, speaker_embeddings)
+        waveform = waveform.squeeze(1)
+        sequence_lengths = predicted_lengths * np.prod(self.config.upsample_rates)
+
+        if not return_dict:
+            outputs = (waveform, sequence_lengths, spectrogram) + text_encoder_output[3:]
+            return outputs
+
+        return VitsModelOutput(
+            waveform=waveform,
+            sequence_lengths=sequence_lengths,
+            spectrogram=spectrogram,
+            hidden_states=text_encoder_output.hidden_states,
+            attentions=text_encoder_output.attentions,
+        )
+
+    #....................................
+
+    def forward_k(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        speaker_id: Optional[int] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        labels: Optional[torch.FloatTensor] = None,
+        labels_attention_mask: Optional[torch.Tensor] = None,
+        monotonic_alignment_function: Optional[Callable] = None,
+    ) -> Union[Tuple[Any], VitsModelOutput]:
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        monotonic_alignment_function = (
+            self.monotonic_align_max_path if monotonic_alignment_function is None else monotonic_alignment_function
+        )
+
+        if attention_mask is not None:
+            input_padding_mask = attention_mask.unsqueeze(-1).float()
+        else:
+            input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).float()
+
+        if self.config.num_speakers > 1 and speaker_id is not None:
+            if isinstance(speaker_id, int):
+                speaker_id = torch.full(size=(1,), fill_value=speaker_id, device=self.device)
+            elif isinstance(speaker_id, (list, tuple, np.ndarray)):
+                speaker_id = torch.tensor(speaker_id, device=self.device)
+
+            if not ((0 <= speaker_id).all() and (speaker_id < self.config.num_speakers).all()).item():
+                raise ValueError(f"Set `speaker_id` in the range 0-{self.config.num_speakers - 1}.")
+            if not (len(speaker_id) == 1 or len(speaker_id == len(input_ids))):
+                raise ValueError(
+                    f"You passed {len(speaker_id)} `speaker_id` but you should either pass one speaker id or `batch_size` `speaker_id`."
+                )
+
+            speaker_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1)
+        else:
+            speaker_embeddings = None
+
+        # if inference, return inference forward of VitsModel
+        if labels is None:
+            return self._inference_forward(
+                input_ids,
+                attention_mask,
+                speaker_embeddings,
+                output_attentions,
+                output_hidden_states,
+                return_dict,
+                input_padding_mask,
+            )
+
+        if labels_attention_mask is not None:
+            labels_padding_mask = labels_attention_mask.unsqueeze(1).float()
+        else:
+            labels_attention_mask = torch.ones((labels.shape[0], labels.shape[2])).float().to(self.device)
+            labels_padding_mask = labels_attention_mask.unsqueeze(1)
+
+        text_encoder_output = self.text_encoder(
+            input_ids=input_ids,
+            padding_mask=input_padding_mask,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
+        hidden_states = hidden_states.transpose(1, 2)
+        input_padding_mask = input_padding_mask.transpose(1, 2)
+        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
+        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances
+
+        latents, posterior_means, posterior_log_variances = self.posterior_encoder(
+            labels, labels_padding_mask, speaker_embeddings
+        )
+        prior_latents = self.flow(latents, labels_padding_mask, speaker_embeddings, reverse=False)
+
+        prior_means, prior_log_variances = prior_means.transpose(1, 2), prior_log_variances.transpose(1, 2)
+        with torch.no_grad():
+            # negative cross-entropy
+
+            # [batch_size, d, latent_length]
+            prior_variances = torch.exp(-2 * prior_log_variances)
+            # [batch_size, 1, latent_length]
+            neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - prior_log_variances, [1], keepdim=True)
+            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
+            neg_cent2 = torch.matmul(-0.5 * (prior_latents**2).transpose(1, 2), prior_variances)
+            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
+            neg_cent3 = torch.matmul(prior_latents.transpose(1, 2), (prior_means * prior_variances))
+            # [batch_size, 1, latent_length]
+            neg_cent4 = torch.sum(-0.5 * (prior_means**2) * prior_variances, [1], keepdim=True)
+
+            # [batch_size, text_length, latent_length]
+            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+
+            attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(labels_padding_mask, -1)
+
+            attn = monotonic_alignment_function(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()
+
+        durations = attn.sum(2)
+
+        if self.config.use_stochastic_duration_prediction:
+            log_duration = self.duration_predictor(
+                hidden_states, input_padding_mask, speaker_embeddings, durations=durations, reverse=False
+            )
+            log_duration = log_duration / torch.sum(input_padding_mask)
+        else:
+            log_duration_padded = torch.log(durations + 1e-6) * input_padding_mask
+            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
+            log_duration = torch.sum((log_duration - log_duration_padded) ** 2, [1, 2]) / torch.sum(input_padding_mask)
+
+        # expand priors
+        prior_means = torch.matmul(attn.squeeze(1), prior_means.transpose(1, 2)).transpose(1, 2)
+        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances.transpose(1, 2)).transpose(1, 2)
+
+        label_lengths = labels_attention_mask.sum(dim=1)
+        latents_slice, ids_slice = self.rand_slice_segments(latents, label_lengths, segment_size=self.segment_size)
+
+        waveform = self.decoder(latents_slice, speaker_embeddings)
+
+        if not return_dict:
+            outputs = (
+                waveform,
+                log_duration,
+                attn,
+                ids_slice,
+                input_padding_mask,
+                labels_padding_mask,
+                latents,
+                prior_latents,
+                prior_means,
+                prior_log_variances,
+                posterior_means,
+                posterior_log_variances,
+            )
+            return outputs
+
+        return VitsTrainingOutput(
+            waveform=waveform,
+            log_duration=log_duration,
+            attn=attn,
+            ids_slice=ids_slice,
+            input_padding_mask=input_padding_mask,
+            labels_padding_mask=labels_padding_mask,
+            latents=latents,
+            prior_latents=prior_latents,
+            prior_means=prior_means,
+            prior_log_variances=prior_log_variances,
+            posterior_means=posterior_means,
+            posterior_log_variances=posterior_log_variances,
+        )
+
+
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        speaker_id: Optional[int] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        labels: Optional[torch.FloatTensor] = None,
+        labels_attention_mask: Optional[torch.Tensor] = None,
+        text_encoder_output=None,
+        posterior_encode_output=None,
+        monotonic_alignment_function: Optional[Callable] = None,
+        speaker_embeddings=None
+    ) -> Union[Tuple[Any], VitsModelOutput]:
+
+        #output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states# if output_hidden_states is not None else self.config.output_hidden_states
+        )
+       # return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+
+       # if attention_mask is not None:
+        input_padding_mask = attention_mask.unsqueeze(-1).float()
+        #else:
+         #   input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).float()
+
+        # speaker_embeddings=None
+       # if labels_attention_mask is not None:
+        labels_padding_mask = labels_attention_mask.unsqueeze(1).float()
+        # else:
+        #     labels_attention_mask = torch.ones((labels.shape[0], labels.shape[2])).float().to(self.device)
+        #     labels_padding_mask = labels_attention_mask.unsqueeze(1)
+        if text_encoder_output is None:
+          text_encoder_output = self.text_encoder(
+                  input_ids=input_ids,
+                  padding_mask=input_padding_mask,
+                  attention_mask=attention_mask,
+                  output_attentions=output_attentions,
+                  output_hidden_states=output_hidden_states,
+                  return_dict=return_dict,
+              )
+        #hidden_states = text_encoder_output[0] #if not return_dict else text_encoder_output.last_hidden_state
+        hidden_states = text_encoder_output[0].transpose(1, 2)
+        input_padding_mask = input_padding_mask.transpose(1, 2)
+        prior_means = text_encoder_output[1].transpose(1, 2) #if not return_dict else text_encoder_output.prior_means
+        prior_log_variances = text_encoder_output[2].transpose(1, 2) #if not return_dict else text_encoder_output.prior_log_variances
+
+        # if posterior_encode_output is None:
+        #     latents, posterior_means, posterior_log_variances = self.posterior_encoder(
+        #         labels, labels_padding_mask, speaker_embeddings
+        #     )
+        # else:
+        latents=posterior_encode_output['posterior_latents']
+        posterior_means=posterior_encode_output['posterior_means']
+        posterior_log_variances=posterior_encode_output['posterior_log_variances']
+
+        prior_latents = self.flow(latents, labels_padding_mask, speaker_embeddings, reverse=False)
+
+        # prior_means, prior_log_variances = prior_means.transpose(1, 2), prior_log_variances.transpose(1, 2)
+        with torch.no_grad():
+            # negative cross-entropy
+
+            # [batch_size, d, latent_length]
+            prior_variances = torch.exp(-2 * prior_log_variances)
+            # [batch_size, 1, latent_length]
+            neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - prior_log_variances, [1], keepdim=True)
+            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
+            neg_cent2 = torch.matmul(-0.5 * (prior_latents**2).transpose(1, 2), prior_variances)
+            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
+            neg_cent3 = torch.matmul(prior_latents.transpose(1, 2), (prior_means * prior_variances))
+            # [batch_size, 1, latent_length]
+            neg_cent4 = torch.sum(-0.5 * (prior_means**2) * prior_variances, [1], keepdim=True)
+
+            # [batch_size, text_length, latent_length]
+            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+
+            attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(labels_padding_mask, -1)
+
+            attn = monotonic_alignment_function(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()
+
+        durations = attn.sum(2)
+
+        #if self.config.use_stochastic_duration_prediction:
+        log_duration = self.duration_predictor(
+                hidden_states, input_padding_mask, speaker_embeddings, durations=durations, reverse=False
+            )
+        log_duration = log_duration / torch.sum(input_padding_mask)
+        # else:
+        #     log_duration_padded = torch.log(durations + 1e-6) * input_padding_mask
+        #     log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
+        #     log_duration = torch.sum((log_duration - log_duration_padded) ** 2, [1, 2]) / torch.sum(input_padding_mask)
+
+        # expand priors
+        prior_means = torch.matmul(attn.squeeze(1), prior_means.transpose(1, 2)).transpose(1, 2)
+        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances.transpose(1, 2)).transpose(1, 2)
+
+        label_lengths = labels_attention_mask.sum(dim=1)
+        latents_slice, ids_slice = self.rand_slice_segments(latents, label_lengths, segment_size=self.segment_size)
+        waveform = self.decoder(latents_slice, speaker_embeddings)
+        return waveform,ids_slice,log_duration,prior_latents,posterior_log_variances,prior_means,prior_log_variances,labels_padding_mask

VitsModelSplit/vits_output.py

@@ -0,0 +1,84 @@
+from typing import Any, Optional, Tuple, Union,List,Dict
+import torch
+from dataclasses import dataclass
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    ModelOutput,
+)
+#.............................................
+
+
+
+@dataclass
+class PosteriorDecoderModelOutput(ModelOutput):
+    labels_padding_mask: torch.FloatTensor = None
+    posterior_latents: torch.FloatTensor = None
+    posterior_means: torch.FloatTensor = None
+    posterior_log_variances: torch.FloatTensor = None
+    latents_slice : torch.FloatTensor = None
+    ids_slice: torch.FloatTensor = None
+    waveform: torch.FloatTensor = None
+    
+#.............................................................................................    
+    
+    
+@dataclass
+class VitsModelOutput(ModelOutput):
+    waveform: torch.FloatTensor = None
+    sequence_lengths: torch.FloatTensor = None
+    spectrogram: Optional[Tuple[torch.FloatTensor]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+#.............................................................................................
+
+@dataclass
+class VitsTrainingOutput(ModelOutput):
+    waveform: torch.FloatTensor = None
+    log_duration: torch.FloatTensor = None
+    attn: torch.FloatTensor = None
+    ids_slice: torch.FloatTensor = None
+    input_padding_mask: torch.FloatTensor = None
+    labels_padding_mask: torch.FloatTensor = None
+    latents: torch.FloatTensor = None
+    prior_latents: torch.FloatTensor = None
+    prior_means: torch.FloatTensor = None
+    prior_log_variances: torch.FloatTensor = None
+    posterior_means: torch.FloatTensor = None
+    posterior_log_variances: torch.FloatTensor = None
+
+
+#.............................................................................................
+
+@dataclass
+class VitsTextEncoderOutput(ModelOutput):
+    """
+    Describes the outputs for the VITS text encoder model, with potential hidden states and attentions.
+
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        prior_means (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            The predicted mean values of the prior distribution for the latent text variables.
+        prior_log_variances (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            The predicted log-variance values of the prior distribution for the latent text variables.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attention weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    last_hidden_state: torch.FloatTensor = None
+    prior_means: torch.FloatTensor = None
+    prior_log_variances: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+#.............................................................................................