chore: import upstream snapshot with attribution

scripts/ssl_tts/make_supdata.py

@@ -0,0 +1,513 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Example Run Command: python make_supdata.py --ssl_model_ckpt_path <PATH TO CKPT> --manifest_path <PATH TO MANIFEST>
+
+import argparse
+import json
+import os
+import time
+from multiprocessing import Pool
+from pathlib import Path
+
+import librosa
+import numpy as np
+import torch
+from omegaconf import open_dict
+from tqdm import tqdm
+
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
+from nemo.collections.tts.models import ssl_tts
+from nemo.collections.tts.parts.utils.tts_dataset_utils import get_base_dir
+from nemo.core.classes import Dataset
+from nemo.core.classes.common import safe_instantiate
+from nemo.utils import logging
+
+
+class AudioDataset(Dataset):
+    def __init__(
+        self,
+        manifest_paths,
+        min_duration=0.5,
+        max_duration=16.0,
+        pad_multiple=1024,
+        sample_rate=22050,
+        sup_data_dir=None,
+    ):
+        self.data = []
+        for manifest_path in manifest_paths:
+            with open(manifest_path, "r") as f:
+                for line in f:
+                    record = json.loads(line)
+                    if record['duration'] < min_duration or record['duration'] > max_duration:
+                        continue
+                    self.data.append(json.loads(line))
+
+        self.base_data_dir = get_base_dir([item["audio_filepath"] for item in self.data])
+        if sup_data_dir is not None:
+            self.sup_data_dir = sup_data_dir
+        else:
+            self.sup_data_dir = os.path.join(self.base_data_dir, "sup_data")
+        if not os.path.exists(self.sup_data_dir):
+            os.makedirs(self.sup_data_dir)
+
+        self.pad_multiple = pad_multiple
+        self.sample_rate = sample_rate
+
+    def __len__(self):
+        return len(self.data)
+
+    def _get_wav_from_filepath(self, audio_filepath):
+        features = AudioSegment.segment_from_file(
+            audio_filepath,
+            target_sr=self.sample_rate,
+            n_segments=-1,
+            trim=False,
+        )
+        audio_samples = features.samples
+        audio, audio_length = torch.tensor(audio_samples), torch.tensor(audio_samples.shape[0]).long()
+
+        # pad audio to a multiple of self.pad_multiple
+        if audio.shape[0] % self.pad_multiple != 0:
+            audio = torch.cat(
+                [audio, torch.zeros(self.pad_multiple - audio.shape[0] % self.pad_multiple, dtype=torch.float)]
+            )
+            audio_length = torch.tensor(audio.shape[0]).long()
+
+        return audio, audio_length
+
+    def pad_collate_fn(self, batch):
+        final_batch = {}
+        for row in batch:
+            for key in row:
+                if key not in final_batch:
+                    final_batch[key] = []
+                final_batch[key].append(row[key])
+
+        max_audio_len = max([_audio_len.item() for _audio_len in final_batch["audio_len"]])
+
+        audios_padded = []
+        for audio in final_batch["audio"]:
+            audio_padded = torch.nn.functional.pad(audio, (0, max_audio_len - audio.size(0)), value=0)
+            audios_padded.append(audio_padded)
+
+        final_batch["audio"] = audios_padded
+        for key in final_batch:
+            if key not in ["rel_audio_path_as_text_id", "wav_path"]:
+                final_batch[key] = torch.stack(final_batch[key])
+
+        return final_batch
+
+    def __getitem__(self, index):
+        sample = self.data[index]
+        rel_audio_path = Path(sample["audio_filepath"]).relative_to(self.base_data_dir).with_suffix("")
+        rel_audio_path_as_text_id = str(rel_audio_path).replace("/", "_")
+        speaker = torch.tensor(sample["speaker"]).long()
+
+        audio, audio_length = self._get_wav_from_filepath(sample["audio_filepath"])
+
+        return {
+            "audio": audio,
+            "audio_len": audio_length,
+            "rel_audio_path_as_text_id": rel_audio_path_as_text_id,
+            "wav_path": sample["audio_filepath"],
+            "speaker": speaker,
+        }
+
+
+def segment_wav(wav, segment_length, segment_hop_size, min_segment_length):
+    if len(wav) < segment_length:
+        pad = torch.zeros(segment_length - len(wav))
+        segment = torch.cat([wav, pad])
+        return [segment]
+    else:
+        si = 0
+        segments = []
+        while si < len(wav) - min_segment_length:
+            segment = wav[si : si + segment_length]
+            if len(segment) < segment_length:
+                pad = torch.zeros(segment_length - len(segment))
+                segment = torch.cat([segment, pad])
+            segments.append(segment)
+            si += segment_hop_size
+        return segments
+
+
+def segment_batch(batch, segment_length=44100, segment_hop_size=22050, min_segment_length=22050):
+    all_segments = []
+    segment_indices = []
+    si = 0
+    for bidx in range(len(batch['audio'])):
+        audio = batch['audio'][bidx]
+        audio_length = batch['audio_len'][bidx]
+        audio_actual = audio[:audio_length]
+        audio_segments = segment_wav(audio_actual, segment_length, segment_hop_size, min_segment_length)
+        all_segments += audio_segments
+        segment_indices.append((si, si + len(audio_segments) - 1))
+        si += len(audio_segments)
+
+    return torch.stack(all_segments), segment_indices
+
+
+def get_mel_spectrogram(fb, wav, stft_params):
+    EPSILON = 1e-9
+    window_fn = torch.hann_window
+
+    spec = torch.stft(
+        input=wav,
+        n_fft=stft_params['n_fft'],  # 1024
+        hop_length=stft_params['hop_length'],  # 256
+        win_length=stft_params['win_length'],  # 1024
+        window=window_fn(stft_params['win_length'], periodic=False).to(torch.float).to('cuda') if window_fn else None,
+        return_complex=True,
+        center=True,
+    )
+
+    if spec.dtype in [torch.cfloat, torch.cdouble]:
+        spec = torch.view_as_real(spec)
+    spec = torch.sqrt(spec.pow(2).sum(-1) + EPSILON)
+
+    mel = torch.matmul(fb.to(spec.dtype), spec)
+    log_mel = torch.log(torch.clamp(mel, min=torch.finfo(mel.dtype).tiny))
+
+    return log_mel
+
+
+def load_wav(wav_path, sample_rate=22050, pad_multiple=1024):
+    wav = AudioSegment.segment_from_file(
+        wav_path,
+        target_sr=sample_rate,
+        n_segments=-1,
+        trim=False,
+    ).samples
+
+    if wav.shape[0] % pad_multiple != 0:
+        wav = np.concatenate([wav, np.zeros(pad_multiple - wav.shape[0] % pad_multiple)])
+    wav = wav[:-1]
+
+    return wav
+
+
+def save_pitch_contour(record):
+    wav_path = record['wav_path']
+    wav_text_id = record['wav_id']
+    sup_data_dir = record['sup_data_dir']
+    stft_params = record['stft_params']
+    wav = load_wav(wav_path, stft_params['sample_rate'], stft_params['pad_multiple'])
+    pitch_contour_fn = f"pitch_contour_{wav_text_id}.pt"
+    pitch_contour_fp = os.path.join(sup_data_dir, pitch_contour_fn)
+
+    f0, _, _ = librosa.pyin(
+        wav,
+        fmin=librosa.note_to_hz('C2'),
+        fmax=stft_params['yin_fmax'],
+        frame_length=stft_params['win_length'],
+        hop_length=stft_params['hop_length'],
+        sr=stft_params['sample_rate'],
+        center=True,
+        fill_na=0.0,
+    )
+
+    pitch_contour = torch.tensor(f0, dtype=torch.float32)
+    torch.save(pitch_contour, pitch_contour_fp)
+    logging.info("saved {}".format(pitch_contour_fp))
+
+    return pitch_contour
+
+
+def compute_pitch_stats(records):
+    def _is_valid_pitch(pitch_mean, pitch_std):
+        c1 = pitch_mean > 0 and pitch_mean < 1000
+        c2 = pitch_std > 0 and pitch_std < 1000
+        return c1 and c2
+
+    speaker_wise_pitch_contours = {}
+    for item in records:
+        wav_id = item['wav_id']
+        speaker = item['speaker']
+        sup_data_dir = item['sup_data_dir']
+        pitch_contour_fn = f"pitch_contour_{wav_id}.pt"
+        pitch_contour_fp = os.path.join(sup_data_dir, pitch_contour_fn)
+        if speaker not in speaker_wise_pitch_contours:
+            speaker_wise_pitch_contours[speaker] = []
+        speaker_wise_pitch_contours[speaker].append(pitch_contour_fp)
+
+    speaker_pitch_stats = {}
+    for speaker in speaker_wise_pitch_contours:
+        non_zero_pc = []
+        for pitch_contour_fp in speaker_wise_pitch_contours[speaker][:50]:
+            pitch_contour = torch.load(pitch_contour_fp)
+            pitch_contour_nonzero = pitch_contour[pitch_contour != 0]
+            if len(pitch_contour_nonzero) > 0:
+                non_zero_pc.append(pitch_contour_nonzero)
+
+        if len(non_zero_pc) > 0:
+            non_zero_pc = torch.cat(non_zero_pc)
+            pitch_mean = non_zero_pc.mean().item()
+            pitch_std = non_zero_pc.std().item()
+            valid = True
+
+            if not _is_valid_pitch(pitch_mean, pitch_std):
+                logging.warning("invalid pitch: {}".format(speaker))
+                pitch_mean = 212.0
+                pitch_std = 70.0
+                valid = "False"
+        else:
+            logging.warning("could not find pitch contour for speaker {}".format(speaker))
+            valid = "False"
+            pitch_mean = 212.0
+            pitch_std = 70.0
+
+        speaker_pitch_stats[speaker] = {"pitch_mean": pitch_mean, "pitch_std": pitch_std, "valid": valid}
+
+    with open(os.path.join(sup_data_dir, "speaker_pitch_stats.json"), "w") as f:
+        json.dump(speaker_pitch_stats, f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Evaluate the model')
+    parser.add_argument(
+        '--ssl_model_ckpt_path',
+        type=str,
+        required=True,
+    )
+    parser.add_argument('--manifest_paths', type=str, required=True)
+    parser.add_argument('--sup_data_dir', type=str, default=None)
+    parser.add_argument('--batch_size', type=int, default=32)
+    parser.add_argument('--ssl_content_emb_type', type=str, default="embedding_and_probs")
+    parser.add_argument('--use_unique_tokens', type=int, default=1)
+    parser.add_argument('--num_workers', type=int, default=8)
+    parser.add_argument('--pool_workers', type=int, default=30)
+    parser.add_argument('--compute_pitch_contours', type=int, default=1)
+    parser.add_argument('--num_pitch_per_speaker', type=int, default=None)  # saves time.
+    parser.add_argument('--sample_rate', type=int, default=22050)
+    parser.add_argument('--pad_multiple', type=int, default=1024)
+    parser.add_argument('--ssl_downsampling_factor', type=int, default=4)
+    parser.add_argument('--stft_n_fft', type=int, default=1024)
+    parser.add_argument('--stft_hop_length', type=int, default=256)
+    parser.add_argument('--stft_win_length', type=int, default=1024)
+    parser.add_argument('--stft_n_mel', type=int, default=80)
+    parser.add_argument('--stft_fmin', type=int, default=0)
+    parser.add_argument('--stft_fmax', type=int, default=8000)
+    parser.add_argument('--yin_fmax', type=int, default=500)
+    parser.add_argument('--segment_length', type=int, default=44100)
+    parser.add_argument('--segment_hop_size', type=int, default=22050)
+    parser.add_argument('--min_segment_length', type=int, default=22050)
+
+    args = parser.parse_args()
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+    manifest_paths = args.manifest_paths.split(",")
+    ssl_model_ckpt_path = args.ssl_model_ckpt_path
+
+    dataset = AudioDataset(
+        manifest_paths, pad_multiple=args.pad_multiple, sample_rate=args.sample_rate, sup_data_dir=args.sup_data_dir
+    )
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        collate_fn=dataset.pad_collate_fn,
+        num_workers=args.num_workers,
+    )
+
+    ssl_model = ssl_tts.SSLDisentangler.load_from_checkpoint(ssl_model_ckpt_path, strict=False)
+    with open_dict(ssl_model.cfg):
+        ssl_model.cfg.preprocessor.exact_pad = True
+    ssl_model.preprocessor = safe_instantiate(ssl_model.cfg.preprocessor)
+    ssl_model.preprocessor_disentangler = ssl_model.preprocessor
+    ssl_model.eval()
+    ssl_model.to(device)
+
+    sample_rate = args.sample_rate
+    stft_params = {
+        "n_fft": args.stft_n_fft,
+        "hop_length": args.stft_hop_length,
+        "win_length": args.stft_win_length,
+        "n_mel": args.stft_n_mel,
+        "sample_rate": sample_rate,
+        "pad_multiple": args.pad_multiple,
+        "fmin": args.stft_fmin,
+        "fmax": args.stft_fmax,
+        "yin_fmax": args.yin_fmax,
+    }
+
+    fb = (
+        torch.tensor(
+            librosa.filters.mel(
+                sr=sample_rate,
+                n_fft=stft_params['n_fft'],
+                n_mels=stft_params['n_mel'],
+                fmin=stft_params['fmin'],
+                fmax=stft_params['fmax'],
+            ),
+            dtype=torch.float,
+        )
+        .unsqueeze(0)
+        .to(device)
+    )
+
+    st = time.time()
+    bidx = 0
+    wav_and_id_list = []
+
+    for batch in tqdm(dataloader):
+        bidx += 1
+        with torch.no_grad():
+            (
+                _,
+                _,
+                batch_content_embedding,
+                batch_content_log_probs,
+                batch_encoded_len,
+            ) = ssl_model.forward_for_export(
+                input_signal=batch['audio'].to(device),
+                input_signal_length=batch['audio_len'].to(device),
+                normalize_content=True,
+            )
+
+            batch_mel_specs = get_mel_spectrogram(fb, batch['audio'][:, :-1].to(device), stft_params)
+            audio_segmented, segment_indices = segment_batch(
+                batch, args.segment_length, args.segment_hop_size, args.min_segment_length
+            )
+            audio_seg_len = torch.tensor([len(segment) for segment in audio_segmented]).to(device).long()
+
+            _, batch_speaker_embeddings, _, _, _ = ssl_model.forward_for_export(
+                input_signal=audio_segmented.to(device),
+                input_signal_length=audio_seg_len,
+                normalize_content=True,
+            )
+
+            for idx in range(batch['audio'].shape[0]):
+                _speaker = batch['speaker'][idx].item()
+                wav_path = batch['wav_path'][idx]
+
+                wav_id = batch['rel_audio_path_as_text_id'][idx]
+                wav_and_id_list.append((wav_path, wav_id, _speaker))
+                content_embedding = batch_content_embedding[idx].detach()
+                content_log_probs = batch_content_log_probs[:, idx, :].detach()  # (content lob prob is (t, b, c))
+                encoded_len = batch_encoded_len[idx].detach()
+                content_embedding = content_embedding[: encoded_len.item()]
+                content_embedding = content_embedding.t()
+                content_log_probs = content_log_probs[: encoded_len.item()]
+                content_log_probs = content_log_probs.t()
+                content_probs = torch.exp(content_log_probs)
+
+                duration = torch.ones(content_embedding.shape[1]) * args.ssl_downsampling_factor
+
+                bsi_start = segment_indices[idx][0]
+                bsi_end = segment_indices[idx][1]
+                speaker_embedding = torch.mean(batch_speaker_embeddings[bsi_start : bsi_end + 1], dim=0)
+
+                l2_norm = torch.norm(speaker_embedding, p=2)
+                speaker_embedding = speaker_embedding / l2_norm
+
+                if args.ssl_content_emb_type == "probs":
+                    # content embedding is only character probabilities
+                    final_content_embedding = content_probs
+                elif args.ssl_content_emb_type == "embedding":
+                    # content embedding is only output of content head of SSL backbone
+                    final_content_embedding = content_embedding
+                elif args.ssl_content_emb_type == "log_probs":
+                    # content embedding is only log of character probabilities
+                    final_content_embedding = content_log_probs
+                elif args.ssl_content_emb_type == "embedding_and_probs":
+                    # content embedding is the concatenation of character probabilities and output of content head of SSL backbone
+                    final_content_embedding = torch.cat([content_embedding, content_probs], dim=0)
+
+                if args.use_unique_tokens == 1:
+                    # group content embeddings with same predicted token (by averaging) and add the durations of the grouped embeddings
+                    # Eg. By default each content embedding corresponds to 4 frames of spectrogram (ssl_downsampling_factor)
+                    # If we group 3 content embeddings, the duration of the grouped embedding will be 12 frames.
+                    # This is useful for adapting the duration during inference based on the speaker.
+                    token_predictions = torch.argmax(content_probs, dim=0)
+                    content_buffer = [final_content_embedding[:, 0]]
+                    unique_content_embeddings = []
+                    unique_tokens = []
+                    durations = []
+                    for _t in range(1, final_content_embedding.shape[1]):
+                        if token_predictions[_t] == token_predictions[_t - 1]:
+                            content_buffer.append(final_content_embedding[:, _t])
+                        else:
+                            durations.append(len(content_buffer) * args.ssl_downsampling_factor)
+                            unique_content_embeddings.append(torch.mean(torch.stack(content_buffer), dim=0))
+                            content_buffer = [final_content_embedding[:, _t]]
+                            unique_tokens.append(token_predictions[_t].item())
+
+                    if len(content_buffer) > 0:
+                        durations.append(len(content_buffer) * args.ssl_downsampling_factor)
+                        unique_content_embeddings.append(torch.mean(torch.stack(content_buffer), dim=0))
+                        unique_tokens.append(token_predictions[_t].item())
+
+                    unique_content_embedding = torch.stack(unique_content_embeddings)
+                    final_content_embedding = unique_content_embedding.t()
+                    duration = torch.tensor(durations).float()
+
+                mel_len = int(batch['audio_len'][idx].item() / stft_params['hop_length'])
+                item_mel = batch_mel_specs[idx][:, :mel_len]
+
+                wav_text_id = batch["rel_audio_path_as_text_id"][idx]
+                content_emb_fn = f"{args.ssl_content_emb_type}_content_embedding_{wav_text_id}.pt"
+                speaker_emb_fn = f"speaker_embedding_{wav_text_id}.pt"
+                duration_fn = f"duration_embedding_{wav_text_id}.pt"  # embedding just for namesake
+                content_emb_fp = os.path.join(dataset.sup_data_dir, content_emb_fn)
+                speaker_emb_fp = os.path.join(dataset.sup_data_dir, speaker_emb_fn)
+                duration_fp = os.path.join(dataset.sup_data_dir, duration_fn)
+
+                mel_spec_fn = f"mel_spec_{wav_text_id}.pt"
+                mel_spec_fp = os.path.join(dataset.sup_data_dir, mel_spec_fn)
+
+                torch.save(item_mel.cpu(), mel_spec_fp)
+                torch.save(final_content_embedding.cpu(), content_emb_fp)
+                torch.save(speaker_embedding.cpu(), speaker_emb_fp)
+                torch.save(duration.cpu(), duration_fp)
+
+            et = time.time()
+            logging.info(
+                "Processed Batch {} of {} | Time per batch: {:.4f} s".format(
+                    bidx + 1, len(dataloader), (et - st) / bidx
+                )
+            )
+
+    if args.compute_pitch_contours == 1:
+        speaker_wise_records = {}
+        for row in wav_and_id_list:
+            wav_path, wav_id, speaker = row
+            if speaker not in speaker_wise_records:
+                speaker_wise_records[speaker] = []
+            speaker_wise_records[speaker].append(
+                {
+                    "wav_path": wav_path,
+                    "wav_id": wav_id,
+                    "sup_data_dir": dataset.sup_data_dir,
+                    "stft_params": stft_params,
+                    "speaker": speaker,
+                }
+            )
+
+        filtered_records = []
+        for speaker in speaker_wise_records:
+            if args.num_pitch_per_speaker is not None:
+                filtered_records += speaker_wise_records[speaker][: args.num_pitch_per_speaker]
+            else:
+                filtered_records += speaker_wise_records[speaker]
+
+        with Pool(args.pool_workers) as p:
+            p.map(save_pitch_contour, filtered_records)
+
+        compute_pitch_stats(filtered_records)
+
+
+if __name__ == '__main__':
+    main()

scripts/ssl_tts/ssl_tts_vc.py

@@ -0,0 +1,297 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Example Run Command: python ssl_tts_vc.py --ssl_model_ckpt_path <PATH TO CKPT> --hifi_ckpt_path <PATH TO CKPT> \
+# --fastpitch_ckpt_path <PATH TO CKPT> --source_audio_path <SOURCE CONTENT WAV PATH> --target_audio_path \
+# <TARGET SPEAKER WAV PATH> --out_path <PATH TO OUTPUT WAV>
+
+import argparse
+import os
+
+import librosa
+import soundfile
+import torch
+
+from nemo.collections.asr.parts.preprocessing.features import WaveformFeaturizer
+from nemo.collections.tts.models import fastpitch_ssl, hifigan, ssl_tts
+from nemo.collections.tts.parts.utils.tts_dataset_utils import get_base_dir
+
+
+def load_wav(wav_path, wav_featurizer, pad_multiple=1024):
+    wav = wav_featurizer.process(wav_path)
+    if (wav.shape[0] % pad_multiple) != 0:
+        wav = torch.cat([wav, torch.zeros(pad_multiple - wav.shape[0] % pad_multiple, dtype=torch.float)])
+    wav = wav[:-1]
+
+    return wav
+
+
+def get_pitch_contour(wav, pitch_mean=None, pitch_std=None, compute_mean_std=False, sample_rate=22050):
+    f0, _, _ = librosa.pyin(
+        wav.numpy(),
+        fmin=librosa.note_to_hz('C2'),
+        fmax=librosa.note_to_hz('C7'),
+        frame_length=1024,
+        hop_length=256,
+        sr=sample_rate,
+        center=True,
+        fill_na=0.0,
+    )
+    pitch_contour = torch.tensor(f0, dtype=torch.float32)
+    _pitch_mean = pitch_contour.mean().item()
+    _pitch_std = pitch_contour.std().item()
+    if compute_mean_std:
+        pitch_mean = _pitch_mean
+        pitch_std = _pitch_std
+    if (pitch_mean is not None) and (pitch_std is not None):
+        pitch_contour = pitch_contour - pitch_mean
+        pitch_contour[pitch_contour == -pitch_mean] = 0.0
+        pitch_contour = pitch_contour / pitch_std
+
+    return pitch_contour
+
+
+def segment_wav(wav, segment_length=44100, hop_size=44100, min_segment_size=22050):
+    if len(wav) < segment_length:
+        pad = torch.zeros(segment_length - len(wav))
+        segment = torch.cat([wav, pad])
+        return [segment]
+    else:
+        si = 0
+        segments = []
+        while si < len(wav) - min_segment_size:
+            segment = wav[si : si + segment_length]
+            if len(segment) < segment_length:
+                pad = torch.zeros(segment_length - len(segment))
+                segment = torch.cat([segment, pad])
+
+            segments.append(segment)
+            si += hop_size
+        return segments
+
+
+def get_speaker_embedding(ssl_model, wav_featurizer, audio_paths, duration=None, device="cpu"):
+    all_segments = []
+    all_wavs = []
+    for audio_path in audio_paths:
+        wav = load_wav(audio_path, wav_featurizer)
+        segments = segment_wav(wav)
+        all_segments += segments
+        all_wavs.append(wav)
+        if duration is not None and len(all_segments) >= duration:
+            # each segment is 2 seconds with one second overlap.
+            # so 10 segments would mean 0 to 2, 1 to 3.. 9 to 11 (11 seconds.)
+            all_segments = all_segments[: int(duration)]
+            break
+
+    signal_batch = torch.stack(all_segments)
+    signal_length_batch = torch.stack([torch.tensor(signal_batch.shape[1]) for _ in range(len(all_segments))])
+    signal_batch = signal_batch.to(device)
+    signal_length_batch = signal_length_batch.to(device)
+    _, speaker_embeddings, _, _, _ = ssl_model.forward_for_export(
+        input_signal=signal_batch, input_signal_length=signal_length_batch, normalize_content=True
+    )
+
+    speaker_embedding = torch.mean(speaker_embeddings, dim=0)
+    l2_norm = torch.norm(speaker_embedding, p=2)
+    speaker_embedding = speaker_embedding / l2_norm
+
+    return speaker_embedding[None]
+
+
+def get_ssl_features_disentangled(
+    ssl_model, wav_featurizer, audio_path, emb_type="embedding_and_probs", use_unique_tokens=False, device="cpu"
+):
+    """
+    Extracts content embedding, speaker embedding and duration tokens to be used as inputs for FastPitchModel_SSL
+    synthesizer. Content embedding and speaker embedding extracted using SSLDisentangler model.
+    Args:
+        ssl_model: SSLDisentangler model
+        wav_featurizer: WaveformFeaturizer object
+        audio_path: path to audio file
+        emb_type: Can be one of embedding_and_probs, embedding, probs, log_probs
+        use_unique_tokens: If True, content embeddings with same predicted token are grouped and duration is different.
+        device: device to run the model on
+    Returns:
+        content_embedding, speaker_embedding, duration
+    """
+    wav = load_wav(audio_path, wav_featurizer)
+    audio_signal = wav[None]
+    audio_signal_length = torch.tensor([wav.shape[0]])
+    audio_signal = audio_signal.to(device)
+    audio_signal_length = audio_signal_length.to(device)
+    _, speaker_embedding, content_embedding, content_log_probs, encoded_len = ssl_model.forward_for_export(
+        input_signal=audio_signal, input_signal_length=audio_signal_length, normalize_content=True
+    )
+
+    content_embedding = content_embedding[0, : encoded_len[0].item()]
+    content_log_probs = content_log_probs[: encoded_len[0].item(), 0, :]
+    content_embedding = content_embedding.t()
+    content_log_probs = content_log_probs.t()
+    content_probs = torch.exp(content_log_probs)
+
+    ssl_downsampling_factor = ssl_model._cfg.encoder.subsampling_factor
+
+    if emb_type == "probs":
+        # content embedding is only character probabilities
+        final_content_embedding = content_probs
+
+    elif emb_type == "embedding":
+        # content embedding is only output of content head of SSL backbone
+        final_content_embedding = content_embedding
+
+    elif emb_type == "log_probs":
+        # content embedding is only log of character probabilities
+        final_content_embedding = content_log_probs
+
+    elif emb_type == "embedding_and_probs":
+        # content embedding is the concatenation of character probabilities and output of content head of SSL backbone
+        final_content_embedding = torch.cat([content_embedding, content_probs], dim=0)
+
+    else:
+        raise ValueError(
+            f"{emb_type} is not valid. Valid emb_type includes probs, embedding, log_probs or embedding_and_probs."
+        )
+
+    duration = torch.ones(final_content_embedding.shape[1]) * ssl_downsampling_factor
+    if use_unique_tokens:
+        # group content embeddings with same predicted token (by averaging) and add the durations of the grouped embeddings
+        # Eg. By default each content embedding corresponds to 4 frames of spectrogram (ssl_downsampling_factor)
+        # If we group 3 content embeddings, the duration of the grouped embedding will be 12 frames.
+        # This is useful for adapting the duration during inference based on the speaker.
+        token_predictions = torch.argmax(content_probs, dim=0)
+        content_buffer = [final_content_embedding[:, 0]]
+        unique_content_embeddings = []
+        unique_tokens = []
+        durations = []
+        for _t in range(1, final_content_embedding.shape[1]):
+            if token_predictions[_t] == token_predictions[_t - 1]:
+                content_buffer.append(final_content_embedding[:, _t])
+            else:
+                durations.append(len(content_buffer) * ssl_downsampling_factor)
+                unique_content_embeddings.append(torch.mean(torch.stack(content_buffer), dim=0))
+                content_buffer = [final_content_embedding[:, _t]]
+                unique_tokens.append(token_predictions[_t].item())
+
+        if len(content_buffer) > 0:
+            durations.append(len(content_buffer) * ssl_downsampling_factor)
+            unique_content_embeddings.append(torch.mean(torch.stack(content_buffer), dim=0))
+            unique_tokens.append(token_predictions[_t].item())
+
+        unique_content_embedding = torch.stack(unique_content_embeddings)
+        final_content_embedding = unique_content_embedding.t()
+        duration = torch.tensor(durations).float()
+
+    duration = duration.to(device)
+    return final_content_embedding[None], speaker_embedding, duration[None]
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Evaluate the model')
+    parser.add_argument('--ssl_model_ckpt_path', type=str)
+    parser.add_argument('--hifi_ckpt_path', type=str)
+    parser.add_argument('--fastpitch_ckpt_path', type=str)
+    parser.add_argument('--source_audio_path', type=str)
+    parser.add_argument('--target_audio_path', type=str)  # can be a list seperated by comma
+    parser.add_argument('--out_path', type=str)
+    parser.add_argument('--source_target_out_pairs', type=str)
+    parser.add_argument('--use_unique_tokens', type=int, default=0)
+    parser.add_argument('--compute_pitch', type=int, default=0)
+    parser.add_argument('--compute_duration', type=int, default=0)
+    args = parser.parse_args()
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+    if args.source_target_out_pairs is not None:
+        assert args.source_audio_path is None, "source_audio_path and source_target_out_pairs are mutually exclusive"
+        assert args.target_audio_path is None, "target_audio_path and source_target_out_pairs are mutually exclusive"
+        assert args.out_path is None, "out_path and source_target_out_pairs are mutually exclusive"
+        with open(args.source_target_out_pairs, "r") as f:
+            lines = f.readlines()
+            source_target_out_pairs = [line.strip().split(";") for line in lines]
+    else:
+        assert args.source_audio_path is not None, "source_audio_path is required"
+        assert args.target_audio_path is not None, "target_audio_path is required"
+        if args.out_path is None:
+            source_name = os.path.basename(args.source_audio_path).split(".")[0]
+            target_name = os.path.basename(args.target_audio_path).split(".")[0]
+            args.out_path = "swapped_{}_{}.wav".format(source_name, target_name)
+
+        source_target_out_pairs = [(args.source_audio_path, args.target_audio_path, args.out_path)]
+
+    out_paths = [r[2] for r in source_target_out_pairs]
+    out_dir = get_base_dir(out_paths)
+    if not os.path.exists(out_dir):
+        os.makedirs(out_dir)
+
+    ssl_model = ssl_tts.SSLDisentangler.load_from_checkpoint(args.ssl_model_ckpt_path, strict=False)
+    ssl_model = ssl_model.to(device)
+    ssl_model.eval()
+
+    vocoder = hifigan.HifiGanModel.load_from_checkpoint(args.hifi_ckpt_path).to(device)
+    vocoder.eval()
+
+    fastpitch_model = fastpitch_ssl.FastPitchModel_SSL.load_from_checkpoint(args.fastpitch_ckpt_path, strict=False)
+    fastpitch_model = fastpitch_model.to(device)
+    fastpitch_model.eval()
+    fastpitch_model.non_trainable_models = {'vocoder': vocoder}
+    fpssl_sample_rate = fastpitch_model._cfg.sample_rate
+
+    wav_featurizer = WaveformFeaturizer(sample_rate=fpssl_sample_rate, int_values=False, augmentor=None)
+
+    use_unique_tokens = args.use_unique_tokens == 1
+    compute_pitch = args.compute_pitch == 1
+    compute_duration = args.compute_duration == 1
+
+    for source_target_out in source_target_out_pairs:
+        source_audio_path = source_target_out[0]
+        target_audio_paths = source_target_out[1].split(",")
+        out_path = source_target_out[2]
+
+        with torch.no_grad():
+            content_embedding1, _, duration1 = get_ssl_features_disentangled(
+                ssl_model,
+                wav_featurizer,
+                source_audio_path,
+                emb_type="embedding_and_probs",
+                use_unique_tokens=use_unique_tokens,
+                device=device,
+            )
+
+            speaker_embedding2 = get_speaker_embedding(
+                ssl_model, wav_featurizer, target_audio_paths, duration=None, device=device
+            )
+
+            pitch_contour1 = None
+            if not compute_pitch:
+                pitch_contour1 = get_pitch_contour(
+                    load_wav(source_audio_path, wav_featurizer), compute_mean_std=True, sample_rate=fpssl_sample_rate
+                )[None]
+                pitch_contour1 = pitch_contour1.to(device)
+
+            wav_generated = fastpitch_model.generate_wav(
+                content_embedding1,
+                speaker_embedding2,
+                pitch_contour=pitch_contour1,
+                compute_pitch=compute_pitch,
+                compute_duration=compute_duration,
+                durs_gt=duration1,
+                dataset_id=0,
+            )
+            wav_generated = wav_generated[0][0]
+            soundfile.write(out_path, wav_generated, fpssl_sample_rate)
+
+
+if __name__ == "__main__":
+    main()